import RoleFlow.Density
import RoleFlow.Infer

/-! # The inverse program (`docs/roleflow.tex`, `sec:inverse`)

The schedule is a forward closure (`fire`) over the recipes of the lines (`Prog.xrecs`):
starting from the weights and the outputs, it fires a recipe whose sources are all known, until
none is left, and emits each fired recipe as a line of the inverse program. A renaming maps the variables of `P` to those of the inverse
program. Firing keeps the renaming sort-preserving (`Good`) and, on a trace of `P`, the values
of the inverse program equal to those of the trace (`Agree`), since the recipes hold
(`xrecs_holds`). At the end the known set is
closed under the recipes (`fire_closed`), so by `Recov` at `Ret` it contains every input
(`Typed.known_xrecs`, `prop:schedule`).

`Prog.inverse` is the inverse program `P†`, a program of the same signature
(`Prog.inverse_wf`). **Concordance** (`Prog.inverse_sem`, `thm:conc`):
`⟦P†⟧_θ ∘ ⟦P⟧_θ = id`. With the right law, `⟦P⟧_θ` is bijective with inverse `⟦P†⟧_θ`
(`Prog.bijective`, `Prog.sem_inverse`, `Prog.equiv`), and `⟦P†⟧_θ` pushes `μ_O` back to `μ_I`
with the negated log-weight (`Prog.inverse_density`, `thm:invdensity`), which needs only the left
law.

Concordance needs less than the typing: `Prog.Invertible` asks only that the schedule reaches
every input, and the recipes hold under any roles. A program that is invertible but not typed
(say `ret x + b, x + b`) gets a left inverse, without a density. -/

open MeasureTheory

namespace RoleFlow

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

/-- **Existence of the schedule**, the recoverability half: under `Recov`, every variable, in
scope or introduced by the lines, is known from the weights and the returned variables. -/
theorem Typed.known_xrecs {nw : ℕ} {O Γ : List S} {Λ : Finset ℕ} {R : List (XRecipe Sg)}
    {D : List (Finset ℕ)} {ls : List (Line Sg)} {xs : List ℕ} {ρ : List Role}
    (hT : Typed Sg O Γ Λ D ls xs ρ) (hrec : Recov nw Γ.length Λ R D)
    (hX : ∀ x ∈ xrecs Γ.length ls ρ, x ∈ R) :
    ∀ v < Γ.length + (ls.map fun l => l.bs.length).sum, Known R ({w | w < nw} ∪ ↑xs.toFinset) v := by
  refine Typed.induction_recov (nw := nw)
    (motive := fun Γ _ ls xs _ => ∀ v < Γ.length + (ls.map fun l => l.bs.length).sum,
      Known R ({w | w < nw} ∪ ↑xs.toFinset) v)
    ?_ ?_ ?_ hT hrec hX
  · intro Γ Λ D xs hR _ hxs _ v hv
    simp only [List.map_nil, List.sum_nil, Nat.add_zero] at hv
    exact (hR.2 v hv).2.mono
      (Set.union_subset_union_right _ (by rw [hxs]; exact Finset.coe_subset.2 (hR.2 v hv).1))
  · intro _ _ _ l ls _ _ _ _ ih v hv
    exact ih v (by simp only [List.map_cons, List.sum_cons] at hv; simp; omega)
  · intro _ _ _ l ls _ _ _ _ _ _ _ _ _ ih v hv
    exact ih v (by simp only [List.map_cons, List.sum_cons] at hv; simp; omega)

/-! ## The forward closure -/

/-- The first element of a list with a property, and the list without it. -/
def pickFirst {α : Type*} (p : α → Bool) : List α → Option (α × List α)
  | [] => none
  | a :: as => if p a then some (a, as) else (pickFirst p as).map fun r => (r.1, a :: r.2)

theorem pickFirst_none {α : Type*} {p : α → Bool} {as : List α} (h : pickFirst p as = none) :
    ∀ a ∈ as, p a = false := by
  induction as with
  | nil => simp
  | cons b as ih =>
    simp only [pickFirst] at h
    split_ifs at h with hb
    simp only [Option.map_eq_none_iff] at h
    intro a ha
    rcases List.mem_cons.1 ha with rfl | ha
    · simpa using hb
    · exact ih h a ha

theorem pickFirst_some {α : Type*} {p : α → Bool} {as bs : List α} {a : α}
    (h : pickFirst p as = some (a, bs)) :
    p a = true ∧ a ∈ as ∧ bs.length + 1 = as.length ∧ (∀ b ∈ bs, b ∈ as) ∧
      ∀ b ∈ as, b = a ∨ b ∈ bs := by
  induction as generalizing bs with
  | nil => simp [pickFirst] at h
  | cons c as ih =>
    simp only [pickFirst] at h
    split_ifs at h with hc
    · cases h
      exact ⟨hc, List.mem_cons_self .., rfl, fun b hb => List.mem_cons_of_mem _ hb,
        fun b hb => List.mem_cons.1 hb⟩
    · obtain ⟨⟨a', bs'⟩, h', e⟩ := Option.map_eq_some_iff.1 h
      cases e
      obtain ⟨h₁, h₂, h₃, h₄, h₅⟩ := ih h'
      refine ⟨h₁, List.mem_cons_of_mem _ h₂, by simp [h₃], fun b hb => ?_, fun b hb => ?_⟩
      · rcases List.mem_cons.1 hb with rfl | hb
        · exact List.mem_cons_self ..
        · exact List.mem_cons_of_mem _ (h₄ b hb)
      · rcases List.mem_cons.1 hb with rfl | hb
        · exact Or.inr (List.mem_cons_self ..)
        · exact (h₅ b hb).imp id (List.mem_cons_of_mem _)

/-- The renaming after firing a recipe with targets `tgt` at the next free position `m`. -/
def renStep (tgt : List ℕ) (m : ℕ) (ren : ℕ → ℕ) (v : ℕ) : ℕ :=
  if v ∈ tgt then m + tgt.idxOf v else ren v

/-- **The schedule**, as a forward closure: fire a recipe of `X` whose sources are all known,
emit it as a line at the next free position `m`, and continue; `ren` renames the variables of
`P` into the inverse program and `K` is the hash set of known variables. -/
def fire : ℕ → List (XRecipe Sg) → ℕ → (ℕ → ℕ) → Std.HashSet ℕ →
    List (Line Sg) × (ℕ → ℕ) × Std.HashSet ℕ
  | 0, _, _, ren, K => ([], ren, K)
  | f + 1, X, m, ren, K =>
    match pickFirst (fun x => x.line.args.all K.contains) X with
    | none => ([], ren, K)
    | some (x, X') =>
      let r := fire f X' (m + x.line.bs.length) (renStep x.tgt m ren) (insertAll K x.tgt)
      (⟨x.line.op, x.line.args.map ren⟩ :: r.1, r.2)

theorem fire_mono : ∀ (f : ℕ) (X : List (XRecipe Sg)) (m : ℕ) (ren : ℕ → ℕ) (K : Std.HashSet ℕ),
    ∀ v ∈ K, v ∈ (fire f X m ren K).2.2
  | 0, _, _, _, _ => fun _ h => h
  | f + 1, X, m, ren, K => by
    simp only [fire]
    split
    · exact fun _ h => h
    · exact fun v h => fire_mono f _ _ _ _ v (mem_insertAll.2 (.inr h))

/-- After enough rounds, the known set is closed under the recipes. -/
theorem fire_closed : ∀ (f : ℕ) (X : List (XRecipe Sg)) (m : ℕ) (ren : ℕ → ℕ) (K : Std.HashSet ℕ),
    X.length ≤ f → ∀ x ∈ X, (∀ a ∈ x.line.args, a ∈ (fire f X m ren K).2.2) →
      ∀ t ∈ x.tgt, t ∈ (fire f X m ren K).2.2
  | 0, X, _, _, _, hX, x, hx => by simp_all
  | f + 1, X, m, ren, K, hX, x, hx => by
    simp only [fire]
    split
    · next h =>
      intro hargs
      have := pickFirst_none h x hx
      simp only [List.all_eq_false] at this
      obtain ⟨a, ha, hn⟩ := this
      exact absurd (hargs a ha) (by simpa using hn)
    · next y X' h =>
      obtain ⟨_, _, hlen, _, hsplit⟩ := pickFirst_some h
      rcases hsplit x hx with rfl | hx'
      · intro _ t ht
        exact fire_mono _ _ _ _ _ _ (mem_insertAll.2 (.inl ht))
      · exact fire_closed f X' _ _ _ (by omega) x hx'

/-- The inverse program's known set contains everything `Known` from its start. -/
theorem known_fire {f : ℕ} {X : List (XRecipe Sg)} {m : ℕ} {ren : ℕ → ℕ} {K : Std.HashSet ℕ}
    (hf : X.length ≤ f) {B : Set ℕ} (hB : ∀ v ∈ B, v ∈ K) {v : ℕ}
    (hv : Known X B v) : v ∈ (fire f X m ren K).2.2 := by
  induction hv with
  | base hv => exact fire_mono _ _ _ _ _ _ (hB _ hv)
  | fire hx _ hv ih => exact fire_closed f X m ren K hf _ hx ih _ hv

/-- The converse of `known_fire`: everything the schedule adds is `Known` from its start. -/
theorem fire_known {X₀ : List (XRecipe Sg)} {B : Set ℕ} :
    ∀ (f : ℕ) (X : List (XRecipe Sg)) (m : ℕ) (ren : ℕ → ℕ) (K : Std.HashSet ℕ),
      (∀ x ∈ X, x ∈ X₀) → (∀ v ∈ K, Known X₀ B v) → ∀ v ∈ (fire f X m ren K).2.2, Known X₀ B v
  | 0, _, _, _, _, _, hK => hK
  | f + 1, X, m, ren, K, hX, hK => by
    simp only [fire]
    split
    · exact hK
    · next x X' h =>
      obtain ⟨hp, hx, _, hsub, _⟩ := pickFirst_some h
      have hargs : ∀ a ∈ x.line.args, a ∈ K := by simpa using hp
      refine fire_known f X' _ _ _ (fun y hy => hX y (hsub y hy)) fun v hv => ?_
      rcases mem_insertAll.1 hv with hv | hv
      · exact .fire (hX x hx) (fun w hw => hK w (hargs w hw)) hv
      · exact hK v hv

/-! ## Sorts and values along the schedule -/

/-- The renaming `ren` takes the known variables of `Γ` to positions of `Γ'` of the same sort. -/
def Good (Γ Γ' : List S) (ren : ℕ → ℕ) (K : Std.HashSet ℕ) : Prop :=
  ∀ v ∈ K, v < Γ.length ∧ Γ'[ren v]? = Γ[v]?

theorem Good.args {Γ Γ' : List S} {ren : ℕ → ℕ} {K : Std.HashSet ℕ} (hG : Good Γ Γ' ren K)
    {xs : List ℕ} {as : List S} (h : Args Γ xs as) (hK : ∀ x ∈ xs, x ∈ K) :
    Args Γ' (xs.map ren) as := by
  induction h with
  | nil => exact List.Forall₂.nil
  | cons h₁ _ ih =>
    refine List.Forall₂.cons ?_ (ih fun x hx => hK x (List.mem_cons_of_mem _ hx))
    rw [(hG _ (hK _ (List.mem_cons_self ..))).2]; exact h₁

theorem lt_of_good {Γ Γ' : List S} {ren : ℕ → ℕ} {K : Std.HashSet ℕ} (hG : Good Γ Γ' ren K) {v : ℕ}
    (hv : v ∈ K) : ren v < Γ'.length := by
  have := (hG v hv).2
  rw [List.getElem?_eq_getElem (hG v hv).1] at this
  exact (List.getElem?_eq_some_iff.1 this).1

theorem Good.step {Γ Γ' : List S} {ren : ℕ → ℕ} {K : Std.HashSet ℕ} (hG : Good Γ Γ' ren K)
    {tgt : List ℕ} {bs : List S} (h : Args Γ tgt bs) :
    Good Γ (Γ' ++ bs) (renStep tgt Γ'.length ren) (insertAll K tgt) := by
  intro v hv
  by_cases ht : v ∈ tgt
  · have hi := List.idxOf_lt_length_of_mem ht
    have hi' : tgt.idxOf v < bs.length := h.length_eq ▸ hi
    refine ⟨h.lt ht, ?_⟩
    rw [renStep, if_pos ht, List.getElem?_append_right (by omega), Nat.add_sub_cancel_left,
      List.getElem?_eq_getElem hi']
    have := h.getElem? _ hi
    rw [List.getElem_idxOf hi] at this
    exact this.symm
  · have hK : v ∈ K := (mem_insertAll.1 hv).resolve_left ht
    rw [renStep, if_neg ht, List.getElem?_append_left (lt_of_good hG hK)]
    exact hG v hK

section Agree
variable [Sorts S]

/-- On known variables, the value `u` of the inverse program agrees with the trace `t`. -/
def Agree {Γ Γ' : List S} (t : Val Γ) (u : Val Γ') (ren : ℕ → ℕ) (K : Std.HashSet ℕ) : Prop :=
  ∀ v ∈ K, ∀ (a : S) (h : Γ[v]? = some a) (h' : Γ'[ren v]? = some a), u.get (ren v) h' = t.get v h

theorem Agree.pick {Γ Γ' : List S} {t : Val Γ} {u : Val Γ'} {ren : ℕ → ℕ} {K : Std.HashSet ℕ}
    (hA : Agree t u ren K) {xs : List ℕ} {as : List S} (h : Args Γ xs as)
    (h' : Args Γ' (xs.map ren) as) (hK : ∀ x ∈ xs, x ∈ K) : u.pick _ h' = t.pick xs h := by
  ext j a hj
  rw [Val.get_pick, Val.get_pick]
  have key : ∀ (i k : ℕ), k ∈ xs → i = ren k → ∀ (hi : Γ'[i]? = some a) (hk : Γ[k]? = some a),
      u.get i hi = t.get k hk := by
    rintro _ k hk rfl hi hk'; exact hA k (hK k hk) a hk' hi
  exact key _ _ (List.getElem_mem _) (List.getElem_map ..) _ _

theorem Agree.step {Γ Γ' : List S} {t : Val Γ} {u : Val Γ'} {ren : ℕ → ℕ} {K : Std.HashSet ℕ}
    (hA : Agree t u ren K) (hG : Good Γ Γ' ren K) {tgt : List ℕ} {bs : List S}
    (h : Args Γ tgt bs) : Agree t (u.append (t.pick tgt h)) (renStep tgt Γ'.length ren)
      (insertAll K tgt) := by
  intro v hv a hva hva'
  by_cases ht : v ∈ tgt
  · have hi := List.idxOf_lt_length_of_mem ht
    have e : renStep tgt Γ'.length ren v = Γ'.length + tgt.idxOf v := by rw [renStep, if_pos ht]
    rw [Val.get_congr _ e _ (e ▸ hva'), Val.get_append_right, Val.get_pick]
    exact Val.get_congr _ (List.getElem_idxOf hi) _ _
  · have hK : v ∈ K := (mem_insertAll.1 hv).resolve_left ht
    have e : renStep tgt Γ'.length ren v = ren v := by rw [renStep, if_neg ht]
    have hlt := lt_of_good hG hK
    rw [Val.get_congr _ e _ (e ▸ hva'), Val.get_append_left _ _ hlt]
    exact hA v hK a hva _

end Agree

/-- The lines emitted by the schedule are well formed, and the renaming stays good. -/
theorem fire_wf {Γ : List S} {X₀ : List (XRecipe Sg)}
    (hX : ∀ x ∈ X₀, Args Γ x.line.args x.line.as ∧ Args Γ x.tgt x.line.bs) :
    ∀ (f : ℕ) (X : List (XRecipe Sg)), (∀ x ∈ X, x ∈ X₀) → ∀ (Γ' : List S) (m : ℕ),
      Γ'.length = m → ∀ (ren : ℕ → ℕ) (K : Std.HashSet ℕ), Good Γ Γ' ren K →
      LinesWF Γ' (fire f X m ren K).1 ∧
        Good Γ (ctxAfter Γ' (fire f X m ren K).1) (fire f X m ren K).2.1 (fire f X m ren K).2.2
  | 0, _, _, _, _, _, _, _, hG => ⟨trivial, hG⟩
  | f + 1, X, hX', Γ', m, hm, ren, K, hG => by
    subst hm
    simp only [fire]
    split
    · exact ⟨trivial, hG⟩
    · next x X' h =>
      obtain ⟨hp, hx, _, hsub, _⟩ := pickFirst_some h
      have hK : ∀ a ∈ x.line.args, a ∈ K := by simpa using hp
      have hwf := hX x (hX' x hx)
      have ih := fire_wf hX f X' (fun y hy => hX' y (hsub y hy)) (Γ' ++ x.line.bs) (Γ'.length + x.line.bs.length) (by simp)
        (renStep x.tgt Γ'.length ren) (insertAll K x.tgt) (hG.step hwf.2)
      exact ⟨⟨hG.args hwf.1 hK, ih.1⟩, ih.2⟩

section FireAgree
variable [Sorts S] (M : Model Sg)

/-- On a trace on which the recipes hold, the inverse program reproduces the trace on every
known variable. -/
theorem fire_agree {Γ : List S} (t : Val Γ) {X₀ : List (XRecipe Sg)}
    (hX : ∀ x ∈ X₀, Args Γ x.line.args x.line.as ∧ Args Γ x.tgt x.line.bs)
    (hH : ∀ x ∈ X₀, ∀ (h₁ : Args Γ x.line.args x.line.as) (h₂ : Args Γ x.tgt x.line.bs),
      M.sem x.line.op (t.pick _ h₁) = t.pick _ h₂) :
    ∀ (f : ℕ) (X : List (XRecipe Sg)), (∀ x ∈ X, x ∈ X₀) → ∀ (Γ' : List S) (m : ℕ),
      Γ'.length = m → ∀ (ren : ℕ → ℕ) (K : Std.HashSet ℕ), Good Γ Γ' ren K →
      ∀ (hwf : LinesWF Γ' (fire f X m ren K).1) (u : Val Γ'), Agree t u ren K →
      Agree t (trace M _ hwf u) (fire f X m ren K).2.1 (fire f X m ren K).2.2
  | 0, _, _, _, _, _, _, _, _, _, _, hA => hA
  | f + 1, X, hX', Γ', m, hm, ren, K, hG, hwf, u, hA => by
    subst hm
    generalize hr : fire (f + 1) X Γ'.length ren K = r at hwf ⊢
    rcases hp : pickFirst (fun x => x.line.args.all K.contains) X with _ | ⟨x, X'⟩
    · simp only [fire, hp] at hr; subst hr; exact hA
    · simp only [fire, hp] at hr; subst hr
      obtain ⟨hp, hx, _, hsub, _⟩ := pickFirst_some hp
      have hK : ∀ a ∈ x.line.args, a ∈ K := by simpa using hp
      have hx₀ := hX x (hX' x hx)
      have hstep : stepLine M ⟨x.line.op, x.line.args.map ren⟩ hwf.1 u =
          u.append (t.pick x.tgt hx₀.2) := by
        rw [stepLine]
        congr 1
        rw [hA.pick hx₀.1 _ hK]
        exact hH x (hX' x hx) _ _
      change Agree t (trace M _ hwf.2 (stepLine M ⟨x.line.op, x.line.args.map ren⟩ hwf.1 u)) _ _
      rw [hstep]
      exact fire_agree t hX hH f X' (fun y hy => hX' y (hsub y hy)) (Γ' ++ x.line.bs) (Γ'.length + x.line.bs.length) (by simp)
        (renStep x.tgt Γ'.length ren) (insertAll K x.tgt) (hG.step hx₀.2) hwf.2
        (u.append (t.pick x.tgt hx₀.2)) (hA.step hG hx₀.2)

end FireAgree

/-! ## The inverse program -/

namespace Prog

/-- The schedule of `P`: the forward closure from the weights and the outputs. The outputs
get the positions after the weights, so the renaming starts as `renStep P.ret P.nw id`. -/
def schedule (P : Prog Sg) (ρ : List Role) : List (Line Sg) × (ℕ → ℕ) × Std.HashSet ℕ :=
  fire (P.xrecs ρ).length (P.xrecs ρ) (P.nw + P.outs.length) (renStep P.ret P.nw id)
    (insertAll (insertAll ∅ (List.range P.nw)) P.ret)

/-- **The inverse program** `P† = (Θ, O, e†)`: it fires the schedule and returns the inputs. -/
def inverse (P : Prog Sg) (ρ : List Role) : Prog Sg where
  wts := P.wts
  ins := P.outs
  outs := P.ins
  lines := (P.schedule ρ).1
  ret := (List.range' P.nw P.ins.length).map (P.schedule ρ).2.1

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

theorem length_ctx₀_le (P : Prog Sg) : P.ctx₀.length ≤ P.ctx.length := by
  obtain ⟨Δ, e⟩ := ctxAfter_append P.ctx₀ P.lines
  show _ ≤ (ctxAfter P.ctx₀ P.lines).length; rw [e]; simp

/-- The weights are known at the start, at their own positions. -/
theorem good_wts (P : Prog Sg) : Good P.ctx P.wts id (insertAll ∅ (List.range P.nw)) :=
    fun v hv => by
  have hlen := P.length_ctx₀_le
  have hv : v < P.wts.length := by simpa [mem_insertAll] using hv
  refine ⟨by simp at hlen; omega, ?_⟩
  rw [id, getElem?_ctxAfter _ (by simp; omega), List.getElem?_append_left hv]

theorem good_init (h : P.WF) :
    Good P.ctx (P.wts ++ P.outs) (renStep P.ret P.nw id)
      (insertAll (insertAll ∅ (List.range P.nw)) P.ret) :=
  P.good_wts.step h.2

theorem inputs_known (hT : P.Typed ρ) : ∀ v < P.ctx₀.length, v ∈ (P.schedule ρ).2.2 := by
  intro v hv
  have hk := hT.known_xrecs (P.recov_init _) (fun _ hx => hx) v (Nat.lt_add_right _ hv)
  exact known_fire le_rfl (fun w hw => by simpa [mem_insertAll, or_comm] using hw) hk

theorem args_inputs : Args P.ctx (List.range' P.nw P.ins.length) P.ins :=
  (Args.range' P.wts P.ins).ctxAfter P.lines

theorem inputs_mem (hT : P.Typed ρ) : ∀ x ∈ List.range' P.nw P.ins.length,
    x ∈ (P.schedule ρ).2.2 := fun x hx =>
  P.inputs_known hT x (by simp only [List.mem_range'_1, Prog.nw] at hx; simp; omega)

/-- `P` is invertible under the roles `ρ`: it is well formed and the schedule reaches every
input. Decidable, and implied by the typing (`Typed.invertible`). -/
def Invertible (P : Prog Sg) (ρ : List Role) : Prop :=
  P.WF ∧ ∀ x ∈ List.range' P.nw P.ins.length, x ∈ (P.schedule ρ).2.2

instance [DecidableEq S] (P : Prog Sg) (ρ : List Role) : Decidable (P.Invertible ρ) :=
  inferInstanceAs (Decidable (_ ∧ _))

theorem Typed.invertible (hT : P.Typed ρ) : P.Invertible ρ := ⟨hT.wf, inputs_mem hT⟩

/-- **Invertibility is `Known`**: the schedule reaches exactly the variables that are `Known`
from the weights and the returned variables (`known_fire`, `fire_known`). -/
theorem invertible_iff_known : P.Invertible ρ ↔ P.WF ∧ ∀ x ∈ List.range' P.nw P.ins.length,
    Known (P.xrecs ρ) ({w | w < P.nw} ∪ ↑P.ret.toFinset) x := by
  have hB : ∀ v, v ∈ ({w | w < P.nw} ∪ ↑P.ret.toFinset : Set ℕ) ↔
      v ∈ insertAll (insertAll ∅ (List.range P.nw)) P.ret := fun v => by
    simp [mem_insertAll, or_comm]
  refine and_congr_right fun _ => forall₂_congr fun x _ => ⟨fun h => ?_, fun h => ?_⟩
  · exact fire_known _ _ _ _ _ (fun _ hx => hx) (fun v hv => .base ((hB v).2 hv)) x h
  · exact known_fire le_rfl (fun v hv => (hB v).1 hv) h

/-- **The inverse program is well formed** (a program over the same signature). -/
theorem Invertible.inverse_wf (hI : P.Invertible ρ) : (P.inverse ρ).WF := by
  have hwf := fire_wf (wf_xrecs (ρ := ρ) hI.1) (P.xrecs ρ).length (P.xrecs ρ) (fun _ hx => hx) _
    (P.nw + P.outs.length)
    (by simp) _ _ (good_init hI.1)
  exact ⟨hwf.1, hwf.2.args args_inputs hI.2⟩

theorem inverse_wf (hT : P.Typed ρ) : (P.inverse ρ).WF := hT.invertible.inverse_wf

section Sem
variable [Sorts S] (M : Model Sg)

/-- At the start, the weights of the inverse program agree with the trace. -/
theorem agree_wts (h : P.WF) (θ : Val P.wts) (x : Val P.ins) :
    Agree (P.tr M h θ x) θ id (insertAll ∅ (List.range P.nw)) := fun v hv a ha ha' => by
  have hv : v < P.wts.length := by simpa [mem_insertAll] using hv
  show θ.get v ha' = (trace M P.lines h.1 (θ.append x)).get v ha
  rw [get_trace M _ _ _ _ (by simp; omega), Val.get_append_left _ _ hv]

/-- **Concordance** (`thm:conc`): the inverse program undoes the program. -/
theorem Invertible.inverse_sem (hI : P.Invertible ρ) (θ : Val P.wts) (x : Val P.ins) :
    (P.inverse ρ).sem M hI.inverse_wf θ (P.sem M hI.1 θ x) = x := by
  have h := hI.1
  set t := P.tr M h θ x
  have hH : ∀ y ∈ P.xrecs ρ, ∀ (h₁ : Args P.ctx y.line.args y.line.as)
      (h₂ : Args P.ctx y.tgt y.line.bs), M.sem y.line.op (t.pick _ h₁) = t.pick _ h₂ :=
    xrecs_holds M h.1 ρ rfl (θ.append x)
  have hA : Agree t (trace M (P.schedule ρ).1 hI.inverse_wf.1 (θ.append (t.pick P.ret h.2)))
      (P.schedule ρ).2.1 (P.schedule ρ).2.2 :=
    fire_agree M t (wf_xrecs (ρ := ρ) h) hH (P.xrecs ρ).length (P.xrecs ρ) (fun _ hx => hx) _
      (P.nw + P.outs.length) (by simp) _ _ (good_init h) hI.inverse_wf.1 _
      ((agree_wts M h θ x).step P.good_wts h.2)
  calc (P.inverse ρ).sem M hI.inverse_wf θ (P.sem M h θ x)
      = (trace M (P.schedule ρ).1 hI.inverse_wf.1 (θ.append (t.pick P.ret h.2))).pick
          ((List.range' P.nw P.ins.length).map (P.schedule ρ).2.1) hI.inverse_wf.2 := rfl
    _ = t.pick _ args_inputs := hA.pick args_inputs _ hI.2
    _ = x := (pick_trace M P.lines h.1 (θ.append x) (Args.range' P.wts P.ins) _).trans
      (Val.pick_range'_append θ x)

theorem inverse_sem (hT : P.Typed ρ) (θ : Val P.wts) (x : Val P.ins) :
    (P.inverse ρ).sem M (inverse_wf hT) θ (P.sem M hT.wf θ x) = x :=
  hT.invertible.inverse_sem M θ x

/-- The inverse program is measurable at fixed weights. -/
theorem Invertible.measurable_inverse (hI : P.Invertible ρ) (θ : Val P.wts) :
    Measurable ((P.inverse ρ).sem M hI.inverse_wf θ) :=
  ((P.inverse ρ).measurable_sem M hI.inverse_wf).comp (measurable_const.prodMk measurable_id)

theorem measurable_inverse (hT : P.Typed ρ) (θ : Val P.wts) :
    Measurable ((P.inverse ρ).sem M (inverse_wf hT) θ) :=
  hT.invertible.measurable_inverse M θ

/-- **Density of the inverse** (`thm:invdensity`): `⟦P†⟧_θ` pushes `μ_O` back to
`μ_I` with the negated log-weight at the recomputed inputs. Only the left law is used. -/
theorem inverse_density (hT : P.Typed ρ) (θ : Val P.wts) :
    Cov (vol P.outs) (vol P.ins) ((P.inverse ρ).sem M (inverse_wf hT) θ)
      fun y => expw (-P.logw M hT.wf ρ θ ((P.inverse ρ).sem M (inverse_wf hT) θ y)) :=
  Cov.inv (measurable_inverse M hT θ)
    ((P.measurable_logw M hT.wf ρ).comp (measurable_const.prodMk measurable_id))
    (inverse_sem M hT θ) (P.density M hT θ)

/-- **Model density** (`cor:modeldensity`): `e^{L_θ}` is the density of `(⟦P†⟧_θ)_* μ_O` with
respect to `μ_I`, by the density theorem and the left inverse. Here and not in `Density.lean`,
since it mentions `P†`. -/
theorem map_inverse (hT : P.Typed ρ) (θ : Val P.wts) :
    (vol P.outs).map ((P.inverse ρ).sem M (inverse_wf hT) θ) =
      (vol P.ins).withDensity fun x => expw (P.logw M hT.wf ρ θ x) :=
  Cov.map_eq (measurable_inverse M hT θ) (inverse_sem M hT θ) (P.density M hT θ)

variable {M}

/-- With the right law, the inverse program is also a right inverse. -/
theorem sem_inverse (hM : M.RightLaw) (hT : P.Typed ρ) (θ : Val P.wts) (y : Val P.outs) :
    P.sem M hT.wf θ ((P.inverse ρ).sem M (inverse_wf hT) θ y) = y := by
  obtain ⟨x, rfl⟩ := P.surjective M hM hT θ y
  rw [inverse_sem M hT]

/-- **Bijectivity** (`thm:bij`): with the right law, `⟦P⟧_θ` is a bijection. -/
theorem bijective (hM : M.RightLaw) (hT : P.Typed ρ) (θ : Val P.wts) :
    Function.Bijective (P.sem M hT.wf θ) :=
  ⟨Function.LeftInverse.injective (inverse_sem M hT θ), P.surjective M hM hT θ⟩

/-- `⟦P⟧_θ` as a measurable equivalence, with inverse `⟦P†⟧_θ`. -/
def equiv (hM : M.RightLaw) (hT : P.Typed ρ) (θ : Val P.wts) : Val P.ins ≃ᵐ Val P.outs where
  toFun := P.sem M hT.wf θ
  invFun := (P.inverse ρ).sem M (inverse_wf hT) θ
  left_inv := inverse_sem M hT θ
  right_inv := sem_inverse hM hT θ
  measurable_toFun := (P.measurable_sem M hT.wf).comp (measurable_const.prodMk measurable_id)
  measurable_invFun := measurable_inverse M hT θ

end Sem

end Prog

end RoleFlow