import Spa.Analysis.Forward
import Spa.Analysis.Utils
import Spa.Interp
import Spa.Showable

namespace Spa

abbrev ConstLattice : Type := AboveBelow ℤ

namespace ConstAnalysis

open AboveBelow in
def plus : ConstLattice → ConstLattice → ConstLattice
  | bot, _ => bot
  | _, bot => bot
  | top, _ => top
  | _, top => top
  | mk z₁, mk z₂ => mk (z₁ + z₂)

open AboveBelow in
def minus : ConstLattice → ConstLattice → ConstLattice
  | bot, _ => bot
  | _, bot => bot
  | top, _ => top
  | _, top => top
  | mk z₁, mk z₂ => mk (z₁ - z₂)

theorem plus_mono₂ : Monotone₂ plus :=
  AboveBelow.monotone₂_of_strict plus
    (fun y => by cases y <;> rfl) (fun x => by cases x <;> rfl)
    (fun y hy => by cases y <;> first | exact absurd rfl hy | rfl)
    (fun x hx => by cases x <;> first | exact absurd rfl hx | rfl)

theorem minus_mono₂ : Monotone₂ minus :=
  AboveBelow.monotone₂_of_strict minus
    (fun y => by cases y <;> rfl) (fun x => by cases x <;> rfl)
    (fun y hy => by cases y <;> first | exact absurd rfl hy | rfl)
    (fun x hx => by cases x <;> first | exact absurd rfl hx | rfl)

def interpConst : ConstLattice → Value → Prop
  | .bot, _ => False
  | .top, _ => True
  | .mk z, v => v = .int z

theorem interpConst_mk_disjoint {z₁ z₂ : ℤ} (hne : z₁ ≠ z₂) {v : Value} :
    ¬(interpConst (.mk z₁) v ∧ interpConst (.mk z₂) v) := by
  rintro ⟨h₁, h₂⟩
  rw [h₁] at h₂
  injection h₂ with hz
  exact hne hz

instance constInterpretation : LatticeInterpretation ConstLattice where
  interp := interpConst
  interp_sup := fun v h => AboveBelow.interp_sup_of (fun _ h => h) (fun _ => trivial) v h
  interp_inf := fun v h => AboveBelow.interp_inf_of (fun hne _ => interpConst_mk_disjoint hne) v h

variable (prog : Program)

def eval : Expr → VariableValues ConstLattice prog → ConstLattice
  | .add e₁ e₂, vs => plus (eval e₁ vs) (eval e₂ vs)
  | .sub e₁ e₂, vs => minus (eval e₁ vs) (eval e₂ vs)
  | .var k, vs =>
    if h : FiniteMap.MemKey k vs then (FiniteMap.locate h).1 else .top
  | .num n, _ => .mk n

theorem eval_mono (e : Expr) : Monotone (eval prog e) := by
  induction e with
  | add e₁ e₂ ih₁ ih₂ =>
    intro vs₁ vs₂ h
    exact eval_combine₂ plus_mono₂ (ih₁ h) (ih₂ h)
  | sub e₁ e₂ ih₁ ih₂ =>
    intro vs₁ vs₂ h
    exact eval_combine₂ minus_mono₂ (ih₁ h) (ih₂ h)
  | var k =>
    intro vs₁ vs₂ h
    simp only [eval]
    by_cases hk : k ∈ prog.vars
    · rw [dif_pos (FiniteMap.memKey_iff.mpr hk),
          dif_pos (FiniteMap.memKey_iff.mpr hk)]
      exact FiniteMap.le_of_mem_mem prog.vars_nodup h
        (FiniteMap.locate _).2 (FiniteMap.locate _).2
    · rw [dif_neg (fun hm => hk (FiniteMap.memKey_iff.mp hm)),
          dif_neg (fun hm => hk (FiniteMap.memKey_iff.mp hm))]
  | num n =>
    intro vs₁ vs₂ _
    exact le_refl _

instance exprEvaluator : ExprEvaluator ConstLattice prog :=
  ⟨eval prog, eval_mono prog⟩

def output : String :=
  show' (result ConstLattice prog)

theorem plus_valid {g₁ g₂ : ConstLattice} {z₁ z₂ : ℤ}
    (h₁ : ⟦g₁⟧ (.int z₁)) (h₂ : ⟦g₂⟧ (.int z₂)) :
    ⟦plus g₁ g₂⟧ (.int (z₁ + z₂)) := by
  rcases g₁ with _ | _ | c₁
  · exact h₁.elim
  · rcases g₂ with _ | _ | c₂
    · exact h₂.elim
    · exact trivial
    · exact trivial
  · rcases g₂ with _ | _ | c₂
    · exact h₂.elim
    · exact trivial
    · injection h₁ with hz₁
      injection h₂ with hz₂
      show Value.int (z₁ + z₂) = Value.int (c₁ + c₂)
      rw [hz₁, hz₂]

theorem minus_valid {g₁ g₂ : ConstLattice} {z₁ z₂ : ℤ}
    (h₁ : ⟦g₁⟧ (.int z₁)) (h₂ : ⟦g₂⟧ (.int z₂)) :
    ⟦minus g₁ g₂⟧ (.int (z₁ - z₂)) := by
  rcases g₁ with _ | _ | c₁
  · exact h₁.elim
  · rcases g₂ with _ | _ | c₂
    · exact h₂.elim
    · exact trivial
    · exact trivial
  · rcases g₂ with _ | _ | c₂
    · exact h₂.elim
    · exact trivial
    · injection h₁ with hz₁
      injection h₂ with hz₂
      show Value.int (z₁ - z₂) = Value.int (c₁ - c₂)
      rw [hz₁, hz₂]

instance eval_valid : ValidExprEvaluator ConstLattice prog := by
  constructor
  intro vs ρ e v hev
  induction hev with
  | num n =>
    intro _
    show ⟦eval prog (.num n) vs⟧ (.int n)
    rfl
  | var x v hxv =>
    intro hvs
    show ⟦eval prog (.var x) vs⟧ v
    simp only [eval]
    by_cases hk : FiniteMap.MemKey x vs
    · rw [dif_pos hk]
      exact hvs _ _ (FiniteMap.locate hk).2 _ hxv
    · rw [dif_neg hk]
      exact trivial
  | add e₁ e₂ z₁ z₂ _ _ ih₁ ih₂ =>
    intro hvs
    have h₁ : ⟦eval prog e₁ vs⟧ (.int z₁) := ih₁ hvs
    have h₂ : ⟦eval prog e₂ vs⟧ (.int z₂) := ih₂ hvs
    show ⟦eval prog (.add e₁ e₂) vs⟧ (.int (z₁ + z₂))
    exact plus_valid h₁ h₂
  | sub e₁ e₂ z₁ z₂ _ _ ih₁ ih₂ =>
    intro hvs
    have h₁ : ⟦eval prog e₁ vs⟧ (.int z₁) := ih₁ hvs
    have h₂ : ⟦eval prog e₂ vs⟧ (.int z₂) := ih₂ hvs
    show ⟦eval prog (.sub e₁ e₂) vs⟧ (.int (z₁ - z₂))
    exact minus_valid h₁ h₂

theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
    ⟦ variablesAt prog.finalState (result ConstLattice prog) ⟧ ρ :=
  Spa.analyze_correct ConstLattice prog hrun

end ConstAnalysis

end Spa