import RoleFlow.Cov

/-! # Signatures and models (`docs/roleflow.tex`, `sec:sig`)

Sorts carry a measurable space and a σ-finite measure each (`Sorts`). Contexts are flat lists
of sorts, and a value of a context `Γ` is a dependent tuple `Val Γ`, measured by the product
`vol Γ`. Values are accessed by position through `Val.get`, which takes the sort as an
argument with a proof `Γ[i]? = some a`, so no caller ever casts between sorts.

A signature (`Sig`) has operations `Op as bs` with `bs ≠ []`, and declares for some of them
one semi-inverse (`SemiInv`): the first `cs` arguments are consumed and solved from the
results and the remaining arguments `ps`. A model (`Model`) interprets operations by
measurable maps and gives each operation a log-weight; for operations with a semi-inverse it
must satisfy the left law and the leaf (change of variables at each parameter). `ModelInv`
adds the right law. `leaf_dag` derives the leaf of a semi-inverse with the derived log-weight
`dagWeight`; it needs only the left law. -/

open MeasureTheory

namespace RoleFlow

/-- Sorts, each with a measurable space and a σ-finite reference measure. -/
class Sorts (S : Type) where
  /-- The space of a sort. -/
  X : S → Type
  meas : ∀ s, MeasurableSpace (X s)
  /-- The reference measure of a sort. -/
  μ : ∀ s, Measure (X s)
  sigmaFinite : ∀ s, SigmaFinite (μ s)

attribute [reducible, instance] Sorts.meas
attribute [instance] Sorts.sigmaFinite

export Sorts (X μ)

/-! ## Signatures -/

/-- A semi-inverse of an operation `Op as bs`: it solves the first arguments `cs` from the
results and the remaining arguments `ps` (the parameters). The split point of the design is
`k = cs.length`, and `cs ≠ []` is `k ≥ 1`. -/
structure SemiInv {S : Type} (Op : List S → List S → Type) (as bs : List S) where
  cs : List S
  ps : List S
  split : cs ++ ps = as
  cs_ne : cs ≠ []
  op : Op (bs ++ ps) cs

/-- A signature: operations, each with at most one semi-inverse. -/
structure Sig (S : Type) where
  Op : List S → List S → Type
  bs_ne : ∀ {as bs}, Op as bs → bs ≠ []
  inv : ∀ {as bs}, Op as bs → Option (SemiInv Op as bs)

/-- The split point: how many leading arguments a semi-inverse consumes. -/
abbrev SemiInv.k {S : Type} {Op : List S → List S → Type} {as bs : List S}
    (si : SemiInv Op as bs) : ℕ := si.cs.length

/-! ## Arguments -/

/-- The arguments `xs` (positions of `Γ`) have the sorts `as`. -/
abbrev Args {S : Type} (Γ : List S) (xs : List ℕ) (as : List S) : Prop :=
  List.Forall₂ (fun x a => Γ[x]? = some a) xs as

section Args
variable {S : Type} {Γ : List S}

theorem Args.lt {xs : List ℕ} {as : List S} (h : Args Γ xs as) {x : ℕ} (hx : x ∈ xs) :
    x < Γ.length := by
  induction h with
  | nil => simp at hx
  | cons h₁ _ ih =>
    rcases List.mem_cons.1 hx with rfl | hx
    · exact (List.getElem?_eq_some_iff.1 h₁).1
    · exact ih hx

/-- Arguments stay valid when the context grows. -/
theorem Args.append {xs : List ℕ} {as : List S} (h : Args Γ xs as) (Δ : List S) :
    Args (Γ ++ Δ) xs as :=
  h.imp fun {x _} hx => by
    rw [List.getElem?_append_left (List.getElem?_eq_some_iff.1 hx).1]; exact hx

theorem Args.lt_length {xs : List ℕ} {as : List S} (h : Args Γ xs as) {j : ℕ}
    (hj : j < as.length) : j < xs.length := by rw [h.length_eq]; exact hj

/-- The `j`-th argument has the `j`-th sort. -/
theorem Args.getElem? {xs : List ℕ} {as : List S} (h : Args Γ xs as) (j : ℕ)
    (hj : j < xs.length) : Γ[xs[j]]? = some (as[j]'(by rw [← h.length_eq]; exact hj)) :=
  (List.forall₂_iff_get.1 h).2 j hj _

/-- The positions `|Γ|, …, |Γ| + |Δ| - 1` of `Γ ++ Δ` have the sorts `Δ`. -/
theorem Args.range' (Γ Δ : List S) : Args (Γ ++ Δ) (List.range' Γ.length Δ.length) Δ := by
  refine List.forall₂_iff_get.2 ⟨by simp, fun j h₁ h₂ => ?_⟩
  simp only [List.get_eq_getElem, List.getElem_range', Nat.one_mul]
  rw [List.getElem?_append_right (by omega), Nat.add_sub_cancel_left, List.getElem?_eq_getElem]

theorem Args.append' {xs ys : List ℕ} {as bs : List S} (h₁ : Args Γ xs as) (h₂ : Args Γ ys bs) :
    Args Γ (xs ++ ys) (as ++ bs) := List.rel_append h₁ h₂

end Args

/-! ## Values of contexts -/

variable {S : Type} [Sorts S]

/-- A value of the context `Γ`: one value per position. -/
abbrev Val (Γ : List S) : Type := (i : Fin Γ.length) → X (Γ.get i)

/-- The product measure on a context. -/
noncomputable def vol (Γ : List S) : Measure (Val Γ) := Measure.pi fun i => μ (Γ.get i)

instance (Γ : List S) : SigmaFinite (vol Γ) := by unfold vol; infer_instance

namespace Val

variable {Γ Δ : List S} {a b : S}

/-- The empty tuple. -/
def nil : Val ([] : List S) := fun i => i.elim0

/-- Prepend a value. -/
def cons (x : X a) (v : Val Γ) : Val (a :: Γ) := Fin.cons (α := fun i => X ((a :: Γ).get i)) x v

def head (v : Val (a :: Γ)) : X a := v 0

def tail (v : Val (a :: Γ)) : Val Γ := fun i => v i.succ

@[simp] theorem head_cons (x : X a) (v : Val Γ) : (cons x v).head = x := rfl
@[simp] theorem tail_cons (x : X a) (v : Val Γ) : (cons x v).tail = v := rfl
@[simp] theorem cons_head_tail (v : Val (a :: Γ)) : cons v.head v.tail = v :=
  Fin.cons_self_tail (α := fun i => X ((a :: Γ).get i)) v

omit [Sorts S] in
theorem lt_of_getElem? {i : ℕ} (h : Γ[i]? = some a) : i < Γ.length :=
  (List.getElem?_eq_some_iff.1 h).1

omit [Sorts S] in
theorem get_eq_of_getElem? {i : ℕ} (h : Γ[i]? = some a) :
    Γ.get ⟨i, lt_of_getElem? h⟩ = a :=
  (List.getElem?_eq_some_iff.1 h).2

/-- The value at position `i`, of the sort `a` that `Γ` has there. -/
def get (v : Val Γ) (i : ℕ) (h : Γ[i]? = some a) : X a :=
  tr X (get_eq_of_getElem? h) (v ⟨i, lt_of_getElem? h⟩)

theorem apply_eq_get (v : Val Γ) (i : Fin Γ.length) : v i = v.get i.1 (by simp) := rfl

@[ext] theorem ext {v w : Val Γ} (h : ∀ (i : ℕ) (a : S) (hi : Γ[i]? = some a), v.get i hi = w.get i hi) :
    v = w := by
  funext i
  rw [apply_eq_get v, apply_eq_get w]; exact h _ _ _

@[simp] theorem get_cons_zero (x : X b) (v : Val Γ) (h : (b :: Γ)[0]? = some a) :
    (cons x v).get 0 h = tr X (by simpa using h) x := rfl

@[simp] theorem get_cons_succ (x : X b) (v : Val Γ) {i : ℕ} (h : (b :: Γ)[i + 1]? = some a) :
    (cons x v).get (i + 1) h = v.get i (by simpa using h) := rfl

/-- Transport along an equality of contexts. -/
def cast (h : Γ = Δ) (v : Val Γ) : Val Δ := h ▸ v

@[simp] theorem cast_rfl (v : Val Γ) : v.cast rfl = v := rfl

@[simp] theorem get_cast (h : Γ = Δ) (v : Val Γ) {i : ℕ} (hi : Δ[i]? = some a) :
    (v.cast h).get i hi = v.get i (h ▸ hi) := by subst h; rfl

/-- Concatenation. -/
def append : {Γ : List S} → Val Γ → Val Δ → Val (Γ ++ Δ)
  | [], _, w => w
  | _ :: _, v, w => cons v.head (append v.tail w)

/-- The part of a value of `Γ ++ Δ` that belongs to `Δ`. -/
def right : {Γ : List S} → Val (Γ ++ Δ) → Val Δ
  | [], w => w
  | _ :: _, w => right (Γ := _) w.tail

@[simp] theorem right_append (v : Val Γ) (w : Val Δ) : (v.append w).right = w := by
  induction Γ with
  | nil => rfl
  | cons b Γ ih => exact ih v.tail

theorem get_append_left (v : Val Γ) (w : Val Δ) {i : ℕ} (hi : i < Γ.length)
    (h : (Γ ++ Δ)[i]? = some a) :
    (v.append w).get i h = v.get i (by rwa [List.getElem?_append_left hi] at h) := by
  induction Γ generalizing i with
  | nil => simp at hi
  | cons b Γ ih =>
    cases i with
    | zero => rfl
    | succ i => exact ih v.tail (by simpa using hi) (by simpa using h)

theorem get_append_right (v : Val Γ) (w : Val Δ) {i : ℕ} (h : (Γ ++ Δ)[Γ.length + i]? = some a) :
    (v.append w).get (Γ.length + i) h =
      w.get i (by rwa [List.getElem?_append_right (by omega), Nat.add_sub_cancel_left] at h) := by
  induction Γ with
  | nil => simp only [List.length_nil, Nat.zero_add]; rfl
  | cons b Γ ih =>
    have e : (b :: Γ).length + i = (Γ.length + i) + 1 := by simp; omega
    simp only [e] at h ⊢
    exact ih v.tail (by simpa using h)

/-- Select the positions `xs`. -/
def pick (v : Val Γ) : (xs : List ℕ) → {as : List S} → Args Γ xs as → Val as
  | [], [], _ => nil
  | x :: xs, _ :: _, h => cons (v.get x (List.forall₂_cons.1 h).1) (pick v xs (List.forall₂_cons.1 h).2)
  | [], _ :: _, h => absurd h (by simp)
  | _ :: _, [], h => absurd h (by simp)

/-! ### Picking -/

theorem tr_get (v : Val Γ) (e : a = b) {i : ℕ} (h : Γ[i]? = some a) :
    tr X e (v.get i h) = v.get i (e ▸ h) := by subst e; rfl

theorem get_congr (v : Val Γ) {i j : ℕ} (e : i = j) (hi : Γ[i]? = some a) (hj : Γ[j]? = some a) :
    v.get i hi = v.get j hj := by subst e; rfl

theorem get_pick (v : Val Γ) {xs : List ℕ} {as : List S} (h : Args Γ xs as) {j : ℕ}
    (hj : as[j]? = some a) :
    (v.pick xs h).get j hj = v.get (xs[j]'(h.lt_length (lt_of_getElem? hj)))
      (by have := h.getElem? j (h.lt_length (lt_of_getElem? hj))
          rwa [(List.getElem?_eq_some_iff.1 hj).2] at this) := by
  induction xs generalizing as j with
  | nil => cases as with
    | nil => simp at hj
    | cons => simp at h
  | cons x xs ih => cases as with
    | nil => simp at h
    | cons b as =>
      cases j with
      | zero => simp only [pick, get_cons_zero, tr_get]; rfl
      | succ j => simp only [pick, get_cons_succ]; exact ih _ _

/-- Picking from two values that agree at the picked positions. -/
theorem pick_eq_of_get {Γ' : List S} (v : Val Γ) (w : Val Γ') {xs : List ℕ} {as : List S}
    (h : Args Γ xs as) (h' : Args Γ' xs as)
    (H : ∀ x ∈ xs, ∀ (a : S) (hx : Γ[x]? = some a) (hx' : Γ'[x]? = some a), v.get x hx = w.get x hx') :
    v.pick xs h = w.pick xs h' := by
  induction xs generalizing as with
  | nil => cases as with
    | nil => rfl
    | cons => simp at h
  | cons x xs ih => cases as with
    | nil => simp at h
    | cons b as =>
      simp only [pick]
      rw [H x (List.mem_cons_self ..), ih _ _ fun y hy => H y (List.mem_cons_of_mem _ hy)]

theorem pick_append {xs ys : List ℕ} {as bs : List S} (v : Val Γ) (h₁ : Args Γ xs as)
    (h₂ : Args Γ ys bs) : v.pick (xs ++ ys) (h₁.append' h₂) = (v.pick xs h₁).append (v.pick ys h₂) := by
  induction xs generalizing as with
  | nil => cases as with
    | nil => rfl
    | cons => simp at h₁
  | cons x xs ih => cases as with
    | nil => simp at h₁
    | cons b as =>
      simp only [List.cons_append, pick, append, head_cons, tail_cons]
      rw [← ih (List.forall₂_cons.1 h₁).2]

theorem pick_cast (e : Γ = Δ) (v : Val Γ) {xs : List ℕ} {as : List S} (h : Args Δ xs as) :
    (v.cast e).pick xs h = v.pick xs (e ▸ h) := by subst e; rfl

/-- Picking the positions of `Δ` from a value of `Γ ++ Δ`. -/
theorem pick_range'_append (v : Val Γ) (w : Val Δ) :
    (v.append w).pick (List.range' Γ.length Δ.length) (Args.range' Γ Δ) = w := by
  ext j a hj
  have e : (List.range' Γ.length Δ.length)[j]'(by simpa using lt_of_getElem? hj) = Γ.length + j :=
    by simp
  rw [get_pick, get_congr _ e _ (by rwa [List.getElem?_append_right (by omega),
    Nat.add_sub_cancel_left]), get_append_right]

/-! ### Measurability -/

@[fun_prop] theorem measurable_get {i : ℕ} (h : Γ[i]? = some a) :
    Measurable fun v : Val Γ => v.get i h :=
  (measurable_tr _).comp (measurable_pi_apply _)

@[fun_prop] theorem measurable_head : Measurable fun v : Val (a :: Γ) => v.head :=
  measurable_pi_apply 0

@[fun_prop] theorem measurable_tail : Measurable fun v : Val (a :: Γ) => v.tail :=
  measurable_pi_lambda _ fun i => measurable_pi_apply i.succ

@[fun_prop] theorem measurable_cons : Measurable fun p : X a × Val Γ => cons p.1 p.2 := by
  refine measurable_pi_lambda _ fun i => ?_
  refine Fin.cases ?_ (fun j => ?_) i
  · exact measurable_fst
  · exact (measurable_pi_apply j).comp measurable_snd

/-- `cons` of measurable maps, in the form `fun_prop` uses. -/
@[fun_prop] theorem measurable_cons' {Y : Type*} [MeasurableSpace Y] {g : Y → X a} {h : Y → Val Γ}
    (hg : Measurable g) (hh : Measurable h) : Measurable fun y => cons (g y) (h y) :=
  measurable_cons.comp (hg.prodMk hh)

@[fun_prop] theorem measurable_cast (h : Γ = Δ) : Measurable (cast h) := by
  subst h; exact measurable_id

@[fun_prop] theorem measurable_append : Measurable fun p : Val Γ × Val Δ => p.1.append p.2 := by
  induction Γ with
  | nil => exact measurable_snd
  | cons b Γ ih =>
    exact measurable_cons.comp ((measurable_head.comp measurable_fst).prodMk
      (ih.comp ((measurable_tail.comp measurable_fst).prodMk measurable_snd)))

@[fun_prop] theorem measurable_right : Measurable fun w : Val (Γ ++ Δ) => w.right := by
  induction Γ with
  | nil => exact measurable_id
  | cons b Γ ih => exact ih.comp measurable_tail

@[fun_prop] theorem measurable_pick {xs : List ℕ} {as : List S} (h : Args Γ xs as) :
    Measurable fun v : Val Γ => v.pick xs h := by
  induction xs generalizing as with
  | nil => cases as with
    | nil => exact measurable_const
    | cons => simp at h
  | cons x xs ih => cases as with
    | nil => simp at h
    | cons a as =>
      exact measurable_cons.comp ((measurable_get (List.forall₂_cons.1 h).1).prodMk (ih _))

/-! ### Products over contexts -/

theorem measurePreserving_uncons :
    MeasurePreserving (fun v : Val (a :: Γ) => (v.head, v.tail)) (vol (a :: Γ)) ((μ a).prod (vol Γ)) := by
  have := measurePreserving_piFinSuccAbove (fun i => μ ((a :: Γ).get i)) 0
  refine ⟨measurable_head.prodMk measurable_tail, ?_⟩
  rw [show (fun v : Val (a :: Γ) => (v.head, v.tail)) =
      MeasurableEquiv.piFinSuccAbove (fun i => X ((a :: Γ).get i)) 0 from ?_]
  · exact this.map_eq
  · funext v; refine Prod.ext rfl ?_; funext j; rfl

theorem eq_nil (v : Val ([] : List S)) : v = nil := funext fun i => i.elim0

theorem vol_nil : vol ([] : List S) = Measure.dirac nil := by
  haveI : IsEmpty (Fin ([] : List S).length) := inferInstanceAs (IsEmpty (Fin 0))
  unfold vol; rw [Measure.pi_of_empty (x := nil)]

instance : IsProbabilityMeasure (vol ([] : List S)) := by rw [vol_nil]; infer_instance

theorem measurePreserving_cons :
    MeasurePreserving (fun p : X a × Val Γ => cons p.1 p.2) ((μ a).prod (vol Γ)) (vol (a :: Γ)) :=
  measurePreserving_uncons.of_inverse measurable_cons fun _ => cons_head_tail _

/-- A one-element context is its sort. -/
theorem measurePreserving_head :
    MeasurePreserving (fun v : Val [a] => v.head) (vol [a]) (μ a) :=
  measurePreserving_fst.comp measurePreserving_uncons

theorem measurePreserving_single :
    MeasurePreserving (fun x : X a => cons x nil) (μ a) (vol [a]) :=
  measurePreserving_head.of_inverse (measurable_cons.comp (measurable_id.prodMk measurable_const))
    fun v => by rw [← cons_head_tail v, eq_nil v.tail]; rfl

end Val

/-! ## Models -/

/-- Glue the consumed arguments and the parameters into the argument tuple. -/
def SemiInv.glue {Op : List S → List S → Type} {as bs : List S} (si : SemiInv Op as bs)
    (c : Val si.cs) (p : Val si.ps) : Val as :=
  (c.append p).cast si.split

theorem SemiInv.measurable_glue {Op : List S → List S → Type} {as bs : List S}
    (si : SemiInv Op as bs) : Measurable fun q : Val si.cs × Val si.ps => si.glue q.1 q.2 := by
  unfold glue; fun_prop

omit [Sorts S] in
theorem SemiInv.args_take {Op : List S → List S → Type} {as bs : List S} (si : SemiInv Op as bs)
    {Γ : List S} {xs : List ℕ} (h : Args Γ xs as) : Args Γ (xs.take si.k) si.cs :=
  by
  rw [← si.split] at h; exact List.forall₂_take_append _ _ _ h

omit [Sorts S] in
theorem SemiInv.args_drop {Op : List S → List S → Type} {as bs : List S} (si : SemiInv Op as bs)
    {Γ : List S} {xs : List ℕ} (h : Args Γ xs as) : Args Γ (xs.drop si.k) si.ps :=
  by
  rw [← si.split] at h; exact List.forall₂_drop_append _ _ _ h

/-- Gluing the consumed arguments and the parameters of a picked argument tuple. -/
theorem SemiInv.glue_pick {Op : List S → List S → Type} {as bs : List S} (si : SemiInv Op as bs)
    {Γ : List S} (v : Val Γ) {xs : List ℕ} (h : Args Γ xs as) :
    si.glue (v.pick _ (si.args_take h)) (v.pick _ (si.args_drop h)) = v.pick xs h := by
  have h₁ : ∀ (as' : List S) (e : as' = as) (ys : List ℕ) (H : Args Γ ys as'),
      (v.pick ys H).cast e = v.pick ys (e ▸ H) := by rintro _ rfl _ _; rfl
  have h₂ : ∀ (ys : List ℕ) (e : ys = xs) (H : Args Γ ys as), v.pick ys H = v.pick xs h := by
    rintro _ rfl _; rfl
  rw [glue, ← Val.pick_append, h₁ _ si.split, h₂ _ (List.take_append_drop _ _)]

/-- A model: measurable interpretations and log-weights; for every operation with a
semi-inverse, the left law and the leaf. The log-weight `ℓ o` is read on the whole argument
tuple, `ℓ_o(c, p) = ℓ o (glue c p)`; it is irrelevant for operations without a semi-inverse. -/
structure Model (Sg : Sig S) where
  sem : ∀ {as bs}, Sg.Op as bs → Val as → Val bs
  ℓ : ∀ {as bs}, Sg.Op as bs → Val as → ℝ
  sem_meas : ∀ {as bs} (o : Sg.Op as bs), Measurable (sem o)
  ℓ_meas : ∀ {as bs} (o : Sg.Op as bs), Measurable (ℓ o)
  left : ∀ {as bs} (o : Sg.Op as bs) (si : SemiInv Sg.Op as bs), Sg.inv o = some si →
    ∀ c p, sem si.op ((sem o (si.glue c p)).append p) = c
  leaf : ∀ {as bs} (o : Sg.Op as bs) (si : SemiInv Sg.Op as bs), Sg.inv o = some si →
    ∀ p, Cov (vol si.cs) (vol bs) (fun c => sem o (si.glue c p)) (fun c => expw (ℓ o (si.glue c p)))

/-- The right law of a model: every semi-inverse is also a right inverse. -/
def Model.RightLaw {Sg : Sig S} (M : Model Sg) : Prop :=
  ∀ {as bs} (o : Sg.Op as bs) (si : SemiInv Sg.Op as bs), Sg.inv o = some si →
    ∀ b p, M.sem o (si.glue (M.sem si.op (b.append p)) p) = b

/-- A model with the right law as well (`Model±` of the design). -/
structure ModelInv (Sg : Sig S) extends Model Sg where
  right : toModel.RightLaw

/-! ## Derived log-weights of semi-inverses -/

section Dag
variable {as bs cs ps : List S}

/-- The derived log-weight of a semi-inverse `G` of `F`, on `G`'s argument tuple `(b, p)`:
`-ℓ_F(G(b, p), p)`. -/
def dagWeight (G : Val (bs ++ ps) → Val cs) (ℓ : Val as → ℝ) (split : cs ++ ps = as)
    (w : Val (bs ++ ps)) : ℝ :=
  -ℓ (((G w).append w.right).cast split)

theorem measurable_dagWeight {G : Val (bs ++ ps) → Val cs} {ℓ : Val as → ℝ}
    (hG : Measurable G) (hℓ : Measurable ℓ) (split : cs ++ ps = as) :
    Measurable (dagWeight G ℓ split) := by
  unfold dagWeight
  have : Measurable fun w : Val (bs ++ ps) => (G w, w.right) := hG.prodMk Val.measurable_right
  exact (hℓ.comp ((Val.measurable_cast split).comp (Val.measurable_append.comp this))).neg

/-- The leaf of a semi-inverse, with the derived log-weight, by **Inv**. Only the left law of
the forward operation is used. -/
theorem leaf_dag {F : Val as → Val bs} {G : Val (bs ++ ps) → Val cs} {ℓ : Val as → ℝ}
    {split : cs ++ ps = as} (hG : Measurable G) (hℓ : Measurable ℓ)
    (left : ∀ c p, G ((F ((c.append p).cast split)).append p) = c)
    (leaf : ∀ p, Cov (vol cs) (vol bs) (fun c => F ((c.append p).cast split))
      (fun c => expw (ℓ ((c.append p).cast split))))
    (p : Val ps) :
    Cov (vol bs) (vol cs) (fun b => G (b.append p)) (fun b => expw (dagWeight G ℓ split (b.append p))) := by
  have hG' : Measurable fun b : Val bs => G (b.append p) :=
    hG.comp (Val.measurable_append.comp (measurable_id.prodMk measurable_const))
  have hℓ' : Measurable fun c : Val cs => ℓ ((c.append p).cast split) :=
    hℓ.comp ((Val.measurable_cast split).comp
      (Val.measurable_append.comp (measurable_id.prodMk measurable_const)))
  exact (Cov.inv hG' hℓ' (fun c => left c p) (leaf p)).congr (fun _ => rfl)
    (fun b => by simp [dagWeight])

end Dag

end RoleFlow