import RoleFlow.Typing
import RoleFlow.Live

/-! # The density theorem (`docs/roleflow.tex`, `sec:density`)

Fix the context `Γ` of all variables and a family of traces `T : Ω → Val Γ` (for a program at
fixed weights: the trace of each input). `Det T A v` says that `v` is *determined* by the
positions `A`: some measurable function of the live state on `A` gives its value on every
trace. A recipe *holds* (`Holds`) if its targets are determined by its sources.

**Factorization** (`det_of_known`, `Recov.det`, `Recov.det_params`): if all recipes hold and
the weights are determined by nothing (they are constant), then everything in
`Known R (Θ ∪ K)` is determined by `K`. With `Recov` this gives the recomputation `r_v` of
every variable in scope from the live state, and at a spine step that of the parameters from
the live variables the step does not consume.

**Density** (`density_aux`, by `Typed.induction_recov`): given that the recipes of the lines
hold (`xrecs_holds`), at each node of a derivation (`NodeInv`) there are a measurable
`F : X_Λ → X_O` and `L : X_Λ → ℝ` with `Cov (lvol Λ) (vol O) F e^L` that give the returned tuple and the remaining log-weight on every trace. At a spine step,
`F = F' ∘ step` with `step(c, q) = (q, ⟦o⟧(c, z q))`, where `z` recomputes the parameters from
`q`; its `Cov` is **Tri** with the leaf at `z q`, between measure-preserving splittings.
`Prog.density` reads this at the root, and `Prog.density_joint` is the joint statement over
the weights, which covers conditional flows. Only the left law is used. -/

open MeasureTheory ENNReal

namespace RoleFlow

variable {S : Type} [Sorts S]

/-! ## Determined positions -/

section Det
variable {Sg : Sig S} {Γ : List S} {Ω : Type*} (T : Ω → Val Γ)

/-- `v` is determined by the positions `A` on the traces `T`. -/
def Det (A : Finset ℕ) (v : ℕ) : Prop :=
  ∀ (a : S) (h : Γ[v]? = some a), ∃ g : LState Γ A → X a, Measurable g ∧
    ∀ ω, (T ω).get v h = g (LState.restrict A (T ω))

/-- A recipe holds on the traces `T`: its targets are determined by its sources. -/
def Holds (r : XRecipe Sg) : Prop := ∀ v ∈ r.tgt, Det T r.line.args.toFinset v

variable {T} {A B : Finset ℕ} {v : ℕ}

theorem Det.of_mem (hv : v ∈ A) : Det T A v := fun _ h =>
  ⟨fun s => s.get v hv h, LState.measurable_get hv h,
    fun _ => (LState.get_restrict _ hv h).symm⟩

/-- If every position of `B` is determined by `A`, so is the live state on `B`. -/
theorem Det.state (hB : ∀ w ∈ B, Det T A w) :
    ∃ G : LState Γ A → LState Γ B, Measurable G ∧
      ∀ ω, G (LState.restrict A (T ω)) = LState.restrict B (T ω) := by
  choose g hg hgT using hB
  refine ⟨fun s i => g i.1.1 i.2 _ (by simp) s, measurable_pi_lambda _ fun i => hg _ _ _ _,
    fun ω => funext fun i => (hgT _ _ _ _ ω).symm⟩

theorem Det.comp (hB : ∀ w ∈ B, Det T A w) (h : Det T B v) : Det T A v := fun a ha => by
  obtain ⟨G, hG, hGT⟩ := Det.state hB
  obtain ⟨g, hg, hgT⟩ := h a ha
  exact ⟨fun s => g (G s), hg.comp hG, fun ω => by rw [hgT]; exact congrArg g (hGT ω).symm⟩

/-- Determined positions can be picked by a measurable function of the live state. -/
theorem Det.pick {xs : List ℕ} {as : List S} (hxs : ∀ x ∈ xs, Det T A x) (h : Args Γ xs as) :
    ∃ Z : LState Γ A → Val as, Measurable Z ∧
      ∀ ω, Z (LState.restrict A (T ω)) = (T ω).pick xs h := by
  obtain ⟨G, hG, hGT⟩ := Det.state (B := xs.toFinset) fun w hw => hxs w (List.mem_toFinset.1 hw)
  refine ⟨fun s => (G s).pick xs (fun x hx => List.mem_toFinset.2 hx) h,
    (LState.measurable_pick _ h).comp hG, fun ω => ?_⟩
  dsimp only
  rw [hGT, LState.pick_restrict]

/-- If the values at `ys` are a measurable function of those at `xs` on every trace, then every
position of `ys` is determined by `xs`. -/
theorem det_of_pick {xs ys : List ℕ} {as bs : List S} (hx : Args Γ xs as) (hy : Args Γ ys bs)
    {G : Val as → Val bs} (hG : Measurable G)
    (hT : ∀ ω, (T ω).pick ys hy = G ((T ω).pick xs hx)) : ∀ v ∈ ys, Det T xs.toFinset v := by
  intro v hv a ha
  obtain ⟨j, hj, rfl⟩ := List.mem_iff_getElem.1 hv
  have hb : bs[j]? = some a :=
    List.getElem?_eq_some_iff.2 ⟨_, Option.some.inj ((hy.getElem? j hj).symm.trans ha)⟩
  refine ⟨fun s => (G (s.pick xs (fun x hx => List.mem_toFinset.2 hx) hx)).get j hb,
    (Val.measurable_get hb).comp (hG.comp (LState.measurable_pick _ hx)), fun ω => ?_⟩
  simp only [LState.pick_restrict]
  rw [← hT, Val.get_pick]

/-- **Factorization**: if the recipes hold and the weights are determined by nothing, everything
known from the weights and `K` is determined by `K`. -/
theorem det_of_known {nw : ℕ} {R : List (XRecipe Sg)} {K : Finset ℕ} (hR : ∀ r ∈ R, Holds T r)
    (hw : ∀ w < nw, Det T ∅ w) (hv : Known R ({w | w < nw} ∪ ↑K) v) : Det T K v := by
  induction hv with
  | base hv =>
    rcases hv with hv | hv
    · intro a ha
      obtain ⟨g, hg, hgT⟩ := hw _ hv a ha
      exact ⟨fun s => g (LState.sub (Finset.empty_subset _) s),
        hg.comp (LState.measurable_sub _), hgT⟩
    · exact Det.of_mem hv
  | fire hr _ hv ih => exact Det.comp (fun w hw => ih w (List.mem_toFinset.1 hw)) (hR _ hr _ hv)

variable {nw n : ℕ} {Λ : Finset ℕ} {R : List (XRecipe Sg)} {D : List (Finset ℕ)}

/-- **Factorization**, the live-state recomputation `r_v`: under `Recov`, every variable in
scope is determined by the live state. -/
theorem Recov.det (hrec : Recov nw n Λ R D) (hR : ∀ r ∈ R, Holds T r)
    (hw : ∀ w < nw, Det T ∅ w) (hv : v < n) : Det T Λ v :=
  det_of_known hR hw ((hrec.2 v hv).2.mono
    (Set.union_subset_union_right _ (Finset.coe_subset.2 (hrec.2 v hv).1)))

/-- **Factorization** at a spine step: the parameters factor through the live variables the
step does not consume, since their supports avoid them. -/
theorem Recov.det_params {l : Line Sg} {k p : ℕ} (hrec : Recov nw n Λ R D)
    (hR : ∀ r ∈ R, Holds T r) (hw : ∀ w < nw, Det T ∅ w) (hargs : ∀ x ∈ l.args, x < n)
    (hacyc : ∀ p ∈ l.args.drop k, ∀ c ∈ l.args.take k, c ∉ D.getD p ∅)
    (hp : p ∈ l.args.drop k) : Det T (Λ \ (l.args.take k).toFinset) p := by
  obtain ⟨hΛ, hK⟩ := hrec.2 p (hargs p (List.mem_of_mem_drop hp))
  refine det_of_known hR hw (hK.mono (Set.union_subset_union_right _
    (Finset.coe_subset.2 fun x hx => Finset.mem_sdiff.2 ⟨hΛ hx, fun hxc => ?_⟩)))
  exact hacyc p hp x (List.mem_toFinset.1 hxc) hx

end Det

/-! ## Density along a derivation -/

section Density
variable {Sg : Sig S} (M : Model Sg) {Γ outs : List S} {Ω : Type*} (T : Ω → Val Γ) (nw : ℕ)

/-- The invariant at a node of a derivation, with live set `Λ`, remaining lines `ls`, returned `xs` and roles `ρ`: a measurable `F` and `L` on live states with
`Cov (lvol Λ) (vol O) F e^L`, which give the returned tuple and the remaining log-weight on
every trace. The node starts in the context `Γ'`, and `pre ω` is the value of `Γ'` from which
the trace `T ω` continues. -/
def NodeInv (Γ' : List S) (Λ : Finset ℕ) (ls : List (Line Sg)) (xs : List ℕ) (ρ : List Role) :
    Prop :=
  ∀ (hwf : LinesWF Γ' ls) (hctx : ctxAfter Γ' ls = Γ) (hret : Args Γ xs outs) (pre : Ω → Val Γ'),
    (∀ ω, (trace M ls hwf (pre ω)).cast hctx = T ω) → (∀ v ∈ Λ, v < Γ'.length) →
    ∃ F : LState Γ Λ → Val outs, ∃ L : LState Γ Λ → ℝ, Measurable F ∧ Measurable L ∧
      Cov (lvol Γ Λ) (vol outs) F (fun s => expw (L s)) ∧ (M.RightLaw → Function.Surjective F) ∧
      ∀ ω, F (LState.restrict Λ (T ω)) = (T ω).pick xs hret ∧
        L (LState.restrict Λ (T ω)) = logwLines M ls hwf ρ (pre ω)

/-- **Density at every node** of a derivation, given recipes `R` that hold and contain those of
the lines. -/
theorem density_aux (hw : ∀ w < nw, Det T ∅ w) {R : List (XRecipe Sg)} (hR : ∀ r ∈ R, Holds T r)
    {O Γ' : List S} {Λ : Finset ℕ} {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) : NodeInv M T (outs := outs) Γ' Λ ls xs ρ := by
  refine Typed.induction_recov (nw := nw) (motive := NodeInv M T) ?_ ?_ ?_ hT hrec hX
  · -- Ret: a reorder
    intro Γ' Λ D xs _ hnd hxs _ hwf hctx hret pre _ _
    have hA : ∀ x ∈ xs, x ∈ Λ := fun x hx => hxs ▸ List.mem_toFinset.2 hx
    have hcov : ∀ x ∈ Λ, x ∈ xs := fun x hx => List.mem_toFinset.1 (hxs.symm ▸ hx)
    refine ⟨fun s => s.pick xs hA hret, fun _ => 0, LState.measurable_pick hA hret,
      measurable_const, ?_, fun _ => (LState.pickEquiv hA hret hnd hcov).surjective,
      fun ω => ⟨LState.pick_restrict _ _ _, rfl⟩⟩
    exact (Cov.pres (LState.measurePreserving_pick hA hret hnd hcov)).congr (fun _ => rfl)
      (fun _ => by simp)
  · -- Cond: the live set is unchanged
    intro Γ' Λ D l ls xs ρ _ _ ih hwf hctx hret pre hpre hΛ
    obtain ⟨F, L, hF, hL, hcov, hsurj, htr⟩ := ih hwf.2 hctx hret
      (fun ω => stepLine M l hwf.1 (pre ω)) hpre (fun v hv => by have := hΛ v hv; simp; omega)
    refine ⟨F, L, hF, hL, hcov, hsurj, fun ω => ⟨(htr ω).1, ?_⟩⟩
    rw [(htr ω).2]
    simp [logwLines]
  · -- Spine: Tri with the leaf at the recomputed parameters
    intro Γ' Λ D l ls xs ρ si hrec ha hsi hnd hcΛ hacyc ih hwf hctx hret pre hpre hΛ
    have hargs : ∀ x ∈ l.args, x < Γ'.length := fun _ => ha.lt
    have lift : ∀ {ys : List ℕ} {bs : List S}, Args Γ' ys bs → Args Γ ys bs := fun h => by
      rw [← hctx]; exact h.ctxAfter _
    have hx := hwf.1
    have hc := si.args_take hx
    have hp := si.args_drop hx
    have hy : Args Γ (l.ys Γ'.length) l.bs := by
      rw [← hctx]; exact (Args.range' Γ' l.bs).ctxAfter ls
    -- the trace before and at the line
    have node_pick : ∀ ω {ys : List ℕ} {bs : List S} (h : Args Γ' ys bs),
        (T ω).pick ys (lift h) = (pre ω).pick ys h := fun ω {_} {_} h => by
      rw [← hpre ω, Val.pick_cast]; exact pick_trace M _ hwf (pre ω) h _
    have node_res : ∀ ω, (T ω).pick _ hy = M.sem l.op ((pre ω).pick l.args hx) := fun ω => by
      rw [← hpre ω, Val.pick_cast]; exact pick_trace_results M l ls hwf (pre ω) _
    -- the live sets
    set c := l.args.take si.k
    set B := Λ \ c.toFinset
    set C := Finset.Ico Γ'.length (Γ'.length + l.bs.length)
    have hBC : ∀ i ∈ B ∪ C, i ∈ B ↔ i ∉ l.ys Γ'.length := fun i hi => by
      constructor
      · intro hb hy'
        have := hΛ i (Finset.mem_sdiff.1 hb).1
        have := (Line.mem_ys.1 hy').1
        omega
      · intro hny
        exact (Finset.mem_union.1 hi).resolve_right fun h =>
          hny (Line.mem_ys.2 (Finset.mem_Ico.1 h))
    -- the parameters, recomputed from the live variables not consumed
    obtain ⟨Z, hZ, hZT⟩ := Det.pick (fun p hp' => Recov.det_params hrec hR hw hargs hacyc hp')
      (lift hp)
    -- the splittings `X_Λ ≅ X_{Λ∖c} × X_c` and `X_{Λ'} ≅ X_{Λ∖c} × X_y`
    have hB : ∀ i ∈ Λ, i ∈ B ↔ i ∉ c := fun i hi => by simp [B, hi]
    let Φ := LState.extract Finset.sdiff_subset hcΛ (lift hc) hnd hB
    let Ψ := LState.extract Finset.subset_union_left
      (fun _ hx => Finset.mem_union_right _ (Finset.mem_Ico.2 (Line.mem_ys.1 hx))) hy
      List.nodup_range' hBC
    obtain ⟨F, L, hF, hL, hcov, hsurj, htr⟩ := ih hwf.2 hctx hret
      (fun ω => stepLine M l hwf.1 (pre ω)) hpre
      (fun v hv => by
        rcases Finset.mem_union.1 hv with hv | hv
        · have := hΛ v (Finset.mem_sdiff.1 hv).1; simp; omega
        · simpa using (Finset.mem_Ico.1 hv).2)
    have hY : Measurable fun q : LState Γ B × Val si.cs => M.sem l.op (si.glue q.2 (Z q.1)) :=
      (M.sem_meas l.op).comp (si.measurable_glue.comp (measurable_snd.prodMk (hZ.comp measurable_fst)))
    have hJ : Measurable fun q : LState Γ B × Val si.cs => M.ℓ l.op (si.glue q.2 (Z q.1)) :=
      (M.ℓ_meas l.op).comp (si.measurable_glue.comp (measurable_snd.prodMk (hZ.comp measurable_fst)))
    have htri := Cov.tri (μ := lvol Γ B) (g := fun q c => M.sem l.op (si.glue c (Z q)))
      (J := fun q c => expw (M.ℓ l.op (si.glue c (Z q)))) hY (measurable_expw.comp hJ)
      fun q => M.leaf l.op si hsi (Z q)
    -- the step `s ↦ (q, ⟦o⟧(c, z q))` into the next live state
    let step : LState Γ Λ → LState Γ (B ∪ C) :=
      fun s => Ψ.symm ((Φ s).1, M.sem l.op (si.glue (Φ s).2 (Z (Φ s).1)))
    have hstep : Measurable step :=
      Ψ.symm.measurable.comp ((measurable_fst.prodMk hY).comp Φ.measurable)
    refine ⟨fun s => F (step s),
      fun s => M.ℓ l.op (si.glue (Φ s).2 (Z (Φ s).1)) + L (step s),
      hF.comp hstep, (hJ.comp Φ.measurable).add (hL.comp hstep), ?_, fun hright o => ?_,
      fun ω => ?_⟩
    · exact (Cov.comp hF (measurable_expw.comp hL)
        ((Cov.pres_comp (LState.measurePreserving_extract ..) (measurable_fst.prodMk hY)
          (measurable_expw.comp hJ) htri).comp_pres ((LState.measurePreserving_extract ..).symm Ψ))
        hcov).congr (fun _ => rfl) (fun _ => (expw_add _ _).symm)
    · -- surjective, by the right law: solve the consumed arguments with the semi-inverse
      obtain ⟨s', rfl⟩ := hsurj hright o
      obtain ⟨⟨q, b⟩, rfl⟩ := Ψ.symm.surjective s'
      obtain ⟨s, hs⟩ := Φ.surjective (q, M.sem si.op (b.append (Z q)))
      exact ⟨s, by dsimp only [step]; rw [hs, hright l.op si hsi b (Z q)]⟩
    · have hΦT := LState.extract_restrict Finset.sdiff_subset hcΛ (lift hc) hnd hB (T ω)
      have hg : si.glue (Φ (LState.restrict Λ (T ω))).2 (Z (Φ (LState.restrict Λ (T ω))).1)
          = (pre ω).pick l.args hx := by
        rw [hΦT, hZT, si.glue_pick (T ω) (lift hx), node_pick ω hx]
      have hs : step (LState.restrict Λ (T ω)) = LState.restrict (B ∪ C) (T ω) := by
        dsimp only [step]
        rw [hg, ← node_res ω, hΦT, ← Ψ.symm_apply_apply (LState.restrict (B ∪ C) (T ω))]
        exact congrArg Ψ.symm (LState.extract_restrict ..).symm
      dsimp only
      rw [hs, hg, (htr ω).1, (htr ω).2]
      simp [logwLines]

end Density

/-! ## The density theorem -/

namespace Prog
variable {Sg : Sig S} (M : Model Sg)

/-- **Density** (`thm:density`) at the root of the derivation, together with surjectivity under
the right law. -/
theorem density_surj (P : Prog Sg) {ρ : List Role} (hT : P.Typed ρ) (θ : Val P.wts) :
    Cov (vol P.ins) (vol P.outs) (P.sem M hT.wf θ) (fun x => expw (P.logw M hT.wf ρ θ x)) ∧
      (M.RightLaw → Function.Surjective (P.sem M hT.wf θ)) := by
  have h := hT.wf
  set T : Val P.ins → Val P.ctx := P.tr M hT.wf θ with hTdef
  -- the weights are constant on the traces
  have hw : ∀ w < P.nw, Det T ∅ w := fun w hw a ha => by
    have hw' : w < P.wts.length := hw
    have hlt : w < P.ctx₀.length := by simp; omega
    have ha' : P.wts[w]? = some a := by
      rwa [getElem?_ctxAfter _ hlt, List.getElem?_append_left hw'] at ha
    refine ⟨fun _ => θ.get w ha', measurable_const, fun x => ?_⟩
    show (trace M P.lines hT.wf.1 (θ.append x)).get w ha = _
    rw [get_trace M _ _ _ ha hlt, Val.get_append_left _ _ hw']
  -- the recipes hold on the traces
  have hR : ∀ r ∈ P.xrecs ρ, Holds T r := fun r hr =>
    det_of_pick (P.wf_xrecs h r hr).1 (P.wf_xrecs h r hr).2 (M.sem_meas r.line.op) fun x =>
      (xrecs_holds M h.1 ρ rfl (θ.append x) r hr _ _).symm
  obtain ⟨F, L, hF, hL, hcov, hsurj, htr⟩ := density_aux M T P.nw hw hR hT (P.recov_init _)
    (fun _ hx => hx) h.1 rfl h.2 (fun x => θ.append x) (fun _ => rfl)
    (fun v hv => by simp [Prog.inputs] at hv; simp; omega)
  -- the inputs as a live state
  have hin₀ := Args.range' P.wts P.ins
  have hin : Args P.ctx _ P.ins := hin₀.ctxAfter P.lines
  have hA : ∀ x ∈ List.range' P.wts.length P.ins.length, x ∈ P.inputs := fun x hx => by
    simp only [List.mem_range'_1] at hx; simp [Prog.inputs]; omega
  have hcov' : ∀ x ∈ P.inputs, x ∈ List.range' P.wts.length P.ins.length := fun x hx => by
    simp only [Prog.inputs, Finset.mem_Ico] at hx; simp only [List.mem_range'_1]; omega
  set E := LState.pickEquiv hA hin List.nodup_range' hcov'
  have hU : ∀ x, E.symm x = LState.restrict P.inputs (T x) := fun x => by
    refine (congrArg E.symm ?_).trans (E.symm_apply_apply _)
    rw [LState.pickEquiv_apply, LState.pick_restrict]
    exact ((pick_trace M P.lines h.1 (θ.append x) hin₀ hin).trans (Val.pick_range'_append θ x)).symm
  have hsem : ∀ x, F (E.symm x) = P.sem M hT.wf θ x := fun x => by rw [hU, (htr x).1]; rfl
  refine ⟨(Cov.pres_comp ((LState.measurePreserving_pick ..).symm E) hF
    (measurable_expw.comp hL) hcov).congr hsem
    (fun x => by rw [Function.comp_apply, hU, (htr x).2]; rfl), fun hright y => ?_⟩
  obtain ⟨s, rfl⟩ := hsurj hright y
  exact ⟨E s, by rw [← hsem, E.symm_apply_apply]⟩

/-- **Density** (`thm:density`): a typed program pushes `e^{L_θ} · μ_I` forward to `μ_O`, for
every value of the weights. -/
theorem density (P : Prog Sg) {ρ : List Role} (hT : P.Typed ρ) (θ : Val P.wts) :
    Cov (vol P.ins) (vol P.outs) (P.sem M hT.wf θ) fun x => expw (P.logw M hT.wf ρ θ x) :=
  (P.density_surj M hT θ).1

/-- Under the right law, the semantics of a typed program is surjective
(`prop:surj`): every live-state step is solved by the semi-inverse. -/
theorem surjective (hM : M.RightLaw) (P : Prog Sg) {ρ : List Role} (hT : P.Typed ρ)
    (θ : Val P.wts) : Function.Surjective (P.sem M hT.wf θ) :=
  (P.density_surj M hT θ).2 hM

/-- **Joint density** (`cor:joint`): over the weights and the inputs together. Conditioning
values of conditional flows are weights, so this is the density of the joint distribution with
them. -/
theorem density_joint (P : Prog Sg) {ρ : List Role} (hT : P.Typed ρ) :
    Cov ((vol P.wts).prod (vol P.ins)) ((vol P.wts).prod (vol P.outs))
      (fun q => (q.1, P.sem M hT.wf q.1 q.2)) (fun q => expw (P.logw M hT.wf ρ q.1 q.2)) :=
  Cov.tri (P.measurable_sem M hT.wf) (measurable_expw.comp (P.measurable_logw M hT.wf ρ))
    (fun θ => P.density M hT θ)

end Prog

end RoleFlow