import RoleFlow.Flow.Ops
import RoleFlow.Flow.PermS

/-! # The flow instance (`docs/roleflow.tex`, `sec:flow`; module `Flow`)

A sort is a dimension `n` with a reference measure on `Vec n`: Lebesgue (`vec n`, data), the
standard normal (`gauss n`, randomness) or the standard logistic (`unit n`, a point of
`(0, 1)ⁿ` in its logit chart). The operations are those of
SemiInv2's flow signature, flat: arguments and results are separate variables.
Each operation has at most one semi-inverse, which solves its first argument (or, for
`split`/`concat`, all of them). An inverse pair is one constructor with a flag `inv`, and the
semi-inverse flips it:

| operation | `inv = false` | `inv = true` | log-weight |
|---|---|---|---|
| `add` | `x + q` | `x - q` (`sub`) | `0` |
| `scaleexp` | `x ⊙ e^s` | `x ⊙ e^{-s}` (`scalenexp`) | `±Σ s` |
| `triu` (Glow's LU factors) | `M x` | `M⁻¹ x` | `±triuLogdet U` |
| `oddpow k` | `x ^ (2k+1)` | `y ^ (1/(2k+1))` | `opowLogdet k x`, `-opowLogdet k (y ^ (1/(2k+1)))` |

The other pairs are `perm ps`/`perm ps⁻¹` (a permutation given by its steps, `Flow.PermS`; its
index map is read as an `Equiv.Perm` by `permOf`), `split`/`concat` (log-weight `0`), and
`rebase r r'`/`rebase r' r`, the identity from the sort `(r, n)` to `(r', n)` with log-weight
`log ρ_{r'} - log ρ_r` (`score`, `sample`, `sigmoid` and `logit` against Lebesgue measure).
The conditioners (`mvmul`, `mul`, `relu`, `conv3x3`, `attend`, …) have none.

The model `flowModel` is a `ModelInv`. Leaves are proved at one flag (by measure preservation,
the Jacobian formula or a density) and derived at the other by `leaf_dag`, whose derived
log-weight agrees with the one given here; `rebase` has one leaf for all pairs of measures.
Every declared pair is mutual, so the right law of an operation is the left law of its
semi-inverse. -/

open MeasureTheory

namespace RoleFlow.Flow

/-! ## Sorts -/

/-- The reference measures: Lebesgue, the standard normal, and the standard logistic. -/
inductive Ref
  | leb
  | normal
  | logistic
  deriving DecidableEq, Repr

/-- The reference measure on `Vec n`. -/
noncomputable def Ref.μ : Ref → (n : ℕ) → Measure (Vec n)
  | .leb, _ => volume
  | .normal, n => stdNormal n
  | .logistic, n => stdLogistic n

instance (r : Ref) (n : ℕ) : SigmaFinite (r.μ n) := by
  cases r
  exacts [inferInstanceAs (SigmaFinite (volume : Measure (Vec n))),
    inferInstanceAs (SigmaFinite (stdNormal n)), inferInstanceAs (SigmaFinite (stdLogistic n))]

/-- The log-density of the reference measure against Lebesgue measure. -/
noncomputable def Ref.logpdf : Ref → (n : ℕ) → Vec n → ℝ
  | .leb, _ => fun _ => 0
  | .normal, n => fun x => Real.log (stdNormalPDF n x)
  | .logistic, n => fun x => Real.log (logisticPDF n x)

/-- The sorts: a reference measure and a dimension. -/
structure FS where
  ref : Ref
  dim : ℕ
  deriving DecidableEq, Repr

/-- Data. -/
abbrev FS.vec (n : ℕ) : FS := ⟨.leb, n⟩
/-- Standard-normal randomness. -/
abbrev FS.gauss (n : ℕ) : FS := ⟨.normal, n⟩
/-- Points of the unit cube in their logit chart. -/
abbrev FS.unit (n : ℕ) : FS := ⟨.logistic, n⟩

noncomputable instance : Sorts FS where
  X s := Vec s.dim
  meas _ := inferInstance
  μ s := s.ref.μ s.dim
  sigmaFinite _ := inferInstance

open FS

/-! ## Operations -/

/-- The flow operations, flat. The flag `inv` selects the second operation of an inverse
pair. -/
inductive FOp : List FS → List FS → Type
  /-- The empty vector (the initial key-value cache). -/
  | nil : FOp [] [vec 0]
  | split (m n : ℕ) : FOp [vec (m + n)] [vec m, vec n]
  | concat (m n : ℕ) : FOp [vec m, vec n] [vec (m + n)]
  /-- The permutation `ps` (its inverse if `inv`): output entry `i` is input entry `ps.idx i`. -/
  | perm (n : ℕ) (ps : PermS) (inv := false) : FOp [vec n] [vec n]
  /-- The identity, from the reference measure `r` to `r'`. -/
  | rebase (r r' : Ref) (n : ℕ) : FOp [⟨r, n⟩] [⟨r', n⟩]
  | relu (n : ℕ) : FOp [vec n] [vec n]
  | tanh (n : ℕ) : FOp [vec n] [vec n]
  | layernorm (n : ℕ) : FOp [vec n] [vec n]
  | gelu (n : ℕ) : FOp [vec n] [vec n]
  /-- `x ↦ x ^ (2k + 1)`, elementwise; `y ↦ y ^ (1/(2k + 1))` (the odd real root) if `inv`. -/
  | oddpow (n k : ℕ) (inv := false) : FOp [vec n] [vec n]
  | rows2 (m n : ℕ) : FOp [vec ((m + m) * n)] [vec (m * n), vec (m * n)]
  /-- `x + q`; `x - q` if `inv`. -/
  | add (n : ℕ) (inv := false) : FOp [vec n, vec n] [vec n]
  /-- `x ⊙ exp s`; `x ⊙ exp (-s)` if `inv`. -/
  | scaleexp (n : ℕ) (inv := false) : FOp [vec n, vec n] [vec n]
  /-- `M x` (`triuMat`), `M⁻¹ x` if `inv`: with channel permutations, Glow's LU 1 × 1
  convolution. -/
  | triu (c n : ℕ) (inv := false) : FOp [vec (c * n), vec (c * c)] [vec (c * n)]
  /-- `W · x` for a row-major `m × n` matrix `W`. -/
  | mvmul (m n : ℕ) : FOp [vec (m * n), vec n] [vec m]
  /-- `x ⊙ y`, elementwise. -/
  | mul (n : ℕ) : FOp [vec n, vec n] [vec n]
  /-- A per-channel bias. -/
  | addc (c h w : ℕ) : FOp [vec (c * (h * w)), vec c] [vec (c * (h * w))]
  /-- Kernel first, then the image. -/
  | conv1x1 (cin cout h w : ℕ) : FOp [vec (cout * cin), vec (cin * (h * w))] [vec (cout * (h * w))]
  | conv3x3 (cin cout h w : ℕ) :
      FOp [vec (cout * (cin * 9)), vec (cin * (h * w))] [vec (cout * (h * w))]
  /-- A query against a key-value cache. -/
  | attend (m h : ℕ) : FOp [vec h, vec (h * m), vec (h * m)] [vec h]

/-! Names for the flipped flags and the changes of sort that the models use. -/
abbrev FOp.sub (n : ℕ) : FOp [vec n, vec n] [vec n] := .add n true
abbrev FOp.scalenexp (n : ℕ) : FOp [vec n, vec n] [vec n] := .scaleexp n true
abbrev FOp.score (n : ℕ) : FOp [vec n] [gauss n] := .rebase .leb .normal n
abbrev FOp.sample (n : ℕ) : FOp [gauss n] [vec n] := .rebase .normal .leb n
abbrev FOp.logit (n : ℕ) : FOp [unit n] [vec n] := .rebase .logistic .leb n
abbrev FOp.cube (n : ℕ) : FOp [vec n] [vec n] := .oddpow n 1
abbrev FOp.cbrt (n : ℕ) : FOp [vec n] [vec n] := .oddpow n 1 true

/-- A semi-inverse solving the first argument from the result and the second argument. -/
abbrev si1 {a b c : FS} (o : FOp [c, b] [a]) : SemiInv FOp [a, b] [c] :=
  ⟨[a], [b], rfl, by simp, o⟩

/-- A semi-inverse of a one-argument operation. -/
abbrev si0 {a b : FS} (o : FOp [b] [a]) : SemiInv FOp [a] [b] := ⟨[a], [], rfl, by simp, o⟩

/-- The declared semi-inverses. -/
def FOp.inv : {as bs : List FS} → FOp as bs → Option (SemiInv FOp as bs)
  | _, _, .split m n => some ⟨[vec (m + n)], [], rfl, by simp, .concat m n⟩
  | _, _, .concat m n => some ⟨[vec m, vec n], [], rfl, by simp, .split m n⟩
  | _, _, .perm n ps inv => some (si0 (.perm n ps !inv))
  | _, _, .rebase r r' n => some (si0 (.rebase r' r n))
  | _, _, .oddpow n k inv => some (si0 (.oddpow n k !inv))
  | _, _, .add n inv => some (si1 (.add n !inv))
  | _, _, .scaleexp n inv => some (si1 (.scaleexp n !inv))
  | _, _, .triu c n inv => some (si1 (.triu c n !inv))
  | _, _, _ => none

/-- The flow signature. -/
def flowSig : Sig FS where
  Op := FOp
  bs_ne := fun o => by cases o <;> simp_all
  inv := FOp.inv

/-! ## Tuples -/

noncomputable section

/-- A one-element tuple. -/
abbrev one {a : FS} (x : X a) : Val [a] := Val.cons x Val.nil

/-- A two-element tuple. -/
abbrev two {a b : FS} (x : X a) (y : X b) : Val [a, b] := Val.cons x (one y)

/-- The arguments of a tuple. -/
abbrev a₀ {a : FS} {Γ : List FS} (v : Val (a :: Γ)) : X a := v.head
abbrev a₁ {a b : FS} {Γ : List FS} (v : Val (a :: b :: Γ)) : X b := v.tail.head
abbrev a₂ {a b c : FS} {Γ : List FS} (v : Val (a :: b :: c :: Γ)) : X c := v.tail.tail.head

/-! ## The model -/

/-- `t`, or `-t` if `inv`. -/
def negIf : Bool → ℝ → ℝ
  | false, t => t
  | true, t => -t

theorem negIf_not (inv : Bool) (t : ℝ) : negIf (!inv) t = -negIf inv t := by
  cases inv <;> simp [negIf]

/-- The permutation of `Fin n` whose index map is `ps.idx` (taken mod `n`), if that map is
injective; otherwise the identity. A `PermS.valid` permutation is one, but nothing here needs it:
every operation is interpreted, and `Print.certify` refuses invalid ones. -/
def permS (n : ℕ) (ps : PermS) : Equiv.Perm (Fin n) := by
  classical
  let f : Fin n → Fin n := fun i => ⟨ps.idx i % n, Nat.mod_lt _ i.pos⟩
  exact if h : Function.Injective f then Equiv.ofBijective f (Finite.injective_iff_bijective.mp h)
    else Equiv.refl _

/-- `perm n ps inv` as a permutation of the entries. -/
def permOf (n : ℕ) (ps : PermS) (inv : Bool) : Equiv.Perm (Fin n) :=
  if inv then (permS n ps).symm else permS n ps

theorem permOf_not (n : ℕ) (ps : PermS) (inv : Bool) :
    permOf n ps (!inv) = (permOf n ps inv).symm := by
  cases inv <;> simp [permOf]

/-- The interpretation of the operations. -/
def semOp : {as bs : List FS} → FOp as bs → Val as → Val bs
  | _, _, .nil => fun _ => one (a := vec 0) Fin.elim0
  | _, _, .split m n => fun v => two (a := vec m) (b := vec n) (vsplit m n (v.head)).1 (vsplit m n (v.head)).2
  | _, _, .concat m n => fun v => one (a := vec (m + n)) (vjoin (a := (v.head : Vec m)) (v.tail.head : Vec n))
  | _, _, .perm n ps inv => fun v => one (a := vec n) fun i => (v.head : Vec n) (permOf n ps inv i)
  | _, _, .rebase _ r' n => fun v => one (a := ⟨r', n⟩) (v.head)
  | _, _, .relu n => fun v => one (a := vec n) (fvrelu (v.head : Vec n))
  | _, _, .tanh n => fun v => one (a := vec n) (fvtanh (v.head : Vec n))
  | _, _, .layernorm n => fun v => one (a := vec n) (flayernorm (v.head : Vec n))
  | _, _, .gelu n => fun v => one (a := vec n) (fgelu (v.head : Vec n))
  | _, _, .oddpow n k inv => fun v =>
      one (a := vec n) (bif inv then foroot k (v.head : Vec n) else fopow k (v.head : Vec n))
  | _, _, .rows2 m n => fun v => two (a := vec (m * n)) (b := vec (m * n))
      (frows2 m n (v.head)).1 (frows2 m n (v.head)).2
  | _, _, .add n inv => fun v =>
      one (a := vec n) fun i => (v.head : Vec n) i + negIf inv ((v.tail.head : Vec n) i)
  | _, _, .scaleexp n inv => fun v =>
      one (a := vec n) fun i => (v.head : Vec n) i * Real.exp (negIf inv ((v.tail.head : Vec n) i))
  | _, _, .triu c n inv => fun v => one (a := vec (c * n)) (bif inv
      then ftriuInv c n (v.head : Vec (c * n)) (v.tail.head) else ftriu c n (v.head) (v.tail.head))
  | _, _, .mvmul m n => fun v => one (a := vec m) (fmvmul m n (v.head) (v.tail.head))
  | _, _, .mul n => fun v => one (a := vec n) fun i => (v.head : Vec n) i * (v.tail.head : Vec n) i
  | _, _, .addc c h w => fun v => one (a := vec (c * (h * w))) (faddc c h w (v.head) (v.tail.head))
  | _, _, .conv1x1 cin cout h w => fun v =>
      one (a := vec (cout * (h * w))) (fconv1x1 cin cout h w (v.head) (v.tail.head))
  | _, _, .conv3x3 cin cout h w => fun v =>
      one (a := vec (cout * (h * w))) (fconv3x3 cin cout h w (v.head) (v.tail.head))
  | _, _, .attend m h => fun v => one (a := vec h) (fattend m h (v.head) (v.tail.head) (v.tail.tail.head))

/-- The log-weights, on the whole argument tuple; `0` for the structural bijections and for
the conditioners. -/
def ℓOp : {as bs : List FS} → FOp as bs → Val as → ℝ
  | _, _, .rebase r r' n => fun v => r'.logpdf n (v.head) - r.logpdf n (v.head)
  | _, _, .oddpow n k inv => fun v =>
      bif inv then -opowLogdet k (foroot k (v.head : Vec n)) else opowLogdet k (v.head : Vec n)
  | _, _, .scaleexp n inv => fun v => negIf inv (∑ i, (v.tail.head : Vec n) i)
  | _, _, .triu c n inv => fun v => negIf inv (triuLogdet c n (v.tail.head))
  | _, _, _ => fun _ => 0

/-! ### Measurability -/

theorem measurable_a₀ {a : FS} {Γ : List FS} : Measurable fun v : Val (a :: Γ) => a₀ v :=
  Val.measurable_head

theorem measurable_a₁ {a b : FS} {Γ : List FS} : Measurable fun v : Val (a :: b :: Γ) => a₁ v :=
  Val.measurable_head.comp Val.measurable_tail

theorem measurable_a₂ {a b c : FS} {Γ : List FS} :
    Measurable fun v : Val (a :: b :: c :: Γ) => a₂ v :=
  Val.measurable_head.comp (Val.measurable_tail.comp Val.measurable_tail)

theorem measurable_one {a : FS} : Measurable fun x : X a => one x :=
  Val.measurable_cons.comp (measurable_id.prodMk measurable_const)

theorem measurable_one' {a b : FS} {Γ : List FS} {f : X a → X b} (hf : Measurable f) :
    Measurable fun v : Val (a :: Γ) => one (f (a₀ v)) :=
  measurable_one.comp (hf.comp measurable_a₀)

theorem measurable_two' {a b c : FS} {Γ : List FS} {f : X a → X b × X c} (hf : Measurable f) :
    Measurable fun v : Val (a :: Γ) => two (f (a₀ v)).1 (f (a₀ v)).2 :=
  Val.measurable_cons' (measurable_fst.comp (hf.comp measurable_a₀))
    (measurable_one.comp (measurable_snd.comp (hf.comp measurable_a₀)))

theorem measurable_one₂ {a b c : FS} {Γ : List FS} {f : X a → X b → X c}
    (hf : Measurable fun p : X a × X b => f p.1 p.2) :
    Measurable fun v : Val (a :: b :: Γ) => one (f (a₀ v) (a₁ v)) :=
  measurable_one.comp (hf.comp (measurable_a₀.prodMk measurable_a₁))

theorem measurable_one₃ {a b c d : FS} {Γ : List FS} {f : X a → X b → X c → X d}
    (hf : Measurable fun p : X a × X b × X c => f p.1 p.2.1 p.2.2) :
    Measurable fun v : Val (a :: b :: c :: Γ) => one (f (a₀ v) (a₁ v) (a₂ v)) :=
  measurable_one.comp (hf.comp (measurable_a₀.prodMk (measurable_a₁.prodMk measurable_a₂)))

theorem measurable_negIf (inv : Bool) : Measurable (negIf inv) := by
  cases inv <;> unfold negIf <;> fun_prop

theorem measurable_logpdf (r : Ref) (n : ℕ) : Measurable (r.logpdf n) := by
  cases r
  exacts [measurable_const, (measurable_stdNormalPDF n).log, (measurable_logisticPDF n).log]

theorem measurable_semOp {as bs : List FS} (o : FOp as bs) : Measurable (semOp o) := by
  cases o with
  | nil => exact measurable_const
  | split m n => exact measurable_two' (a := vec (m + n)) (b := vec m) (c := vec n) (measurePreserving_vsplit m n).measurable
  | concat m n => exact measurable_one₂ (a := vec m) (b := vec n) (c := vec (m + n)) measurable_vjoin
  | perm n ps inv => exact measurable_one' (a := vec n) (b := vec n) (f := fun x : Vec n => fun i => x (permOf n ps inv i)) (by fun_prop)
  | rebase r r' n => exact measurable_one' (a := ⟨r, n⟩) (b := ⟨r', n⟩) (f := fun x : Vec n => x) measurable_id
  | relu n => exact measurable_one' (a := vec n) (b := vec n) (f := fvrelu (n := n)) (by unfold fvrelu; fun_prop)
  | tanh n => exact measurable_one' (a := vec n) (b := vec n) (f := fvtanh (n := n)) (by unfold fvtanh; fun_prop)
  | layernorm n => exact measurable_one' (a := vec n) (b := vec n) (f := flayernorm (n := n)) (by unfold flayernorm; fun_prop)
  | gelu n => exact measurable_one' (a := vec n) (b := vec n) (f := fgelu (n := n)) (by unfold fgelu; fun_prop)
  | oddpow n k inv =>
    cases inv
    exacts [measurable_one' (a := vec n) (b := vec n) (measurable_fopow k n),
      measurable_one' (a := vec n) (b := vec n) (measurable_foroot k n)]
  | rows2 m n => exact measurable_two' (a := vec ((m + m) * n)) (b := vec (m * n)) (c := vec (m * n)) (f := frows2 m n) (by unfold frows2; fun_prop)
  | add n inv =>
    have := measurable_negIf inv
    exact measurable_one₂ (a := vec n) (b := vec n) (c := vec n) (f := fun x y : Vec n => fun i => x i + negIf inv (y i)) (by fun_prop)
  | scaleexp n inv =>
    have := measurable_negIf inv
    exact measurable_one₂ (a := vec n) (b := vec n) (c := vec n) (f := fun x y : Vec n => fun i => x i * Real.exp (negIf inv (y i))) (by fun_prop)
  | triu c n inv =>
    cases inv
    exacts [measurable_one₂ (a := vec (c * n)) (b := vec (c * c)) (c := vec (c * n)) (measurable_ftriu c n),
      measurable_one₂ (a := vec (c * n)) (b := vec (c * c)) (c := vec (c * n)) (measurable_ftriuInv c n)]
  | mvmul m n => exact measurable_one₂ (a := vec (m * n)) (b := vec n) (c := vec m) (f := fmvmul m n) (by unfold fmvmul; fun_prop)
  | mul n => exact measurable_one₂ (a := vec n) (b := vec n) (c := vec n) (f := fun x y : Vec n => fun i => x i * y i) (by fun_prop)
  | addc c h w => exact measurable_one₂ (a := vec (c * (h * w))) (b := vec c) (c := vec (c * (h * w))) (f := faddc c h w) (by unfold faddc; fun_prop)
  | conv1x1 cin cout h w =>
    exact measurable_one₂ (a := vec (cout * cin)) (b := vec (cin * (h * w))) (c := vec (cout * (h * w))) (f := fconv1x1 cin cout h w) (by unfold fconv1x1; fun_prop)
  | conv3x3 cin cout h w =>
    refine measurable_one₂ (a := vec (cout * (cin * 9))) (b := vec (cin * (h * w))) (c := vec (cout * (h * w))) (f := fconv3x3 cin cout h w) ?_
    unfold fconv3x3
    refine measurable_pi_lambda _ fun i => ?_
    refine Finset.measurable_sum _ fun _ _ => Finset.measurable_sum _ fun _ _ =>
      Finset.measurable_sum _ fun _ _ => ?_
    split_ifs <;> fun_prop
  | attend m h =>
    have : Measurable fun p : Vec h × Vec (h * m) × Vec (h * m) => fattend m h p.1 p.2.1 p.2.2 := by
      unfold fattend; fun_prop
    exact measurable_one₃ (a := vec h) (b := vec (h * m)) (c := vec (h * m)) (d := vec h) this

theorem measurable_ℓOp : ∀ {as bs : List FS} (o : FOp as bs), Measurable (ℓOp o)
  | _, _, .rebase r r' n =>
    ((measurable_logpdf r' n).comp measurable_a₀).sub ((measurable_logpdf r n).comp measurable_a₀)
  | _, _, .oddpow n k false => (measurable_opowLogdet k n).comp measurable_a₀
  | _, _, .oddpow n k true =>
    ((measurable_opowLogdet k n).comp ((measurable_foroot k n).comp measurable_a₀)).neg
  | _, _, .scaleexp n inv => (measurable_negIf inv).comp
      (Finset.measurable_sum _ fun i _ => (measurable_pi_apply i).comp measurable_a₁)
  | _, _, .triu c n inv => (measurable_negIf inv).comp ((measurable_triuLogdet c n).comp measurable_a₁)
  | _, _, .nil | _, _, .split .. | _, _, .concat ..
  | _, _, .perm .. | _, _, .relu _ | _, _, .tanh _ | _, _, .layernorm _ | _, _, .gelu _
  | _, _, .rows2 .. | _, _, .add .. | _, _, .mvmul .. | _, _, .mul _ | _, _, .addc ..
  | _, _, .conv1x1 .. | _, _, .conv3x3 ..
  | _, _, .attend .. => measurable_const

/-! ### Laws -/

theorem eq_one {a : FS} (v : Val [a]) : one v.head = v := by
  conv_rhs => rw [← Val.cons_head_tail v, Val.eq_nil v.tail]

/-- The left law of `o` with the semi-inverse `si`. -/
abbrev LeftLaw {as bs : List FS} (o : FOp as bs) (si : SemiInv FOp as bs) : Prop :=
  ∀ c p, semOp si.op ((semOp o (si.glue c p)).append p) = c

/-- The left law of a one-argument operation. -/
theorem left0 {a b : FS} {o : FOp [a] [b]} {o' : FOp [b] [a]} (F : X a → X b) (G : X b → X a)
    (hF : ∀ v, semOp o v = one (F v.head)) (hG : ∀ v, semOp o' v = one (G v.head))
    (h : ∀ x, G (F x) = x) : LeftLaw o (si0 o') := fun (c : Val [a]) p => by
  show semOp o' ((semOp o (c.append p)).append p) = c
  rw [hF, hG]; show one (G (F c.head)) = c; rw [h]; exact eq_one c

/-- The left law of an operation that consumes its first argument. -/
theorem left1 {a b c : FS} {o : FOp [a, b] [c]} {o' : FOp [c, b] [a]} (F : X a → X b → X c)
    (G : X c → X b → X a) (hF : ∀ v, semOp o v = one (F v.head v.tail.head))
    (hG : ∀ v, semOp o' v = one (G v.head v.tail.head)) (h : ∀ x q, G (F x q) q = x) :
    LeftLaw o (si1 o') := fun (x : Val [a]) (p : Val [b]) => by
  show semOp o' ((semOp o (x.append p)).append p) = x
  rw [hF, hG]; show one (G (F x.head p.head) p.head) = x; rw [h]; exact eq_one x

theorem left_split (m n : ℕ) (c : Val [vec (m + n)]) (p : Val ([] : List FS)) :
    semOp (.concat m n) ((semOp (.split m n) ((c.append p).cast rfl)).append p) = c := by
  conv_rhs => rw [← eq_one c]
  simp only [semOp, Val.append, Val.head_cons, Val.tail_cons, Val.cast_rfl]
  exact congrArg one (vjoin_vsplit m n _)

theorem left_concat (m n : ℕ) (c : Val [vec m, vec n]) (p : Val ([] : List FS)) :
    semOp (.split m n) ((semOp (.concat m n) ((c.append p).cast rfl)).append p) = c := by
  conv_rhs => rw [← Val.cons_head_tail c, ← eq_one c.tail]
  simp [semOp, Val.append]

theorem left_perm (n : ℕ) (ps : PermS) (inv : Bool) :
    LeftLaw (.perm n ps inv) (si0 (.perm n ps !inv)) :=
  left0 (fun x : Vec n => fun i => x (permOf n ps inv i)) (fun x => fun i => x (permOf n ps (!inv) i))
    (fun _ => rfl) (fun _ => rfl) fun x => funext fun i => by
      rw [permOf_not]; exact congrArg x ((permOf n ps inv).apply_symm_apply i)

theorem left_rebase (r r' : Ref) (n : ℕ) : LeftLaw (.rebase r r' n) (si0 (.rebase r' r n)) :=
  left0 (fun x : Vec n => x) (fun x => x) (fun _ => rfl) (fun _ => rfl) fun _ => rfl

theorem left_oddpow (n k : ℕ) (inv : Bool) :
    LeftLaw (.oddpow n k inv) (si0 (.oddpow n k !inv)) := by
  cases inv
  exacts [left0 (fopow k) (foroot k) (fun _ => rfl) (fun _ => rfl) (foroot_fopow k),
    left0 (foroot k) (fopow k) (fun _ => rfl) (fun _ => rfl) (fopow_foroot k)]

theorem left_add (n : ℕ) (inv : Bool) : LeftLaw (.add n inv) (si1 (.add n !inv)) :=
  left1 (fun x q : Vec n => fun i => x i + negIf inv (q i))
    (fun x q => fun i => x i + negIf (!inv) (q i)) (fun _ => rfl) (fun _ => rfl)
    fun _ _ => funext fun _ => by simp [negIf_not]

theorem left_scaleexp (n : ℕ) (inv : Bool) : LeftLaw (.scaleexp n inv) (si1 (.scaleexp n !inv)) :=
  left1 (fun x q : Vec n => fun i => x i * Real.exp (negIf inv (q i)))
    (fun x q => fun i => x i * Real.exp (negIf (!inv) (q i))) (fun _ => rfl) (fun _ => rfl)
    fun _ _ => funext fun _ => by simp [negIf_not, mul_assoc, ← Real.exp_add]

theorem left_triu (c n : ℕ) (inv : Bool) : LeftLaw (.triu c n inv) (si1 (.triu c n !inv)) := by
  cases inv
  exacts [left1 (fun (x : Vec (c * n)) U => ftriu c n x U) (fun x U => ftriuInv c n x U)
      (fun _ => rfl) (fun _ => rfl) ftriuInv_ftriu,
    left1 (fun (x : Vec (c * n)) U => ftriuInv c n x U) (fun x U => ftriu c n x U)
      (fun _ => rfl) (fun _ => rfl) ftriu_ftriuInv]

/-! ### Leaves -/

/-- Change of variables on one value, carried to one-element contexts. -/
theorem cov_one {a b : FS} {f : X a → X b} {J : X a → ENNReal} (hf : Measurable f)
    (hJ : Measurable J) (h : Cov (μ a) (μ b) f J) :
    Cov (vol [a]) (vol [b]) (fun x => one (f x.head)) (fun x => J x.head) := fun φ hφ => by
  have h₁ : Measurable fun y => φ (one (f y)) * J y := ((hφ.comp measurable_one).comp hf).mul hJ
  show ∫⁻ x, (fun y => φ (one (f y)) * J y) x.head ∂vol [a] = _
  rw [(Val.measurePreserving_head (a := a)).lintegral_comp h₁,
    h (fun y => φ (one y)) (hφ.comp measurable_one)]
  exact Val.measurePreserving_single.lintegral_comp hφ

/-- From Lebesgue measure to a reference measure by the identity. -/
theorem cov_toRef (r : Ref) (n : ℕ) :
    Cov volume (r.μ n) (fun x => x) fun x => expw (r.logpdf n x) := by
  cases r
  · exact (Cov.pres (MeasurePreserving.id _)).congr (fun _ => rfl) fun _ => by simp [Ref.logpdf]
  · exact cov_toDensity _ (measurable_stdNormalPDF n) (stdNormalPDF_pos n)
  · exact cov_toDensity _ (measurable_logisticPDF n) (logisticPDF_pos n)

/-- `rebase`: from `r` to Lebesgue measure by `Cov.inv`, then to `r'`. -/
theorem leaf_rebase (r r' : Ref) (n : ℕ) (p : Val ([] : List FS)) :
    Cov (vol [FS.mk r n]) (vol [FS.mk r' n]) (fun c => semOp (.rebase r r' n) ((c.append p).cast rfl))
      (fun c => expw (ℓOp (.rebase r r' n) ((c.append p).cast rfl))) :=
  (cov_one (a := ⟨r, n⟩) (b := ⟨r', n⟩) (f := fun x : Vec n => x) measurable_id
    (measurable_expw.comp ((measurable_logpdf r' n).sub (measurable_logpdf r n)))
    ((((cov_toRef r n).inv measurable_id (measurable_logpdf r n) fun _ => rfl).comp measurable_id
      (measurable_expw.comp (measurable_logpdf r' n)) (cov_toRef r' n)).congr (fun _ => rfl)
      fun _ => by simp only [Function.comp_apply, id, sub_eq_neg_add, expw_add])).congr
    (fun _ => rfl) (fun _ => rfl)

theorem leaf_add (n : ℕ) (p : Val [vec n]) :
    Cov (vol [vec n]) (vol [vec n]) (fun c => semOp (.add n) ((c.append p).cast rfl))
      (fun c => expw (ℓOp (.add n) ((c.append p).cast rfl))) :=
  (cov_one (a := vec n) (b := vec n) (by fun_prop) measurable_const (cov_add (p.head : Vec n))).congr
    (fun _ => rfl) (fun _ => by simp [ℓOp])

theorem leaf_scaleexp (n : ℕ) (p : Val [vec n]) :
    Cov (vol [vec n]) (vol [vec n]) (fun c => semOp (.scaleexp n) ((c.append p).cast rfl))
      (fun c => expw (ℓOp (.scaleexp n) ((c.append p).cast rfl))) :=
  (cov_one (a := vec n) (b := vec n) (by fun_prop) measurable_const
    (cov_scale (p.head : Vec n))).congr (fun _ => rfl) (fun _ => rfl)

theorem leaf_triu (c n : ℕ) (p : Val [vec (c * c)]) :
    Cov (vol [vec (c * n)]) (vol [vec (c * n)]) (fun x => semOp (.triu c n) ((x.append p).cast rfl))
      (fun x => expw (ℓOp (.triu c n) ((x.append p).cast rfl))) :=
  (cov_one (a := vec (c * n)) (b := vec (c * n))
    ((measurable_ftriu c n).comp (measurable_id.prodMk measurable_const)) measurable_const
    (cov_triu (n := n) (p.head : Vec (c * c)))).congr (fun _ => rfl) (fun _ => rfl)

theorem leaf_perm (n : ℕ) (ps : PermS) (inv : Bool) (p : Val ([] : List FS)) :
    Cov (vol [vec n]) (vol [vec n]) (fun c => semOp (.perm n ps inv) ((c.append p).cast rfl))
      (fun c => expw (ℓOp (.perm n ps inv) ((c.append p).cast rfl))) :=
  (cov_one (a := vec n) (b := vec n) (f := fun x : Vec n => fun i => x (permOf n ps inv i))
    (by fun_prop) measurable_const (cov_perm (permOf n ps inv))).congr (fun _ => rfl)
    (fun _ => by simp [ℓOp])

theorem leaf_oddpow (n k : ℕ) (p : Val ([] : List FS)) :
    Cov (vol [vec n]) (vol [vec n]) (fun c => semOp (.oddpow n k) ((c.append p).cast rfl))
      (fun c => expw (ℓOp (.oddpow n k) ((c.append p).cast rfl))) :=
  (cov_one (a := vec n) (b := vec n) (measurable_fopow k n)
    (measurable_expw.comp (measurable_opowLogdet k n)) (cov_fopow k n)).congr
    (fun _ => rfl) (fun _ => rfl)

theorem leaf_split (m n : ℕ) (p : Val ([] : List FS)) :
    Cov (vol [vec (m + n)]) (vol [vec m, vec n]) (fun c => semOp (.split m n) ((c.append p).cast rfl))
      (fun c => expw (ℓOp (.split m n) ((c.append p).cast rfl))) := by
  have h : MeasurePreserving
      (fun c : Val [vec (m + n)] => two (a := vec m) (b := vec n)
        (vsplit m n (c.head : Vec (m + n))).1 (vsplit m n (c.head : Vec (m + n))).2)
      (vol [vec (m + n)]) (vol [vec m, vec n]) :=
    (Val.measurePreserving_cons (a := vec m) (Γ := [vec n])).comp
      (((MeasurePreserving.id (μ (vec m))).prod (Val.measurePreserving_single (a := vec n))).comp
        ((measurePreserving_vsplit m n).comp (Val.measurePreserving_head (a := vec (m + n)))))
  exact (Cov.pres h).congr (fun _ => rfl) (fun _ => by simp [ℓOp])

/-! ### Derived leaves -/

/-- The leaf of a semi-inverse, derived by `leaf_dag` from the leaf and the left law of the
operation it inverts; the derived log-weight is checked against `ℓOp`. -/
theorem leaf_of_dag {as bs : List FS} {o : FOp as bs} {si : SemiInv FOp as bs}
    (left : LeftLaw o si)
    (leaf : ∀ p, Cov (vol si.cs) (vol bs) (fun c => semOp o ((c.append p).cast si.split))
      (fun c => expw (ℓOp o ((c.append p).cast si.split))))
    {L : Val (bs ++ si.ps) → ℝ} (hL : ∀ w, dagWeight (semOp si.op) (ℓOp o) si.split w = L w)
    (p : Val si.ps) :
    Cov (vol bs) (vol si.cs) (fun b => semOp si.op (b.append p)) (fun b => expw (L (b.append p))) :=
  (leaf_dag (measurable_semOp _) (measurable_ℓOp _) left leaf p).congr (fun _ => rfl)
    (fun b => by rw [hL])

/-! ### Right laws -/

/-- The right law: every declared pair is mutual, so the right law of an operation is the left
law of its semi-inverse, at the flipped flag. -/
theorem flow_right {as bs : List FS} (o : FOp as bs) (si : SemiInv FOp as bs) (h : FOp.inv o = some si)
    (b : Val bs) (p : Val si.ps) : semOp o (si.glue (semOp si.op (b.append p)) p) = b := by
  cases o <;> cases h
  case split m n => exact left_concat m n b p
  case concat m n => exact left_split m n b p
  case perm n ps inv => cases inv; exacts [left_perm n ps true b p, left_perm n ps false b p]
  case rebase r r' n => exact left_rebase r' r n b p
  case oddpow n k inv => cases inv; exacts [left_oddpow n k true b p, left_oddpow n k false b p]
  case add n inv => cases inv; exacts [left_add n true b p, left_add n false b p]
  case scaleexp n inv => cases inv; exacts [left_scaleexp n true b p, left_scaleexp n false b p]
  case triu c n inv => cases inv; exacts [left_triu c n true b p, left_triu c n false b p]

/-- The model of the flow signature. -/
def flowModel : ModelInv flowSig where
  sem := semOp
  ℓ := ℓOp
  sem_meas := measurable_semOp
  ℓ_meas := measurable_ℓOp
  left := by
    intro as bs o si h
    cases o <;> cases h
    case split m n => exact left_split m n
    case concat m n => exact left_concat m n
    case perm n ps inv => exact left_perm n ps inv
    case rebase r r' n => exact left_rebase r r' n
    case oddpow n k inv => exact left_oddpow n k inv
    case add n inv => exact left_add n inv
    case scaleexp n inv => exact left_scaleexp n inv
    case triu c n inv => exact left_triu c n inv
  leaf := by
    intro as bs o si h p
    cases o <;> cases h
    case split m n => exact leaf_split m n p
    case concat m n =>
      exact leaf_of_dag (o := .split m n) (si := ⟨[vec (m + n)], [], rfl, by simp, .concat m n⟩)
        (left_split m n) (leaf_split m n) (L := fun _ => 0) (fun _ => by simp [dagWeight, ℓOp]) p
    case perm n ps inv => exact leaf_perm n ps inv p
    case rebase r r' n => exact leaf_rebase r r' n p
    case oddpow n k inv =>
      cases inv
      exacts [leaf_oddpow n k p,
        leaf_of_dag (left_oddpow n k false) (leaf_oddpow n k) (fun _ => rfl) p]
    case add n inv =>
      cases inv
      exacts [leaf_add n p,
        leaf_of_dag (left_add n false) (leaf_add n) (fun _ => by simp [dagWeight, ℓOp]) p]
    case scaleexp n inv =>
      cases inv
      exacts [leaf_scaleexp n p, leaf_of_dag (left_scaleexp n false) (leaf_scaleexp n) (fun _ => rfl) p]
    case triu c n inv =>
      cases inv
      exacts [leaf_triu c n p, leaf_of_dag (left_triu c n false) (leaf_triu c n) (fun _ => rfl) p]
  right := flow_right

end

end RoleFlow.Flow