import Spa.Language.Semantics

namespace Spa

/-- Environments agree on the current values of the selected variables. -/
def Env.AgreeOn (xs : Finset String) (ρ σ : Env) : Prop :=
  ∀ x ∈ xs, ∀ v, Env.Mem (x, v) ρ ↔ Env.Mem (x, v) σ

/-- Observable environment equality, ignoring shadowed bindings. -/
def Env.Equiv (ρ σ : Env) : Prop :=
  ∀ x v, Env.Mem (x, v) ρ ↔ Env.Mem (x, v) σ

lemma Env.Mem.functional {ρ : Env} {x : String} {v w : Value}
    (h : Env.Mem (x, v) ρ) (h' : Env.Mem (x, w) ρ) : v = w := by
  induction ρ with
  | nil => cases h
  | cons pair rest ih =>
      cases h <;> cases h' <;> aesop

lemma Env.mem_cons {ρ : Env} {x y : String} {v w : Value} :
    Env.Mem (x, v) ((y, w) :: ρ) ↔ (x = y ∧ v = w) ∨ (x ≠ y ∧ Env.Mem (x, v) ρ) := by
  constructor
  · intro h; cases h <;> aesop
  · rintro (⟨rfl, rfl⟩ | ⟨hne, h⟩)
    · exact .here _ _ _
    · exact .there _ _ _ _ _ hne h

lemma Env.Equiv.refl (ρ : Env) : Env.Equiv ρ ρ := fun _ _ => Iff.rfl
lemma Env.Equiv.symm {ρ σ : Env} (h : Env.Equiv ρ σ) : Env.Equiv σ ρ :=
  fun x v => (h x v).symm
lemma Env.Equiv.trans {ρ σ τ : Env} (h : Env.Equiv ρ σ) (h' : Env.Equiv σ τ) :
    Env.Equiv ρ τ := fun x v => (h x v).trans (h' x v)
lemma Env.Equiv.cons {ρ σ : Env} (h : Env.Equiv ρ σ) (x : String) (v : Value) :
    Env.Equiv ((x, v) :: ρ) ((x, v) :: σ) := by
  intro y w; simp only [Env.mem_cons, h y w]

lemma Env.cons_equiv_of_mem {ρ : Env} {x : String} {v : Value}
    (h : Env.Mem (x, v) ρ) : Env.Equiv ((x, v) :: ρ) ρ := by
  intro y w
  rw [Env.mem_cons]
  constructor
  · rintro (⟨rfl, rfl⟩ | ⟨_, hw⟩) <;> assumption
  · intro hw
    by_cases he : y = x
    · subst y; exact Or.inl ⟨rfl, hw.functional h⟩
    · exact Or.inr ⟨he, hw⟩

lemma EvalExpr.congr_env {ρ σ : Env} {e : Expr} {v : Value}
    (h : EvalExpr ρ e v) (ha : Env.AgreeOn e.vars ρ σ) : EvalExpr σ e v := by
  induction h with
  | num => exact .num _ _
  | var x v hm => exact .var _ _ _ ((ha x (by simp [Expr.vars]) v).mp hm)
  | add a b u v _ _ iha ihb =>
      exact .add _ _ _ _ _ (iha (fun x hx => ha x (Finset.mem_union_left _ hx)))
        (ihb (fun x hx => ha x (Finset.mem_union_right _ hx)))
  | sub a b u v _ _ iha ihb =>
      exact .sub _ _ _ _ _ (iha (fun x hx => ha x (Finset.mem_union_left _ hx)))
        (ihb (fun x hx => ha x (Finset.mem_union_right _ hx)))

lemma EvalExpr.deterministic {ρ : Env} {e : Expr} {v w : Value}
    (h : EvalExpr ρ e v) (h' : EvalExpr ρ e w) : v = w := by
  induction h generalizing w with
  | num => cases h'; rfl
  | var _ _ hm => cases h' with | var _ _ hm' => exact hm.functional hm'
  | add a b u v h₁ h₂ ih₁ ih₂ =>
      cases h' with
      | add _ _ u' v' h₁' h₂' =>
          have := ih₁ h₁'; have := ih₂ h₂'; aesop
  | sub a b u v h₁ h₂ ih₁ ih₂ =>
      cases h' with
      | sub _ _ u' v' h₁' h₂' =>
          have := ih₁ h₁'; have := ih₂ h₂'; aesop

/-- Variables a statement may assign. -/
def Stmt.writes : Stmt → Finset String
  | .basic .noop => ∅
  | .basic (.assign x _) => {x}
  | .andThen a b => a.writes ∪ b.writes
  | .ifElse _ a b => a.writes ∪ b.writes
  | .whileLoop _ b => b.writes

lemma EvalStmt.preserves_unwritten {ρ σ : Env} {s : Stmt} (h : EvalStmt ρ s σ)
    {x : String} (hx : x ∉ s.writes) :
    ∀ v, Env.Mem (x, v) ρ ↔ Env.Mem (x, v) σ := by
  induction h with
  | basic _ _ _ hb =>
      cases hb with
      | noop => exact fun _ => Iff.rfl
      | assign y _ _ _ =>
          have hn : x ≠ y := by simpa [Stmt.writes] using hx
          intro v; simp [Env.mem_cons, hn]
  | andThen _ _ _ _ _ _ _ ih₁ ih₂ =>
      simp only [Stmt.writes, Finset.mem_union, not_or] at hx
      exact fun v => (ih₁ hx.1 v).trans (ih₂ hx.2 v)
  | ifTrue _ _ _ _ _ _ _ _ _ ih =>
      exact ih (fun hm => hx (Finset.mem_union_left _ hm))
  | ifFalse _ _ _ _ _ _ _ ih =>
      exact ih (fun hm => hx (Finset.mem_union_right _ hm))
  | whileTrue _ _ _ _ _ _ _ _ _ _ ih₁ ih₂ =>
      exact fun v => (ih₁ hx v).trans (ih₂ hx v)
  | whileFalse => exact fun _ => Iff.rfl

/-- Evaluations depend on current bindings, not the list of shadowed bindings. -/
noncomputable def EvalStmt.congr_env {ρ ρ' : Env} {s : Stmt} (h : EvalStmt ρ s ρ') :
    ∀ {σ}, Env.Equiv ρ σ → Σ σ', {_h : EvalStmt σ s σ' // Env.Equiv ρ' σ'} := by
  induction h with
  | basic _ _ _ hb =>
      intro σ he
      cases hb with
      | noop => exact ⟨σ, .basic _ _ _ (.noop _), he⟩
      | assign x rhs v hv =>
          exact ⟨_, .basic _ _ _ (.assign _ _ _ _ (hv.congr_env (fun y _ => he y))), he.cons x v⟩
  | andThen _ _ _ _ _ _ _ ih₁ ih₂ =>
      intro σ he
      obtain ⟨σ₁, h₁, he₁⟩ := ih₁ he
      obtain ⟨σ₂, h₂, he₂⟩ := ih₂ he₁
      exact ⟨σ₂, .andThen _ _ _ _ _ h₁ h₂, he₂⟩
  | ifTrue _ _ _ _ _ _ hc hz _ ih =>
      intro σ he
      obtain ⟨σ', h', he'⟩ := ih he
      exact ⟨σ', .ifTrue _ _ _ _ _ _ (hc.congr_env (fun x _ => he x)) hz h', he'⟩
  | ifFalse _ _ _ _ _ hc _ ih =>
      intro σ he
      obtain ⟨σ', h', he'⟩ := ih he
      exact ⟨σ', .ifFalse _ _ _ _ _ (hc.congr_env (fun x _ => he x)) h', he'⟩
  | whileTrue _ _ _ _ _ _ hc hz _ _ ih₁ ih₂ =>
      intro σ he
      obtain ⟨σ₁, h₁, he₁⟩ := ih₁ he
      obtain ⟨σ₂, h₂, he₂⟩ := ih₂ he₁
      exact ⟨σ₂, .whileTrue _ _ _ _ _ _ (hc.congr_env (fun x _ => he x)) hz h₁ h₂, he₂⟩
  | whileFalse _ _ _ hc =>
      intro σ he
      exact ⟨σ, .whileFalse _ _ _ (hc.congr_env (fun x _ => he x)), he⟩

end Spa