import Spa.Lattice.Tuple
import Mathlib.Data.List.Nodup

/-!

# Finite Maps

This file defines _finite maps_, or key-value maps with a finite domain. This
is encoded as a map from `Fin` into the value type. Finite maps form a
lattice from pointwise composition: $(f \land g) k = f k \land g k$,
and, provided the domain `\beta` is of finite height, so is the map
lattice as a whole.

In fact, the isomorphism is described and proven in `Spa/Lattice/Tuple.lean`.

-/

namespace Spa

/-- Key-value map with domain `α` and codomain `β`, with possible keys $\textit{ks} \subseteq \alpha$. -/
def FiniteMap (α β : Type*) (ks : List α) : Type _ := Fin ks.length → β

namespace FiniteMap

variable {α β : Type*} {ks : List α}

instance [Lattice β] : Lattice (FiniteMap α β ks) :=
  inferInstanceAs (Lattice (Fin ks.length → β))

instance [FiniteHeightLattice β] : FiniteHeightLattice (FiniteMap α β ks) :=
  inferInstanceAs (FiniteHeightLattice (Fin ks.length → β))

instance [DecidableEq β] : DecidableEq (FiniteMap α β ks) :=
  inferInstanceAs (DecidableEq (Fin ks.length → β))

instance : Membership (α × β) (FiniteMap α β ks) :=
  ⟨fun fm p => ∃ i : Fin ks.length, ks.get i = p.1 ∧ fm i = p.2⟩

lemma mem_iff {fm : FiniteMap α β ks} {p : α × β} :
    p ∈ fm ↔ ∃ i : Fin ks.length, ks.get i = p.1 ∧ fm i = p.2 := Iff.rfl

def MemKey (k : α) (_fm : FiniteMap α β ks) : Prop := k ∈ ks

lemma MemKey_iff {k : α} {fm : FiniteMap α β ks} : MemKey k fm ↔ k ∈ ks := Iff.rfl

instance {k : α} {fm : FiniteMap α β ks} [DecidableEq α] : Decidable (MemKey k fm) :=
  decidable_of_iff _ MemKey_iff.symm

lemma mem_key_of_mem {k : α} {v : β} {fm : FiniteMap α β ks}
    (h : (k, v) ∈ fm) : MemKey k fm := by
  obtain ⟨i, hi, _⟩ := h
  have hik : ks.get i = k := hi
  exact hik ▸ ks.get_mem i

def toList (fm : FiniteMap α β ks) : List (α × β) :=
  (List.finRange ks.length).map fun i => (ks.get i, fm i)

lemma le_def [Lattice β] {fm₁ fm₂ : FiniteMap α β ks} :
    fm₁ ≤ fm₂ ↔ ∀ i, fm₁ i ≤ fm₂ i := Iff.rfl

section Locate

variable [DecidableEq α]

/-- Recover the value stored under a present key. -/
def locate {k : α} {fm : FiniteMap α β ks} (h : MemKey k fm) :
    {v : β // (k, v) ∈ fm} :=
  let i : Fin ks.length := ⟨ks.idxOf k, List.idxOf_lt_length_iff.mpr h⟩
  ⟨fm i, i, List.idxOf_get _, rfl⟩

end Locate

variable [Lattice β]

lemma le_of_mem_mem (hks : ks.Nodup) {fm₁ fm₂ : FiniteMap α β ks}
    (hle : fm₁ ≤ fm₂) {k : α} {v₁ v₂ : β}
    (h₁ : (k, v₁) ∈ fm₁) (h₂ : (k, v₂) ∈ fm₂) : v₁ ≤ v₂ := by
  obtain ⟨i, hi, rfl⟩ := h₁
  obtain ⟨j, hj, rfl⟩ := h₂
  have hij : i = j := hks.get_inj_iff.mp (hi.trans hj.symm)
  subst hij
  exact le_def.mp hle i

lemma mem_sup {fm₁ fm₂ : FiniteMap α β ks} {k : α} {v : β}
    (h : (k, v) ∈ fm₁ ⊔ fm₂) :
    ∃ v₁ v₂, v = v₁ ⊔ v₂ ∧ (k, v₁) ∈ fm₁ ∧ (k, v₂) ∈ fm₂ := by
  obtain ⟨i, hi, rfl⟩ := h
  exact ⟨fm₁ i, fm₂ i, rfl, ⟨i, hi, rfl⟩, ⟨i, hi, rfl⟩⟩

lemma mem_inf {fm₁ fm₂ : FiniteMap α β ks} {k : α} {v : β}
    (h : (k, v) ∈ fm₁ ⊓ fm₂) :
    ∃ v₁ v₂, v = v₁ ⊓ v₂ ∧ (k, v₁) ∈ fm₁ ∧ (k, v₂) ∈ fm₂ := by
  obtain ⟨i, hi, rfl⟩ := h
  exact ⟨fm₁ i, fm₂ i, rfl, ⟨i, hi, rfl⟩, ⟨i, hi, rfl⟩⟩

section Updating

variable [DecidableEq α]

def updating (fm : FiniteMap α β ks) (ks' : List α) (g : α → β) : FiniteMap α β ks :=
  fun i => if ks.get i ∈ ks' then g (ks.get i) else fm i

omit [Lattice β] in
lemma eq_of_mem_updating {k : α} {v : β} {fm : FiniteMap α β ks}
    {ks' : List α} {g : α → β} (hk : k ∈ ks')
    (h : (k, v) ∈ updating fm ks' g) : v = g k := by
  obtain ⟨i, hi, rfl⟩ := h
  show (if ks.get i ∈ ks' then g (ks.get i) else fm i) = g k
  rw [if_pos (by rw [hi]; exact hk), hi]

omit [Lattice β] in
lemma mem_of_mem_updating {k : α} {v : β} {fm : FiniteMap α β ks}
    {ks' : List α} {g : α → β} (hk : k ∉ ks')
    (h : (k, v) ∈ updating fm ks' g) : (k, v) ∈ fm := by
  obtain ⟨i, hi, rfl⟩ := h
  refine ⟨i, hi, ?_⟩
  show fm i = (if ks.get i ∈ ks' then g (ks.get i) else fm i)
  rw [if_neg (by rw [hi]; exact hk)]

lemma updating_mono {fm₁ fm₂ : FiniteMap α β ks} {ks' : List α}
    {g₁ g₂ : α → β} (hfm : fm₁ ≤ fm₂) (hg : ∀ k, g₁ k ≤ g₂ k) :
    updating fm₁ ks' g₁ ≤ updating fm₂ ks' g₂ := by
  rw [le_def]
  intro i
  show (if ks.get i ∈ ks' then g₁ (ks.get i) else fm₁ i)
     ≤ (if ks.get i ∈ ks' then g₂ (ks.get i) else fm₂ i)
  split
  · exact hg (ks.get i)
  · exact le_def.mp hfm i

end Updating

section GeneralizedUpdate

variable [DecidableEq α] {L : Type*} [Lattice L]

def generalizedUpdate (f : L → FiniteMap α β ks) (g : α → L → β)
    (ks' : List α) : L → FiniteMap α β ks := fun l =>
  (f l).updating ks' (fun k => g k l)

variable {f : L → FiniteMap α β ks} {g : α → L → β} {ks' : List α}

lemma generalizedUpdate_monotone (hf : Monotone f)
    (hg : ∀ k, Monotone (g k)) : Monotone (generalizedUpdate f g ks') :=
  fun _ _ hl => updating_mono (hf hl) (fun k => hg k hl)

omit [Lattice β] [Lattice L] in
lemma generalizedUpdate_mem_eq {k : α} {v : β} {l : L} (hk : k ∈ ks')
    (h : (k, v) ∈ generalizedUpdate f g ks' l) : v = g k l :=
  eq_of_mem_updating (g := fun k => g k l) hk h

omit [Lattice β] [Lattice L] in
lemma generalizedUpdate_not_mem_backward {k : α} {v : β} {l : L} (hk : k ∉ ks')
    (h : (k, v) ∈ generalizedUpdate f g ks' l) : (k, v) ∈ f l :=
  mem_of_mem_updating hk h

end GeneralizedUpdate

section ValuesAt

variable [DecidableEq α]

/-- The value stored under `k`, if `k` is a key. -/
private def lookup (fm : FiniteMap α β ks) (k : α) : Option β :=
  if h : k ∈ ks then some (fm ⟨ks.idxOf k, List.idxOf_lt_length_iff.mpr h⟩) else none

/-- The values stored under the keys `ks'` (skipping any that are not keys). -/
def valuesAt (fm : FiniteMap α β ks) (ks' : List α) : List β :=
  ks'.filterMap fm.lookup

omit [Lattice β] in
lemma mem_valuesAt (hks : ks.Nodup) {fm : FiniteMap α β ks} {k : α} {v : β}
    {ks' : List α} (hk : k ∈ ks') (h : (k, v) ∈ fm) : v ∈ valuesAt fm ks' := by
  refine List.mem_filterMap.mpr ⟨k, hk, ?_⟩
  obtain ⟨i, hi, rfl⟩ := h
  have hik : ks.get i = k := hi
  have hmem : k ∈ ks := hik ▸ ks.get_mem i
  show (if h : k ∈ ks then
          some (fm ⟨ks.idxOf k, List.idxOf_lt_length_iff.mpr h⟩) else none) = some (fm i)
  rw [dif_pos hmem]
  have : (⟨ks.idxOf k, List.idxOf_lt_length_iff.mpr hmem⟩ : Fin ks.length) = i :=
    hks.get_inj_iff.mp (by rw [List.idxOf_get, hi])
  rw [this]

private lemma lookup_rel {fm₁ fm₂ : FiniteMap α β ks} (hle : fm₁ ≤ fm₂) (k : α) :
    Option.Rel (· ≤ ·) (fm₁.lookup k) (fm₂.lookup k) := by
  show Option.Rel _
    (if h : k ∈ ks then some (fm₁ ⟨ks.idxOf k, List.idxOf_lt_length_iff.mpr h⟩) else none)
    (if h : k ∈ ks then some (fm₂ ⟨ks.idxOf k, List.idxOf_lt_length_iff.mpr h⟩) else none)
  by_cases hk : k ∈ ks
  · rw [dif_pos hk, dif_pos hk]; exact Option.Rel.some (le_def.mp hle _)
  · rw [dif_neg hk, dif_neg hk]; exact Option.Rel.none

lemma valuesAt_le {fm₁ fm₂ : FiniteMap α β ks} (hle : fm₁ ≤ fm₂)
    (ks' : List α) :
    List.Forall₂ (· ≤ ·) (valuesAt fm₁ ks') (valuesAt fm₂ ks') := by
  induction ks' with
  | nil => exact List.Forall₂.nil
  | cons k ks'' ih =>
    have hrel := lookup_rel hle k
    rw [valuesAt, valuesAt, List.filterMap_cons, List.filterMap_cons]
    revert hrel
    generalize fm₁.lookup k = o₁
    generalize fm₂.lookup k = o₂
    intro hrel
    cases hrel with
    | none => simpa [valuesAt] using ih
    | some hv => exact List.Forall₂.cons hv (by simpa [valuesAt] using ih)

end ValuesAt

end FiniteMap

end Spa