import Spa.Language
import Spa.Lattice.FiniteMap
import Spa.Interp

namespace Spa

namespace Forward

variable (L : Type) [Lattice L] (prog : Program)

abbrev VariableValues : Type := FiniteMap String L prog.vars

abbrev StateVariables : Type := FiniteMap prog.State (VariableValues L prog) prog.states

def botV [FiniteHeightLattice L] : VariableValues L prog :=
  (⊥ : VariableValues L prog)

variable {L prog}

omit [Lattice L] in
lemma states_memKey (s : prog.State) (sv : StateVariables L prog) :
    FiniteMap.MemKey s sv :=
  FiniteMap.MemKey_iff.mpr (prog.states_complete s)

def variablesAt (s : prog.State) (sv : StateVariables L prog) :
    VariableValues L prog :=
  (FiniteMap.locate (states_memKey s sv)).1

omit [Lattice L] in
lemma variablesAt_mem (s : prog.State) (sv : StateVariables L prog) :
    (s, variablesAt s sv) ∈ sv :=
  (FiniteMap.locate (states_memKey s sv)).2

lemma variablesAt_le {sv₁ sv₂ : StateVariables L prog} (hle : sv₁ ≤ sv₂)
    (s : prog.State) : variablesAt s sv₁ ≤ variablesAt s sv₂ :=
  FiniteMap.le_of_mem_mem prog.states_nodup hle
    (variablesAt_mem s sv₁) (variablesAt_mem s sv₂)

variable [FiniteHeightLattice L]

def joinForKey (k : prog.State) (sv : StateVariables L prog) :
    VariableValues L prog :=
  (sv.valuesAt (prog.incoming k)).foldr (· ⊔ ·) (botV L prog)

lemma joinForKey_mono (k : prog.State) :
    Monotone (joinForKey (L := L) k) := by
  intro sv₁ sv₂ hle
  exact foldr_mono _ (FiniteMap.valuesAt_le hle (prog.incoming k)) (le_refl _)
    (fun b _ _ hab => sup_le_sup_right hab b)
    (fun a _ _ hab => sup_le_sup_left hab a)

def joinAll (sv : StateVariables L prog) : StateVariables L prog :=
  FiniteMap.generalizedUpdate id joinForKey prog.states sv

lemma joinAll_mono : Monotone (joinAll (L := L) (prog := prog)) :=
  FiniteMap.generalizedUpdate_monotone monotone_id joinForKey_mono

lemma joinAll_mem_eq {s : prog.State} {vs : VariableValues L prog}
    {sv : StateVariables L prog} (h : (s, vs) ∈ joinAll sv) :
    vs = joinForKey s sv :=
  FiniteMap.generalizedUpdate_mem_eq (prog.states_complete s) h

lemma variablesAt_joinAll (s : prog.State) (sv : StateVariables L prog) :
    variablesAt s (joinAll sv) = joinForKey s sv :=
  joinAll_mem_eq (variablesAt_mem s (joinAll sv))

class StateInterp (L : Type) [Lattice L] (prog : Program) where
  St : Env → Type
  init : St []
  interp : VariableValues L prog → (ρ : Env) → St ρ → Prop
  interp_sup : ∀ {vs₁ vs₂ : VariableValues L prog} {ρ : Env} {st : St ρ},
    interp vs₁ ρ st ∨ interp vs₂ ρ st → interp (vs₁ ⊔ vs₂) ρ st
  interp_inf : ∀ {vs₁ vs₂ : VariableValues L prog} {ρ : Env} {st : St ρ},
    interp vs₁ ρ st ∧ interp vs₂ ρ st → interp (vs₁ ⊓ vs₂) ρ st

instance [S : StateInterp L prog] :
    Interp (VariableValues L prog) ((ρ : Env) → S.St ρ → Prop) :=
  ⟨S.interp⟩

lemma interp_foldr [S : StateInterp L prog]
    {vs : VariableValues L prog} {vss : List (VariableValues L prog)}
    {ρ : Env} {st : S.St ρ} (hvs : ⟦ vs ⟧ ρ st) (hmem : vs ∈ vss) :
    ⟦ vss.foldr (· ⊔ ·) (botV L prog) ⟧ ρ st := by
  induction vss with
  | nil => cases hmem
  | cons vs' vss' ih =>
    rcases List.mem_cons.mp hmem with rfl | hmem'
    · exact S.interp_sup (Or.inl hvs)
    · exact S.interp_sup (Or.inr (ih hmem'))

variable [I : LatticeInterpretation L]

instance : StateInterp L prog where
  St := fun _ => PUnit
  init := PUnit.unit
  interp vs ρ _ := ∀ (k : String) (l : L), (k, l) ∈ vs →
    ∀ (v : Value), Env.Mem (k, v) ρ → I.interp l v
  interp_sup := by
    intro vs₁ vs₂ ρ st h k l hmem v hv
    obtain ⟨l₁, l₂, rfl, h₁, h₂⟩ := FiniteMap.mem_sup hmem
    rcases h with h | h
    · exact I.interp_sup v (Or.inl (h _ _ h₁ _ hv))
    · exact I.interp_sup v (Or.inr (h _ _ h₂ _ hv))
  interp_inf := by
    intro vs₁ vs₂ ρ st h k l hmem v hv
    obtain ⟨l₁, l₂, rfl, h₁, h₂⟩ := FiniteMap.mem_inf hmem
    exact I.interp_inf v ⟨h.1 _ _ h₁ _ hv, h.2 _ _ h₂ _ hv⟩

end Forward

end Spa