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

namespace Spa

def FiniteMap (A B : Type*) (ks : List A) : Type _ := Fin ks.length → B

namespace FiniteMap

variable {A B : Type*} {ks : List A}

instance [Lattice B] : Lattice (FiniteMap A B ks) :=
  inferInstanceAs (Lattice (Fin ks.length → B))

instance [Lattice B] [FiniteHeightLattice B] : FiniteHeightLattice (FiniteMap A B ks) :=
  inferInstanceAs (FiniteHeightLattice (Fin ks.length → B))

instance [DecidableEq B] : DecidableEq (FiniteMap A B ks) :=
  inferInstanceAs (DecidableEq (Fin ks.length → B))

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

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

def MemKey (k : A) (_fm : FiniteMap A B ks) : Prop := k ∈ ks

theorem MemKey_iff {k : A} {fm : FiniteMap A B ks} : MemKey k fm ↔ k ∈ ks := Iff.rfl

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

theorem mem_key_of_mem {k : A} {v : B} {fm : FiniteMap A B 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 A B ks) : List (A × B) :=
  (List.finRange ks.length).map fun i => (ks.get i, fm i)

theorem le_def [Lattice B] {fm₁ fm₂ : FiniteMap A B ks} :
    fm₁ ≤ fm₂ ↔ ∀ i, fm₁ i ≤ fm₂ i := Iff.rfl

section Locate

variable [DecidableEq A]

/-- Recover the value stored under a present key. -/
def locate {k : A} {fm : FiniteMap A B ks} (h : MemKey k fm) :
    {v : B // (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 B]

theorem le_of_mem_mem (hks : ks.Nodup) {fm₁ fm₂ : FiniteMap A B ks}
    (hle : fm₁ ≤ fm₂) {k : A} {v₁ v₂ : B}
    (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

theorem mem_sup {fm₁ fm₂ : FiniteMap A B ks} {k : A} {v : B}
    (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 A]

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

omit [Lattice B] in
theorem eq_of_mem_updating {k : A} {v : B} {fm : FiniteMap A B ks}
    {ks' : List A} {g : A → B} (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 B] in
theorem mem_of_mem_updating {k : A} {v : B} {fm : FiniteMap A B ks}
    {ks' : List A} {g : A → B} (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)]

theorem updating_mono {fm₁ fm₂ : FiniteMap A B ks} {ks' : List A}
    {g₁ g₂ : A → B} (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 A] {L : Type*} [Lattice L]

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

variable {f : L → FiniteMap A B ks} {g : A → L → B} {ks' : List A}

theorem 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 B] [Lattice L] in
theorem generalizedUpdate_mem_eq {k : A} {v : B} {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 B] [Lattice L] in
theorem generalizedUpdate_not_mem_backward {k : A} {v : B} {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 A]

/-- The value stored under `k`, if `k` is a key. -/
private def lookup (fm : FiniteMap A B ks) (k : A) : Option B :=
  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 A B ks) (ks' : List A) : List B :=
  ks'.filterMap fm.lookup

omit [Lattice B] in
theorem mem_valuesAt (hks : ks.Nodup) {fm : FiniteMap A B ks} {k : A} {v : B}
    {ks' : List A} (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 theorem lookup_rel {fm₁ fm₂ : FiniteMap A B ks} (hle : fm₁ ≤ fm₂) (k : A) :
    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

theorem valuesAt_le {fm₁ fm₂ : FiniteMap A B ks} (hle : fm₁ ≤ fm₂)
    (ks' : List A) :
    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