import RoleFlow.Sig
import Mathlib.Analysis.SpecialFunctions.Pow.Continuity
import Mathlib.Analysis.SpecialFunctions.Trigonometric.Deriv
import Mathlib.LinearAlgebra.Matrix.Block
import Mathlib.LinearAlgebra.Matrix.NonsingularInverse
import Mathlib.MeasureTheory.Function.Jacobian
import Mathlib.MeasureTheory.Measure.Haar.OfBasis
import Mathlib.MeasureTheory.Group.Measure
import Mathlib.Probability.Distributions.Gaussian.Real

/-! # The flows' operations over `ℝ`

The real functions behind the flow signature (`RoleFlow.Flow.Basic`), on flat vectors `Vec n`, and
the facts its leaves need: the reference measures of the `gauss` and `unit` sorts (Lebesgue
measure with the standard-normal and the logistic density), change of variables for invertible
matrices (through the Jacobian formula), for the elementwise odd powers (off the null set
of the axes) and for the changes of measure, and measurability of every conditioner. Ported from
SemiInv2; only the weights are now `expw` of a log-weight. -/

open MeasureTheory ENNReal

namespace RoleFlow.Flow

/-- A flat vector. -/
abbrev Vec (n : ℕ) := Fin n → ℝ

/-! ## Splitting and joining -/

/-- `Vec (m + n)` cut into its first `m` and last `n` entries. -/
def vsplit (m n : ℕ) (v : Vec (m + n)) : Vec m × Vec n :=
  (fun i => v (Fin.castAdd n i), fun j => v (Fin.natAdd m j))

/-- Concatenation. -/
def vjoin {m n : ℕ} (a : Vec m) (b : Vec n) : Vec (m + n) := Fin.append a b

theorem vjoin_vsplit (m n : ℕ) (v : Vec (m + n)) : vjoin (vsplit m n v).1 (vsplit m n v).2 = v :=
  Fin.append_castAdd_natAdd

@[simp] theorem vsplit_vjoin {m n : ℕ} (a : Vec m) (b : Vec n) : vsplit m n (vjoin a b) = (a, b) := by
  ext i <;> simp [vsplit, vjoin]

theorem measurePreserving_vsplit (m n : ℕ) :
    MeasurePreserving (vsplit m n) volume (volume.prod volume) := by
  have h₁ := (measurePreserving_piCongrLeft (fun _ : Fin (m + n) => (volume : Measure ℝ))
    finSumFinEquiv).symm
  have h₂ := measurePreserving_sumPiEquivProdPi (fun _ : Fin m ⊕ Fin n => (volume : Measure ℝ))
  exact (h₂.comp h₁ : _)

@[fun_prop] theorem measurable_vjoin {m n : ℕ} :
    Measurable fun p : Vec m × Vec n => vjoin p.1 p.2 := by
  refine measurable_pi_lambda _ fun i => ?_
  refine Fin.addCases (fun i => ?_) (fun j => ?_) i
  · simp only [vjoin, Fin.append_left]; fun_prop
  · simp only [vjoin, Fin.append_right]; fun_prop

theorem measurePreserving_vjoin (m n : ℕ) :
    MeasurePreserving (fun p : Vec m × Vec n => vjoin p.1 p.2) (volume.prod volume) volume :=
  (measurePreserving_vsplit m n).of_inverse measurable_vjoin fun v => vjoin_vsplit m n v

/-! ## Conditioners -/

/-- A flat read at a `ℕ` index; `0` out of range. -/
def natGet {k : ℕ} (v : Vec k) (i : ℕ) : ℝ := if h : i < k then v ⟨i, h⟩ else 0

/-- `mvmul`: a row-major `m × n` matrix against a vector. -/
noncomputable def fmvmul (m n : ℕ) (W : Vec (m * n)) (x : Vec n) : Vec m :=
  fun i => ∑ j : Fin n, natGet W (i.val * n + j.val) * x j

/-- `addc`: a per-channel bias broadcast over `h * w` pixels. -/
def faddc (c h w : ℕ) (v : Vec (c * (h * w))) (b : Vec c) : Vec (c * (h * w)) :=
  fun i => v i + natGet b (i.val / (h * w))

/-- `relu`. -/
noncomputable def fvrelu {n : ℕ} (v : Vec n) : Vec n := fun i => max (v i) 0

/-- `tanh`, the bounded log-scale head. -/
noncomputable def fvtanh {n : ℕ} (v : Vec n) : Vec n := fun i => Real.tanh (v i)

/-- `conv1x1`: per-pixel channel mixing. -/
noncomputable def fconv1x1 (cin cout h w : ℕ) (K : Vec (cout * cin))
    (v : Vec (cin * (h * w))) : Vec (cout * (h * w)) :=
  fun i =>
    ∑ ci : Fin cin,
      natGet K (i.val / (h * w) * cin + ci.val) * natGet v (ci.val * (h * w) + i.val % (h * w))

/-- `conv3x3`: SAME-zero-padded cross-correlation; out-of-image offsets contribute `0`. -/
noncomputable def fconv3x3 (cin cout h w : ℕ) (K : Vec (cout * (cin * 9)))
    (v : Vec (cin * (h * w))) : Vec (cout * (h * w)) :=
  fun i =>
    let p := i.val % (h * w)
    ∑ ci : Fin cin, ∑ dy : Fin 3, ∑ dx : Fin 3,
      if 1 ≤ p / w + dy.val ∧ p / w + dy.val ≤ h ∧ 1 ≤ p % w + dx.val ∧ p % w + dx.val ≤ w then
        natGet K (i.val / (h * w) * (cin * 9) + ci.val * 9 + dy.val * 3 + dx.val)
          * natGet v (ci.val * (h * w) + (p / w + dy.val - 1) * w + (p % w + dx.val - 1))
      else 0

/-- `rows2`: the two-row-block split of a matrix. -/
def frows2 (m n : ℕ) (v : Vec ((m + m) * n)) : Vec (m * n) × Vec (m * n) :=
  (fun i => natGet v i.val, fun i => natGet v (m * n + i.val))

/-- Layer norm without gain and bias, with `ε = 10⁻⁵`. -/
noncomputable def flayernorm {n : ℕ} (v : Vec n) : Vec n :=
  let μ := (∑ j, v j) / n
  let var := (∑ j, (v j - μ) ^ 2) / n
  fun i => (v i - μ) / Real.sqrt (var + 1 / 100000)

/-- `gelu`, in its `tanh` approximation. -/
noncomputable def fgelu {n : ℕ} (v : Vec n) : Vec n :=
  fun i => v i / 2 * (1 + Real.tanh (Real.sqrt (2 / Real.pi) * (v i + 0.044715 * v i ^ 3)))

/-- Single-head dot-product attention: key and value `j` are the row `j * h + ·`. -/
noncomputable def fattend (m h : ℕ) (q : Vec h) (K Vs : Vec (h * m)) : Vec h :=
  let s : Fin m → ℝ := fun j => (∑ i : Fin h, q i * natGet K (j.val * h + i.val)) / Real.sqrt h
  fun i => ∑ j : Fin m,
    Real.exp (s j) / (∑ j' : Fin m, Real.exp (s j')) * natGet Vs (j.val * h + i.val)

/-! ## Densities of the reference measures -/

/-- The standard-normal density on `Vec n`. -/
noncomputable def stdNormalPDF (n : ℕ) (x : Vec n) : ℝ :=
  ∏ i, ProbabilityTheory.gaussianPDFReal 0 1 (x i)

theorem stdNormalPDF_pos (n : ℕ) (x : Vec n) : 0 < stdNormalPDF n x :=
  Finset.prod_pos fun i _ => ProbabilityTheory.gaussianPDFReal_pos 0 1 (x i) one_ne_zero

/-- The logistic density `σ'(t) = e⁻ᵗ / (1 + e⁻ᵗ)²` on `Vec n`: the density of a point of
`(0, 1)ⁿ` (uniform) in its logit chart. -/
noncomputable def logisticPDF (n : ℕ) (t : Vec n) : ℝ :=
  ∏ i, Real.exp (-t i) / (1 + Real.exp (-t i)) ^ 2

theorem logisticPDF_pos (n : ℕ) (t : Vec n) : 0 < logisticPDF n t :=
  Finset.prod_pos fun i _ => div_pos (Real.exp_pos _) (by positivity)

theorem measurable_stdNormalPDF (n : ℕ) : Measurable (stdNormalPDF n) :=
  Finset.measurable_prod _ fun i _ =>
    (ProbabilityTheory.measurable_gaussianPDFReal 0 1).comp (measurable_pi_apply i)

theorem measurable_logisticPDF (n : ℕ) : Measurable (logisticPDF n) := by
  unfold logisticPDF; fun_prop

/-- The standard normal on `Vec n`. -/
noncomputable def stdNormal (n : ℕ) : Measure (Vec n) :=
  volume.withDensity fun x => ENNReal.ofReal (stdNormalPDF n x)

/-- The logistic distribution on `Vec n`. -/
noncomputable def stdLogistic (n : ℕ) : Measure (Vec n) :=
  volume.withDensity fun x => ENNReal.ofReal (logisticPDF n x)

instance (n : ℕ) : SigmaFinite (stdNormal n) := by unfold stdNormal; infer_instance
instance (n : ℕ) : SigmaFinite (stdLogistic n) := by unfold stdLogistic; infer_instance

/-- From Lebesgue measure to `volume.withDensity ρ` by the identity, with log-weight `log ρ`. -/
theorem cov_toDensity {n : ℕ} (ρ : Vec n → ℝ) (hm : Measurable ρ) (hp : ∀ x, 0 < ρ x) :
    Cov (volume : Measure (Vec n)) (volume.withDensity fun x => ENNReal.ofReal (ρ x))
      (fun x => x) (fun x => expw (Real.log (ρ x))) := fun ψ hψ => by
  rw [lintegral_withDensity_eq_lintegral_mul _ (f := fun x => ENNReal.ofReal (ρ x))
    (ENNReal.measurable_ofReal.comp hm) hψ]
  refine lintegral_congr fun x => ?_
  simp [expw, Real.exp_log (hp x), mul_comm]

/-! ## Change of variables through the Jacobian -/

/-- A differentiable bijection of a finite-dimensional space with Haar measure satisfies change
of variables with `|det f'|`. -/
theorem cov_of_hasFDerivAt {E : Type*} [NormedAddCommGroup E] [NormedSpace ℝ E]
    [FiniteDimensional ℝ E] [MeasurableSpace E] [BorelSpace E] {μ : Measure E}
    [μ.IsAddHaarMeasure] {f : E → E} {f' : E → E →L[ℝ] E}
    (hf' : ∀ x, HasFDerivAt f (f' x) x) (hf : Function.Bijective f) :
    Cov μ μ f (fun x => ENNReal.ofReal |(f' x).det|) := fun φ _ => by
  have h := lintegral_image_eq_lintegral_abs_det_fderiv_mul μ MeasurableSet.univ
    (fun x _ => (hf' x).hasFDerivWithinAt) hf.injective.injOn φ
  rw [Set.image_univ_of_surjective hf.surjective, Measure.restrict_univ] at h
  rw [h]; congr 1; funext x; ring

/-- An invertible matrix satisfies change of variables with `|det M|`. -/
theorem cov_mulVec {n : ℕ} (M : Matrix (Fin n) (Fin n) ℝ) (h : IsUnit M.det) :
    Cov volume volume M.mulVec (fun _ => ENNReal.ofReal |M.det|) := by
  have hb : Function.Bijective M.mulVec :=
    ⟨Matrix.mulVec_injective_iff_isUnit.mpr ((Matrix.isUnit_iff_isUnit_det M).mpr h),
     Matrix.mulVec_surjective_iff_isUnit.mpr ((Matrix.isUnit_iff_isUnit_det M).mpr h)⟩
  have := cov_of_hasFDerivAt (μ := (volume : Measure (Vec n)))
    (f' := fun _ => LinearMap.toContinuousLinearMap (Matrix.toLin' M))
    (fun x => by
      have hf : M.mulVec = ⇑(LinearMap.toContinuousLinearMap (Matrix.toLin' M)) :=
        funext fun v => by simp [Matrix.toLin'_apply]
      rw [hf]; exact ContinuousLinearMap.hasFDerivAt _) hb
  simpa [LinearMap.det_toLin'] using this

/-- `x ↦ x ⊙ exp s`, with log-weight `Σ s`. -/
theorem cov_scale {n : ℕ} (s : Vec n) :
    Cov volume volume (fun x : Vec n => fun i => x i * Real.exp (s i))
      (fun _ => expw (∑ i, s i)) := by
  have := cov_mulVec (Matrix.diagonal fun i => Real.exp (s i))
    (by simp [Matrix.det_diagonal, Finset.prod_ne_zero_iff, Real.exp_ne_zero])
  refine this.congr (fun x => ?_) fun _ => ?_
  · ext i; simp [Matrix.mulVec_diagonal, mul_comm]
  · simp [Matrix.det_diagonal, expw, Real.exp_sum, abs_of_pos (Finset.prod_pos fun i _ => Real.exp_pos (s i))]

/-- A translation preserves Lebesgue measure. -/
theorem cov_add {n : ℕ} (t : Vec n) :
    Cov volume volume (fun x : Vec n => fun i => x i + t i) (fun _ => 1) :=
  Cov.pres (measurePreserving_add_right volume t)

/-- `v ↦ v ∘ σ` preserves Lebesgue measure. -/
theorem cov_perm {n : ℕ} (σ : Equiv.Perm (Fin n)) :
    Cov volume volume (fun v : Vec n => fun i => v (σ i)) (fun _ => 1) := by
  refine Cov.pres ?_
  convert volume_measurePreserving_piCongrLeft (fun _ : Fin n => ℝ) σ.symm using 1
  funext v i; simp [MeasurableEquiv.piCongrLeft, Equiv.piCongrLeft]

/-! ## The triangular channel map (Glow's LU 1 × 1 convolution) -/

/-- The matrix of `triu c n` on `Vec (c * n)`, `c` channels of `n` pixels. Entry `i` is channel
`i / n` at pixel `i % n`. Within a pixel it is upper triangular, with diagonal `exp U_aa` and
entries `U_ab` above it (`U` row-major `c × c`). -/
noncomputable def triuMat (c n : ℕ) (U : Vec (c * c)) : Matrix (Fin (c * n)) (Fin (c * n)) ℝ :=
  fun i j =>
    if i.val % n = j.val % n then
      if i.val / n = j.val / n then Real.exp (natGet U (i.val / n * c + i.val / n))
      else if i.val / n < j.val / n then natGet U (i.val / n * c + j.val / n) else 0
    else 0

/-- `x ↦ M x`. -/
noncomputable def ftriu (c n : ℕ) (x : Vec (c * n)) (U : Vec (c * c)) : Vec (c * n) :=
  (triuMat c n U).mulVec x

/-- `x ↦ M⁻¹ x`, a triangular solve. -/
noncomputable def ftriuInv (c n : ℕ) (x : Vec (c * n)) (U : Vec (c * c)) : Vec (c * n) :=
  (triuMat c n U)⁻¹.mulVec x

/-- `log det M = Σᵢ log Mᵢᵢ = n · Σ_a U_aa`. -/
noncomputable def triuLogdet (c n : ℕ) (U : Vec (c * c)) : ℝ :=
  ∑ i : Fin (c * n), natGet U (i.val / n * c + i.val / n)

theorem triuMat_blockTriangular {c n : ℕ} (U : Vec (c * c)) :
    (triuMat c n U).BlockTriangular id := by
  intro i j hji
  have hji : j.val < i.val := hji
  unfold triuMat
  split_ifs with h₁ h₂ h₃
  · exfalso
    have : i.val = j.val := by
      rw [← Nat.div_add_mod i.val n, ← Nat.div_add_mod j.val n, h₂, h₁]
    omega
  · exact absurd (Nat.lt_of_div_lt_div h₃) (by omega)
  · rfl
  · rfl

theorem triuMat_det {c n : ℕ} (U : Vec (c * c)) :
    (triuMat c n U).det = Real.exp (triuLogdet c n U) := by
  rw [Matrix.det_of_upperTriangular (triuMat_blockTriangular U), triuLogdet, Real.exp_sum]
  refine Finset.prod_congr rfl fun i _ => ?_
  simp [triuMat]

theorem triuMat_isUnit {c n : ℕ} (U : Vec (c * c)) : IsUnit (triuMat c n U).det := by
  rw [triuMat_det]; exact (Real.exp_pos _).ne'.isUnit

theorem ftriuInv_ftriu {c n : ℕ} (x : Vec (c * n)) (U : Vec (c * c)) :
    ftriuInv c n (ftriu c n x U) U = x := by
  simp [ftriu, ftriuInv, Matrix.mulVec_mulVec, Matrix.nonsing_inv_mul _ (triuMat_isUnit U)]

theorem ftriu_ftriuInv {c n : ℕ} (x : Vec (c * n)) (U : Vec (c * c)) :
    ftriu c n (ftriuInv c n x U) U = x := by
  simp [ftriu, ftriuInv, Matrix.mulVec_mulVec, Matrix.mul_nonsing_inv _ (triuMat_isUnit U)]

theorem cov_triu {c n : ℕ} (U : Vec (c * c)) :
    Cov volume volume (fun x : Vec (c * n) => ftriu c n x U) (fun _ => expw (triuLogdet c n U)) :=
  (cov_mulVec (triuMat c n U) (triuMat_isUnit _)).congr (fun _ => rfl)
    fun _ => by rw [triuMat_det, abs_of_pos (Real.exp_pos _)]; rfl

/-! ## Elementwise odd powers and their roots

`x ↦ x ^ (2k + 1)` for `k : ℕ` (the cube at `k = 1`), a bijection of `ℝ` whose derivative
`(2k + 1) x ^ 2k` vanishes only at `0`. -/

/-- The odd exponent `2k + 1`. -/
abbrev oddExp (k : ℕ) : ℕ := 2 * k + 1

theorem odd_oddExp (k : ℕ) : Odd (oddExp k) := odd_two_mul_add_one k

/-- The real `(2k + 1)`-th root: `x ^ (1/(2k+1))` on `x ≥ 0`, extended oddly. -/
noncomputable def oroot (k : ℕ) (x : ℝ) : ℝ :=
  if 0 ≤ x then x ^ (((oddExp k : ℕ) : ℝ)⁻¹) else -((-x) ^ (((oddExp k : ℕ) : ℝ)⁻¹))

theorem oroot_pow (k : ℕ) (x : ℝ) : oroot k (x ^ oddExp k) = x := by
  unfold oroot
  rcases le_or_gt 0 x with hx | hx
  · rw [if_pos (pow_nonneg hx _), Real.pow_rpow_inv_natCast hx (Nat.succ_ne_zero _)]
  · have h3 : x ^ oddExp k < 0 := Odd.pow_neg (odd_oddExp k) hx
    rw [if_neg (not_le.mpr h3), ← (odd_oddExp k).neg_pow,
      Real.pow_rpow_inv_natCast (by linarith) (Nat.succ_ne_zero _), neg_neg]

theorem pow_oroot (k : ℕ) (y : ℝ) : oroot k y ^ oddExp k = y := by
  unfold oroot
  split_ifs with hy
  · exact Real.rpow_inv_natCast_pow hy (Nat.succ_ne_zero _)
  · rw [(odd_oddExp k).neg_pow, Real.rpow_inv_natCast_pow (by linarith) (Nat.succ_ne_zero _),
      neg_neg]

/-- `x ↦ x ^ (2k + 1)`, elementwise. -/
def fopow (k : ℕ) {n : ℕ} (v : Vec n) : Vec n := fun i => v i ^ oddExp k

/-- `y ↦ y ^ (1/(2k + 1))`, elementwise. -/
noncomputable def foroot (k : ℕ) {n : ℕ} (v : Vec n) : Vec n := fun i => oroot k (v i)

theorem foroot_fopow (k : ℕ) {n : ℕ} (v : Vec n) : foroot k (fopow k v) = v :=
  funext fun i => oroot_pow k (v i)
theorem fopow_foroot (k : ℕ) {n : ℕ} (v : Vec n) : fopow k (foroot k v) = v :=
  funext fun i => pow_oroot k (v i)

/-- `(2k + 1) x ^ 2k`, or `1` at `x = 0`: the derivative off the axes, kept positive so that its
log is finite everywhere. The axes are null, so change of variables holds with it
(`cov_fopow`). -/
noncomputable def opowDet1 (k : ℕ) (x : ℝ) : ℝ :=
  if x = 0 then 1 else (oddExp k : ℝ) * x ^ (2 * k)

theorem opowDet1_pos (k : ℕ) (x : ℝ) : 0 < opowDet1 k x := by
  unfold opowDet1; split_ifs with h
  · exact one_pos
  · have : 0 < x ^ (2 * k) := by rw [pow_mul]; exact pow_pos (by positivity) k
    positivity

/-- `Σᵢ log ((2k + 1) xᵢ ^ 2k)` off the axes: the log-determinant of the elementwise power. -/
noncomputable def opowLogdet (k : ℕ) {n : ℕ} (v : Vec n) : ℝ :=
  ∑ i, Real.log (opowDet1 k (v i))

theorem measurable_oroot (k : ℕ) : Measurable (oroot k) := by
  have hg : Measurable fun x : ℝ => x ^ (((oddExp k : ℕ) : ℝ)⁻¹) :=
    (Real.continuous_rpow_const (by positivity)).measurable
  exact Measurable.ite measurableSet_Ici hg (hg.comp measurable_neg).neg

theorem measurable_fopow (k n : ℕ) : Measurable (fopow k (n := n)) := by unfold fopow; fun_prop

theorem measurable_foroot (k n : ℕ) : Measurable (foroot k (n := n)) :=
  measurable_pi_lambda _ fun i => (measurable_oroot k).comp (measurable_pi_apply i)

theorem measurable_opowDet1 (k : ℕ) : Measurable (opowDet1 k) :=
  Measurable.ite (measurableSet_singleton 0) measurable_const (by fun_prop)

theorem measurable_opowLogdet (k n : ℕ) : Measurable (opowLogdet k (n := n)) :=
  Finset.measurable_sum _ fun i _ =>
    Real.measurable_log.comp ((measurable_opowDet1 k).comp (measurable_pi_apply i))

theorem hasFDerivAt_fopow (k : ℕ) {n : ℕ} (x : Vec n) :
    HasFDerivAt (fopow k) (LinearMap.toContinuousLinearMap
      (Matrix.toLin' (Matrix.diagonal fun i => (oddExp k : ℝ) * x i ^ (2 * k)))) x := by
  rw [hasFDerivAt_pi']
  intro i
  refine ((hasFDerivAt_apply (𝕜 := ℝ) i x).pow (oddExp k)).congr_fderiv ?_
  refine ContinuousLinearMap.ext fun v => ?_
  simp [Matrix.toLin'_apply, Matrix.mulVec_diagonal, oddExp]

/-- The axes are null. -/
theorem ae_ne_zero (n : ℕ) : ∀ᵐ x : Vec n ∂volume, ∀ i, x i ≠ 0 := by
  rw [ae_iff]
  refine measure_mono_null (fun x hx => ?_) (measure_iUnion_null fun i =>
    Measure.addHaar_submodule volume (LinearMap.ker (LinearMap.proj i : (Fin n → ℝ) →ₗ[ℝ] ℝ)) ?_)
  · simp only [not_forall, not_not, Set.mem_iUnion, SetLike.mem_coe, LinearMap.mem_ker,
      LinearMap.proj_apply] at hx ⊢
    exact hx
  · intro h
    have : (Pi.single i (1 : ℝ) : Fin n → ℝ) ∈
        LinearMap.ker (LinearMap.proj i : (Fin n → ℝ) →ₗ[ℝ] ℝ) := h ▸ Submodule.mem_top
    simp at this

/-- `x ↦ x ^ (2k + 1)`, with log-weight `opowLogdet k`. -/
theorem cov_fopow (k n : ℕ) :
    Cov volume volume (fopow k (n := n)) (fun x => expw (opowLogdet k x)) := by
  have hb : Function.Bijective (fopow k (n := n)) :=
    Function.bijective_iff_has_inverse.mpr ⟨foroot k, foroot_fopow k, fopow_foroot k⟩
  have h := cov_of_hasFDerivAt (μ := (volume : Measure (Vec n))) (hasFDerivAt_fopow k) hb
  refine Cov.congr_ae ?_ h
  refine (ae_ne_zero n).mono fun x hx => ?_
  have hp : Real.exp (opowLogdet k x) = ∏ i, (oddExp k : ℝ) * x i ^ (2 * k) := by
    rw [opowLogdet, Real.exp_sum]
    exact Finset.prod_congr rfl fun i _ => by
      rw [Real.exp_log (opowDet1_pos _ _)]; simp [opowDet1, hx i]
  simp only [ContinuousLinearMap.det, LinearMap.coe_toContinuousLinearMap, LinearMap.det_toLin',
    Matrix.det_diagonal, expw]
  rw [← hp, abs_of_pos (Real.exp_pos _)]

/-! ## Measurability of the conditioners -/

@[fun_prop] theorem Measurable.natGet {α : Type*} [MeasurableSpace α] {k : ℕ} {f : α → Vec k}
    (hf : Measurable f) (i : ℕ) : Measurable fun a => natGet (f a) i := by
  unfold RoleFlow.Flow.natGet; split_ifs <;> fun_prop

@[fun_prop] theorem measurable_tanh : Measurable Real.tanh := by
  rw [show Real.tanh = fun x => Real.sinh x / Real.cosh x from funext Real.tanh_eq_sinh_div_cosh]
  fun_prop

@[fun_prop] theorem Continuous.natGet {α : Type*} [TopologicalSpace α] {k : ℕ} {f : α → Vec k}
    (hf : Continuous f) (i : ℕ) : Continuous fun a => natGet (f a) i := by
  unfold RoleFlow.Flow.natGet; split_ifs <;> fun_prop

theorem continuous_triuMat (c n : ℕ) :
    Continuous fun U : Vec (c * c) => triuMat c n U := by
  refine continuous_pi fun i => continuous_pi fun j => ?_
  unfold triuMat; split_ifs <;> fun_prop

theorem continuous_triuMat_inv (c n : ℕ) :
    Continuous fun U : Vec (c * c) => (triuMat c n U)⁻¹ := by
  have h := continuous_triuMat c n
  simp_rw [Matrix.inv_def, Ring.inverse_eq_inv']
  exact (h.matrix_det.inv₀ fun U => (triuMat_isUnit U).ne_zero).smul h.matrix_adjugate

theorem measurable_ftriu (c n : ℕ) :
    Measurable fun p : Vec (c * n) × Vec (c * c) => ftriu c n p.1 p.2 := by
  unfold ftriu
  exact (((continuous_triuMat c n).comp continuous_snd).matrix_mulVec continuous_fst).measurable

theorem measurable_ftriuInv (c n : ℕ) :
    Measurable fun p : Vec (c * n) × Vec (c * c) => ftriuInv c n p.1 p.2 := by
  unfold ftriuInv
  exact (((continuous_triuMat_inv c n).comp continuous_snd).matrix_mulVec
    continuous_fst).measurable

theorem measurable_triuLogdet (c n : ℕ) : Measurable (triuLogdet c n) := by
  unfold triuLogdet; fun_prop

end RoleFlow.Flow