import RoleFlow.Inverse

/-! # Invertibility as a type system

`Prog.Invertible` is a closure condition, decided by the schedule. Here it gets a syntax-directed
judgment `InvTyped`, one rule per line, that is sound and complete for it
(`Prog.invTyped_iff`). It annotates each variable `v` with a monotone Boolean formula `Φ v` over
variables, read as "`v` is recoverable from any set of variables that satisfies `Φ v`".
The supports of role typing are the special case where every formula is one conjunction.

* **Prog**: `Φ v = v`.
* **Cond** `ys = o(xs)`: each `y ∈ ys` becomes `y ∨ ⋀ Φ(xs)` (`condF`).
* **Spine** `ys = o(cs, ps)` with a partial inverse: every formula substitutes
  `c := c ∨ (⋀ ys ∧ ⋀ Φ(ps))` for `c ∈ cs` (`spineF`); without one, `Φ` is unchanged.
* **Ret** `ret xs`: each input's formula holds when the weights and `xs` are true.

Both updates are `upd Φ tgt args` for the line's recipe `tgt ← o'(args)`. The invariant is that
`Φ v` holds at `A` iff `v` is `Known` from `A` under the recipes of the lines so far
(`known_cons`, `eval_upd`); it needs no linearity, acyclicity or scoping.

Role typing is an instance (`Typed.sim`, `Prog.Typed.invTyped`): its invariant `Recov` says
`Θ ∪ D(v) ⊨ Φ(v)` (`recov_iff_eval`), so each support is a conjunction that implies the formula,
and the live set adds what role typing needs beyond invertibility. -/

namespace RoleFlow

variable {S : Type} {Sg : Sig S}

/-! ## Monotone formulas -/

/-- A monotone Boolean formula over variables. -/
inductive MF
  | var (v : ℕ)
  | tt
  | and (φ ψ : MF)
  | or (φ ψ : MF)

namespace MF

/-- `φ` holds when exactly the variables of `A` are true. -/
def eval (A : Set ℕ) : MF → Prop
  | var v => v ∈ A
  | tt => True
  | and φ ψ => eval A φ ∧ eval A ψ
  | or φ ψ => eval A φ ∨ eval A ψ

/-- Substitution of formulas for variables. -/
def bind (σ : ℕ → MF) : MF → MF
  | var v => σ v
  | tt => tt
  | and φ ψ => and (bind σ φ) (bind σ ψ)
  | or φ ψ => or (bind σ φ) (bind σ ψ)

/-- The conjunction of a list. -/
def all : List MF → MF
  | [] => tt
  | φ :: φs => and φ (all φs)

theorem eval_bind (A : Set ℕ) (σ : ℕ → MF) (φ : MF) :
    (φ.bind σ).eval A ↔ φ.eval {t | (σ t).eval A} := by
  induction φ <;> simp_all [bind, eval]

theorem eval_all (A : Set ℕ) (φs : List MF) : (all φs).eval A ↔ ∀ φ ∈ φs, φ.eval A := by
  induction φs <;> simp_all [all, eval]

end MF

/-- Firing the recipe `tgt ← o'(args)`: each target `t` becomes `t ∨ ⋀ Φ(args)` in every
formula. -/
def upd (Φ : ℕ → MF) (tgt args : List ℕ) : ℕ → MF := fun v =>
  (Φ v).bind fun t => if t ∈ tgt then .or (.var t) (.all (args.map Φ)) else .var t

theorem eval_upd (Φ : ℕ → MF) (tgt args : List ℕ) (A : Set ℕ) (v : ℕ) :
    (upd Φ tgt args v).eval A ↔
      (Φ v).eval (A ∪ {t | t ∈ tgt ∧ ∀ w ∈ args, (Φ w).eval A}) := by
  rw [upd, MF.eval_bind]
  congr! 1
  ext t
  by_cases ht : t ∈ tgt <;> simp [ht, MF.eval, MF.eval_all]

/-! ## Adding a recipe to `Known` -/

theorem Known.of_subset {R R' : List (XRecipe Sg)} (hR : ∀ r ∈ R, r ∈ R') {B : Set ℕ} {v : ℕ}
    (h : Known R B v) : Known R' B v := by
  induction h with
  | base hv => exact .base hv
  | fire hr _ hv ih => exact .fire (hR _ hr) ih hv

/-- **One more recipe**: `r` adds its targets to the base if its sources are known without it.
A firing of `r` that needs its own targets adds nothing, since `Known` is a least fixpoint. -/
theorem known_cons (r : XRecipe Sg) (R : List (XRecipe Sg)) (A : Set ℕ) (v : ℕ) :
    Known (r :: R) A v ↔
      Known R (A ∪ {t | t ∈ r.tgt ∧ ∀ w ∈ r.line.args, Known R A w}) v := by
  constructor
  · intro h
    induction h with
    | base hv => exact .base (.inl hv)
    | @fire r' _ hr' _ hv ih =>
      rcases List.mem_cons.1 hr' with rfl | hr'
      · by_cases hk : ∀ w ∈ r'.line.args, Known R A w
        · exact .base (.inr ⟨hv, hk⟩)
        · refine absurd (fun w hw => ?_) hk
          refine (ih w hw).trans fun u hu => ?_
          rcases hu with hu | ⟨-, hu⟩
          · exact .base hu
          · exact absurd hu hk
      · exact .fire hr' ih hv
  · intro h
    refine (h.of_subset fun _ => List.mem_cons_of_mem _).trans fun u hu => ?_
    rcases hu with hu | ⟨hu, hk⟩
    · exact .base hu
    · exact .fire (List.mem_cons_self ..) (fun w hw => (hk w hw).of_subset fun _ =>
        List.mem_cons_of_mem _) hu

/-- `upd` tracks `Known` when a recipe is added. -/
theorem upd_known {Φ : ℕ → MF} {R : List (XRecipe Sg)} (hΦ : ∀ A v, (Φ v).eval A ↔ Known R A v)
    (r : XRecipe Sg) : ∀ A v, (upd Φ r.tgt r.line.args v).eval A ↔ Known (r :: R) A v := by
  intro A v
  rw [eval_upd, hΦ, known_cons]
  simp only [hΦ]

/-! ## The judgment -/

/-- **Cond**: each result `y` becomes `y ∨ ⋀ Φ(xs)`. -/
def condF (Φ : ℕ → MF) (l : Line Sg) (n : ℕ) : ℕ → MF := upd Φ (l.ys n) l.args

/-- **Spine**, solving the first `k` arguments `cs` from the results `ys` and the parameters
`ps`: every formula substitutes `c := c ∨ (⋀ ys ∧ ⋀ Φ(ps))`. -/
def spineF (Φ : ℕ → MF) (l : Line Sg) (k n : ℕ) : ℕ → MF :=
  upd Φ (l.args.take k) (l.ys n ++ l.args.drop k)

/-- `Γ; Φ ⊢ e : O` with roles `ρ`, for the lines `ls`, the returned positions `xs`, `nw`
weights and the inputs `I`: **Ret**, **Cond**, **Spine** (with a partial inverse) and a spine
step without one, by recursion on the lines. Every rule checks the sorts of its line. -/
def InvTyped (Sg : Sig S) (O : List S) (nw : ℕ) (I : List ℕ) :
    List S → (ℕ → MF) → List (Line Sg) → List ℕ → List Role → Prop
  | Γ, Φ, [], xs, [] => Args Γ xs O ∧ ∀ i ∈ I, (Φ i).eval ({w | w < nw} ∪ ↑xs.toFinset)
  | Γ, Φ, l :: ls, xs, .cond :: ρ =>
    Args Γ l.args l.as ∧ InvTyped Sg O nw I (Γ ++ l.bs) (condF Φ l Γ.length) ls xs ρ
  | Γ, Φ, l :: ls, xs, .spine :: ρ =>
    Args Γ l.args l.as ∧
      match Sg.inv l.op with
      | some si => InvTyped Sg O nw I (Γ ++ l.bs) (spineF Φ l si.k Γ.length) ls xs ρ
      | none => InvTyped Sg O nw I (Γ ++ l.bs) Φ ls xs ρ
  | _, _, _, _, _ => False

/-- **Soundness and completeness of `InvTyped`**, for any context: if `Φ` tracks `Known` under
the recipes `R₀` of the lines before, then the judgment holds iff the lines and returned
variables are well formed, there is one role per line, and every input is `Known` from the
weights and the returned variables. -/
theorem invTyped_iff {O : List S} {nw : ℕ} {I : List ℕ} {xs : List ℕ} :
    ∀ {Γ : List S} {Φ : ℕ → MF} {R₀ : List (XRecipe Sg)} {ls : List (Line Sg)} {ρ : List Role},
      (∀ A v, (Φ v).eval A ↔ Known R₀ A v) →
      (InvTyped Sg O nw I Γ Φ ls xs ρ ↔
        (LinesWF Γ ls ∧ Args (ctxAfter Γ ls) xs O) ∧ ρ.length = ls.length ∧
          ∀ i ∈ I, Known (xrecs Γ.length ls ρ ++ R₀) ({w | w < nw} ∪ ↑xs.toFinset) i)
  | Γ, Φ, R₀, [], [], hΦ => by simp [InvTyped, LinesWF, ctxAfter, xrecs, hΦ]
  | _, _, _, [], _ :: _, _ => by simp [InvTyped]
  | _, _, _, _ :: _, [], _ => by simp [InvTyped]
  | Γ, Φ, R₀, l :: ls, .cond :: ρ, hΦ => by
    have ih := invTyped_iff (O := O) (nw := nw) (I := I) (xs := xs) (Γ := Γ ++ l.bs) (ls := ls) (ρ := ρ) (upd_known hΦ ⟨l, l.ys Γ.length⟩)
    have hR : ∀ B i, Known (xrecs (Γ ++ l.bs).length ls ρ ++ ⟨l, l.ys Γ.length⟩ :: R₀) B i ↔
        Known (xrecs Γ.length (l :: ls) (.cond :: ρ) ++ R₀) B i := fun B i =>
      ⟨Known.of_subset (by simp [xrecs_cond]; tauto), Known.of_subset (by simp [xrecs_cond]; tauto)⟩
    simp only [InvTyped, condF, ih, LinesWF, ctxAfter, hR, List.length_cons, add_left_inj]
    tauto
  | Γ, Φ, R₀, l :: ls, .spine :: ρ, hΦ => by
    cases hsi : Sg.inv l.op with
    | none =>
      have ih := invTyped_iff (O := O) (nw := nw) (I := I) (xs := xs) (Γ := Γ ++ l.bs) (ls := ls) (ρ := ρ) hΦ
      have e : xrecs Γ.length (l :: ls) (.spine :: ρ) = xrecs (Γ ++ l.bs).length ls ρ := by
        simp [xrecs, hsi]
      simp only [InvTyped, hsi, ih, LinesWF, ctxAfter, e, List.length_cons, add_left_inj]
      tauto
    | some si =>
      have ih := invTyped_iff (O := O) (nw := nw) (I := I) (xs := xs) (Γ := Γ ++ l.bs) (ls := ls) (ρ := ρ)
        (upd_known hΦ ⟨⟨si.op, l.ys Γ.length ++ l.args.drop si.k⟩, l.args.take si.k⟩)
      have hR : ∀ B i, Known (xrecs (Γ ++ l.bs).length ls ρ ++
            ⟨⟨si.op, l.ys Γ.length ++ l.args.drop si.k⟩, l.args.take si.k⟩ :: R₀) B i ↔
          Known (xrecs Γ.length (l :: ls) (.spine :: ρ) ++ R₀) B i := fun B i =>
        ⟨Known.of_subset (by simp [xrecs_spine _ _ _ _ hsi]; tauto),
          Known.of_subset (by simp [xrecs_spine _ _ _ _ hsi]; tauto)⟩
      simp only [InvTyped, hsi, spineF, ih, LinesWF, ctxAfter, hR, List.length_cons, add_left_inj]
      tauto

/-! ## Role typing as an instance -/

theorem MF.eval_mono {A B : Set ℕ} (h : A ⊆ B) : ∀ φ : MF, φ.eval A → φ.eval B
  | .var _, hv => h hv
  | .tt, _ => trivial
  | .and φ ψ, ⟨h₁, h₂⟩ => ⟨φ.eval_mono h h₁, ψ.eval_mono h h₂⟩
  | .or φ ψ, hv => hv.imp (φ.eval_mono h) (ψ.eval_mono h)

/-- **Role typing's invariant is a conjunction below `Φ`**: when `Φ` tracks `Known`, `Recov`
says that every variable's formula holds at the weights and its support, which is live. -/
theorem recov_iff_eval {nw n : ℕ} {Λ : Finset ℕ} {R : List (XRecipe Sg)} {D : List (Finset ℕ)}
    {Φ : ℕ → MF} (hΦ : ∀ A v, (Φ v).eval A ↔ Known R A v) :
    Recov nw n Λ R D ↔ D.length = n ∧
      ∀ v < n, D.getD v ∅ ⊆ Λ ∧ (Φ v).eval ({w | w < nw} ∪ ↑(D.getD v ∅)) := by
  simp only [Recov, hΦ]

theorem Recov.of_subset {nw n : ℕ} {Λ : Finset ℕ} {R R' : List (XRecipe Sg)}
    {D : List (Finset ℕ)} (hR' : ∀ r ∈ R, r ∈ R') (hR : Recov nw n Λ R D) : Recov nw n Λ R' D :=
  ⟨hR.1, fun v hv => ⟨(hR.2 v hv).1, (hR.2 v hv).2.of_subset hR'⟩⟩

/-- **Simulation**: a role typing whose annotations satisfy `Recov` (equivalently
`Θ ∪ D(v) ⊨ Φ(v)`, `recov_iff_eval`) for the recipes `R₀` of the lines before gives an
invertibility typing line by line: **Cond** and **Spine** become **I-Cond** and **I-Spine**,
`Recov` is preserved for the recipes so far, and at **Ret** the support of every input lies in
the live set, the returned variables. -/
theorem Typed.sim {O : List S} {nw : ℕ} {I : List ℕ} {xs : List ℕ} :
    ∀ {Γ : List S} {Λ : Finset ℕ} {D : List (Finset ℕ)} {ls : List (Line Sg)} {ρ : List Role}
      {Φ : ℕ → MF} {R₀ : List (XRecipe Sg)},
      Typed Sg O Γ Λ D ls xs ρ → (∀ A v, (Φ v).eval A ↔ Known R₀ A v) →
      Recov nw Γ.length Λ R₀ D → (∀ i ∈ I, i < Γ.length) → InvTyped Sg O nw I Γ Φ ls xs ρ
  | Γ, Λ, D, [], [], Φ, R₀, hT, hΦ, hR, hI => by
    refine ⟨hT.2.2, fun i hi => (hΦ _ i).2 ?_⟩
    refine (hR.2 i (hI i hi)).2.mono (Set.union_subset_union_right _ ?_)
    rw [hT.2.1]; exact Finset.coe_subset.2 (hR.2 i (hI i hi)).1
  | Γ, Λ, D, l :: ls, .cond :: ρ, Φ, R₀, ⟨ha, hT⟩, hΦ, hR, hI => by
    refine ⟨ha, Typed.sim hT (upd_known hΦ ⟨l, l.ys Γ.length⟩) ?_ fun i hi => by
      simp; have := hI i hi; omega⟩
    rw [List.length_append]
    exact (hR.of_subset fun _ => List.mem_cons_of_mem _).cond (fun _ => ha.lt)
      (List.mem_cons_self ..)
  | Γ, Λ, D, l :: ls, .spine :: ρ, Φ, R₀, ⟨ha, si, hsi, _, _, hacyc, hT⟩, hΦ, hR, hI => by
    simp only [InvTyped, hsi]
    refine ⟨ha, Typed.sim hT
      (upd_known hΦ ⟨⟨si.op, l.ys Γ.length ++ l.args.drop si.k⟩, l.args.take si.k⟩) ?_
      fun i hi => by simp; have := hI i hi; omega⟩
    rw [List.length_append]
    exact (hR.of_subset fun _ => List.mem_cons_of_mem _).spine (fun _ => ha.lt) hacyc
      (List.mem_cons_self ..)

/-- `⊢^inv_ρ P`: **Prog** starts every formula as its own variable. -/
def Prog.InvTyped (P : Prog Sg) (ρ : List Role) : Prop :=
  RoleFlow.InvTyped Sg P.outs P.nw (List.range' P.nw P.ins.length) P.ctx₀ MF.var P.lines P.ret ρ

namespace Prog

variable {P : Prog Sg} {ρ : List Role}

theorem known_nil (A : Set ℕ) (v : ℕ) : (MF.var v).eval A ↔ Known ([] : List (XRecipe Sg)) A v :=
  ⟨.base, fun h => by cases h with | base h => exact h | fire hr => simp at hr⟩

/-- **The judgment characterizes invertibility**: `⊢^inv_ρ P` iff `P` is invertible under `ρ`
and `ρ` has one role per line. -/
theorem invTyped_iff : P.InvTyped ρ ↔ P.Invertible ρ ∧ ρ.length = P.lines.length := by
  rw [Prog.InvTyped, RoleFlow.invTyped_iff known_nil, invertible_iff_known, Prog.WF, Prog.xrecs]
  simp only [List.append_nil]
  tauto

instance [DecidableEq S] (P : Prog Sg) (ρ : List Role) : Decidable (P.InvTyped ρ) :=
  decidable_of_iff _ invTyped_iff.symm

/-- **Role typing instantiates the invertibility judgment** (`Typed.sim` from **Prog**). -/
theorem Typed.invTyped (hT : P.Typed ρ) : P.InvTyped ρ :=
  RoleFlow.Typed.sim hT known_nil (P.recov_init []) fun i hi => by
    simp only [List.mem_range'_1, Prog.nw] at hi; simp; omega

end Prog

end RoleFlow