import Spa.Analysis.Forward.Lattices
import Spa.Analysis.Forward.Evaluation
import Spa.Analysis.Forward.Adapters
import Spa.Fixedpoint

namespace Spa

variable {L : Type} [Lattice L] {prog : Program} [E : StmtEvaluator L prog]

def updateVariablesForState (s : prog.State) (sv : StateVariables L prog) :
    VariableValues L prog :=
  (prog.code s).foldl (fun vs bs => E.eval s bs vs) (variablesAt s sv)

theorem updateVariablesForState_mono (s : prog.State) :
    Monotone (updateVariablesForState (L := L) s) := fun _ _ hle =>
  foldl_mono' (prog.code s) _ (E.eval_mono s ·) (variablesAt_le hle s)

def updateAll (sv : StateVariables L prog) : StateVariables L prog :=
  FiniteMap.generalizedUpdate id updateVariablesForState
    prog.states sv

theorem updateAll_mono : Monotone (updateAll (L := L) (prog := prog)) :=
  FiniteMap.generalizedUpdate_monotone monotone_id updateVariablesForState_mono

theorem updateAll_mem_eq {s : prog.State} {vs : VariableValues L prog}
    {sv : StateVariables L prog} (hmem : (s, vs) ∈ updateAll sv) :
    vs = updateVariablesForState s sv :=
  FiniteMap.generalizedUpdate_mem_eq (prog.states_complete s) hmem

theorem variablesAt_updateAll (s : prog.State) (sv : StateVariables L prog) :
    variablesAt s (updateAll sv) = updateVariablesForState s sv :=
  updateAll_mem_eq (variablesAt_mem s (updateAll sv))

variable [FiniteHeightLattice L]

def analyze (sv : StateVariables L prog) : StateVariables L prog :=
  updateAll (joinAll sv)

theorem analyze_mono : Monotone (analyze (L := L) (prog := prog)) := fun _ _ hle =>
  updateAll_mono (joinAll_mono hle)

variable [DecidableEq L]

variable (L prog) in
def result : StateVariables L prog :=
  Fixedpoint.aFix analyze analyze_mono

variable (L prog) in
theorem result_eq : result L prog = analyze (result L prog) :=
  Fixedpoint.aFix_eq analyze analyze_mono

theorem joinForKey_initialState :
    joinForKey prog.initialState (result L prog) = botV L prog := by
  rw [joinForKey, prog.incoming_initialState_eq_nil]
  rfl

variable [I : LatticeInterpretation L] [V : ValidStmtEvaluator L prog]

omit [FiniteHeightLattice L] [DecidableEq L] in
theorem eval_fold_valid {s : prog.State} {bss : List BasicStmt}
    {vs : VariableValues L prog} {ρ₁ ρ₂ : Env}
    (hbss : EvalBasicStmts ρ₁ bss ρ₂) (hvs : ⟦ vs ⟧ ρ₁) :
    ⟦ bss.foldl (fun vs bs => E.eval s bs vs) vs ⟧ ρ₂ := by
  induction hbss generalizing vs with
  | nil => exact hvs
  | cons hbs _ ih => exact ih (ValidStmtEvaluator.valid hbs hvs)

omit [FiniteHeightLattice L] [DecidableEq L] in
theorem updateVariablesForState_matches {s : prog.State}
    {sv : StateVariables L prog} {ρ₁ ρ₂ : Env}
    (hbss : EvalBasicStmts ρ₁ (prog.code s) ρ₂)
    (hvs : ⟦ variablesAt s sv ⟧ ρ₁) :
    ⟦ updateVariablesForState s sv ⟧ ρ₂ :=
  eval_fold_valid hbss hvs

omit [FiniteHeightLattice L] [DecidableEq L] in
theorem updateAll_matches {s : prog.State} {sv : StateVariables L prog}
    {ρ₁ ρ₂ : Env} (hbss : EvalBasicStmts ρ₁ (prog.code s) ρ₂)
    (hvs : ⟦ variablesAt s sv ⟧ ρ₁) :
    ⟦ variablesAt s (updateAll sv) ⟧ ρ₂ := by
  rw [variablesAt_updateAll]
  exact updateVariablesForState_matches hbss hvs

theorem stepTrace {s₁ : prog.State} {ρ₁ ρ₂ : Env}
    (hjoin : ⟦ joinForKey s₁ (result L prog) ⟧ ρ₁)
    (hbss : EvalBasicStmts ρ₁ (prog.code s₁) ρ₂) :
    ⟦ variablesAt s₁ (result L prog) ⟧ ρ₂ := by
  rw [result_eq L prog]
  refine updateAll_matches hbss ?_
  rw [variablesAt_joinAll]
  exact hjoin

theorem walkTrace {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env}
    (hjoin : ⟦ joinForKey s₁ (result L prog) ⟧ ρ₁)
    (tr : Trace prog.graph s₁ s₂ ρ₁ ρ₂) :
    ⟦ variablesAt s₂ (result L prog) ⟧ ρ₂ := by
  induction tr with
  | single hbss => exact stepTrace hjoin hbss
  | @edge _ ρ' _ i₁ i₂ _ hbss hedge _ ih =>
    have hstep : ⟦ variablesAt i₁ (result L prog) ⟧ ρ' :=
      stepTrace hjoin hbss
    have hmem : variablesAt i₁ (result L prog)
        ∈ (result L prog).valuesAt (prog.incoming i₂) :=
      FiniteMap.mem_valuesAt prog.states_nodup
        (prog.mem_incoming_of_edge hedge) (variablesAt_mem i₁ (result L prog))
    exact ih (interp_foldr hstep hmem)

omit V in
theorem interp_joinForKey_initialState :
    ⟦ joinForKey prog.initialState (result L prog) ⟧ [] := by
  rw [joinForKey_initialState]
  exact interp_botV_nil

variable (L prog) in
theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
    ⟦ variablesAt prog.finalState (result L prog) ⟧ ρ :=
  walkTrace interp_joinForKey_initialState (prog.trace hrun)

end Spa