import Spa.Analysis.Forward.Evaluation

namespace Spa

namespace Forward

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

def updateVariablesFromExpression (k : String) (e : Expr)
    (vs : VariableValues L prog) : VariableValues L prog :=
  FiniteMap.generalizedUpdate id (fun _ vs => E.eval e vs) [k] vs

lemma updateVariablesFromExpression_mono (k : String) (e : Expr) :
    Monotone (updateVariablesFromExpression (L := L) (prog := prog) k e) :=
  FiniteMap.generalizedUpdate_monotone monotone_id (fun _ => E.eval_mono e)

def evalBasicStmt (bs : BasicStmt)
    (vs : VariableValues L prog) : VariableValues L prog :=
  match bs with
  | .assign k e => updateVariablesFromExpression k e vs
  | .noop => vs

lemma evalBasicStmt_mono (bs : BasicStmt) :
    Monotone (evalBasicStmt (L := L) (prog := prog) bs) := by
  cases bs with
  | assign k e => exact updateVariablesFromExpression_mono k e
  | noop => exact monotone_id

def evalBasicStmtOpt (obs : Option BasicStmt)
    (vs : VariableValues L prog) : VariableValues L prog :=
  match obs with
  | none => vs
  | some bs => evalBasicStmt bs vs

lemma evalBasicStmtOpt_mono (obs : Option BasicStmt) :
    Monotone (evalBasicStmtOpt (L := L) (prog := prog) obs) := by
  cases obs <;> unfold evalBasicStmtOpt
  · exact monotone_id
  · apply evalBasicStmt_mono

instance ExprEvaluator.toStmtEvaluator : StmtEvaluator L prog :=
  ⟨evalBasicStmtOpt ∘ prog.code,
   by intro s; simp; exact (evalBasicStmtOpt_mono (prog.code s))⟩

instance ExprEvaluator.toStmtEvaluator_valid [LatticeInterpretation L]
    [ValidExprEvaluator L prog] : ValidStmtEvaluator L prog := by
  constructor
  simp [StmtEvaluator.eval, evalBasicStmtOpt]
  intro s vs ρ₁ ρ₂; generalize prog.code s = obs; intro hev hvs
  rcases hev with _ | @⟨_,bs,hev⟩ <;> try simpa
  rcases hev with _ | @⟨k, e, v, hev⟩ <;> try simpa
  intros k' l' hkl' v' hρ
  rcases hρ with _ | ⟨_,_,_,_,_,hne,hmem⟩ <;> simp [evalBasicStmt] at hkl'
  · have hl := FiniteMap.generalizedUpdate_mem_eq (f := id)
        (g := fun _ vs => E.eval e vs) (List.mem_singleton_self k) hkl'
    rewrite [hl]; simp
    exact ValidExprEvaluator.valid hev hvs
  · have hl := FiniteMap.generalizedUpdate_not_mem_backward
        (fun hmem => hne (List.mem_singleton.mp hmem)) hkl'
    apply hvs _ _ hl _ hmem

end Forward

end Spa