import RoleFlow.Inverse

/-! # Function calls: a signature extended by a typed program

A typed program `P` (roles `ρ`), whose inverse program `P† = P.inverse ρ` is typed too (roles
`ρi`), becomes two operations of a new signature (`Sig.declare`): `fwd`, on `P`'s inputs then
its weights, and `inv`, on its outputs then its weights. Each is the other's semi-inverse, with
the inputs (outputs) consumed and the weights as parameters. Every operation of the old signature
stays, as `prim`.

The model of the new signature (`ModelInv.link`) interprets `fwd` by `⟦P⟧` with log-weight
`L_θ` of `P`, and `inv` by `⟦P†⟧` with the log-weight of `P†`; its laws are theorems about every
typed program: the left laws are `inverse_sem` and `sem_inverse`, the leaves are `density` for
`P` and for `P†`. So a call needs nothing beyond the typings of the callee and its inverse.

Definitions are added one at a time (`Linked`), each over the signature of the ones before it:
a tower `Sg₀`, `Sg₀.declare P₁`, `(Sg₀.declare P₁).declare P₂`, …, so a body may call any
function defined before it, and none can recurse. A call to a function `k` definitions back is
its `fwd` under `k` `prim`s (`Linked.find`). `Linked.model` is the model at the top, and every
theorem about typed programs holds for programs with calls (`Linked.density`,
`Linked.bijective`). -/

open MeasureTheory

namespace RoleFlow

variable {S : Type}

/-! ## The declared operations -/

/-- The operations of `Sg` with `P` declared: `Sg`'s own, and `P`'s forward and inverse. -/
inductive DeclOp (Sg : Sig S) (P : Prog Sg) : List S → List S → Type
  | prim {as bs : List S} : Sg.Op as bs → DeclOp Sg P as bs
  | fwd : DeclOp Sg P (P.ins ++ P.wts) P.outs
  | inv : DeclOp Sg P (P.outs ++ P.wts) P.ins

/-- A semi-inverse of `Sg` as one of the extended signature. -/
def SemiInv.prim {Sg : Sig S} {P : Prog Sg} {as bs : List S} (si : SemiInv Sg.Op as bs) :
    SemiInv (DeclOp Sg P) as bs :=
  ⟨si.cs, si.ps, si.split, si.cs_ne, .prim si.op⟩

/-- `Sg` with `P` declared. -/
def Sig.declare (Sg : Sig S) (P : Prog Sg) (hi : P.ins ≠ []) (ho : P.outs ≠ []) : Sig S where
  Op := DeclOp Sg P
  bs_ne
    | .prim o => Sg.bs_ne o
    | .fwd => ho
    | .inv => hi
  inv
    | .prim o => (Sg.inv o).map SemiInv.prim
    | .fwd => some ⟨P.ins, P.wts, rfl, hi, .inv⟩
    | .inv => some ⟨P.outs, P.wts, rfl, ho, .fwd⟩

/-- An operation of `Sg` applied to every line. -/
def Prog.mapOp {Sg Sg' : Sig S} (f : ∀ {as bs : List S}, Sg.Op as bs → Sg'.Op as bs)
    (P : Prog Sg) : Prog Sg' :=
  { P with lines := P.lines.map fun l => ⟨f l.op, l.args⟩ }

/-! ## Values of concatenated contexts -/

namespace Val
variable [Sorts S] {Γ Δ : List S}

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

@[simp] theorem left_append (v : Val Γ) (w : Val Δ) : (v.append w).left = v := by
  induction Γ with
  | nil => exact (eq_nil v).symm
  | cons b Γ ih => exact (congrArg (cons v.head) (ih v.tail)).trans (cons_head_tail v)

@[fun_prop] theorem measurable_left : Measurable fun w : Val (Γ ++ Δ) => w.left := by
  induction Γ with
  | nil => exact measurable_const
  | cons b Γ ih => exact measurable_cons.comp (measurable_head.prodMk (ih.comp measurable_tail))

end Val

/-! ## The model -/

section Link
variable [Sorts S] {Sg : Sig S} (M : ModelInv Sg) (P : Prog Sg) {ρ ρi : List Role}
  (hT : P.Typed ρ) (hTi : (P.inverse ρ).Typed ρi) (hi : P.ins ≠ []) (ho : P.outs ≠ [])

/-- The interpretation: `fwd` is `⟦P⟧`, `inv` is `⟦P†⟧`, on the inputs (outputs), then the
weights. -/
def linkSem : ∀ {as bs : List S}, DeclOp Sg P as bs → Val as → Val bs
  | _, _, .prim o => M.sem o
  | _, _, .fwd => fun v => P.sem M.toModel hT.wf v.right v.left
  | _, _, .inv => fun v => (P.inverse ρ).sem M.toModel hTi.wf v.right v.left

/-- The log-weights: `L_θ` of `P` for `fwd`, of `P†` for `inv`. -/
noncomputable def linkℓ : ∀ {as bs : List S}, DeclOp Sg P as bs → Val as → ℝ
  | _, _, .prim o => M.ℓ o
  | _, _, .fwd => fun v => P.logw M.toModel hT.wf ρ v.right v.left
  | _, _, .inv => fun v => (P.inverse ρ).logw M.toModel hTi.wf ρi v.right v.left

theorem measurable_linkSem {as bs : List S} (o : DeclOp Sg P as bs) :
    Measurable (linkSem M P hT hTi o) := by
  cases o with
  | prim o => exact M.sem_meas o
  | fwd => exact (P.measurable_sem _ hT.wf).comp (Val.measurable_right.prodMk Val.measurable_left)
  | inv =>
    exact ((P.inverse ρ).measurable_sem _ hTi.wf).comp (Val.measurable_right.prodMk Val.measurable_left)

theorem measurable_linkℓ {as bs : List S} (o : DeclOp Sg P as bs) :
    Measurable (linkℓ M P hT hTi o) := by
  cases o with
  | prim o => exact M.ℓ_meas o
  | fwd => exact (P.measurable_logw _ hT.wf ρ).comp (Val.measurable_right.prodMk Val.measurable_left)
  | inv =>
    exact ((P.inverse ρ).measurable_logw _ hTi.wf ρi).comp
      (Val.measurable_right.prodMk Val.measurable_left)

/-- **The model of `Sg` with `P` declared**: the laws of `fwd` and `inv` are concordance
(`inverse_sem`, `sem_inverse`) and density, for `P` and for `P†`. -/
noncomputable def ModelInv.link : ModelInv (Sg.declare P hi ho) where
  sem := linkSem M P hT hTi
  ℓ := linkℓ M P hT hTi
  sem_meas := measurable_linkSem M P hT hTi
  ℓ_meas := measurable_linkℓ M P hT hTi
  left := by
    intro as bs o si h c p
    cases o with
    | prim o =>
      obtain ⟨si₀, h₀, rfl⟩ := Option.map_eq_some_iff.1 (show (Sg.inv o).map SemiInv.prim = _ from h)
      exact M.left o si₀ h₀ c p
    | fwd =>
      cases h
      simp only [linkSem, SemiInv.glue, Val.cast_rfl, Val.left_append, Val.right_append]
      exact Prog.inverse_sem _ hT p c
    | inv =>
      cases h
      simp only [linkSem, SemiInv.glue, Val.cast_rfl, Val.left_append, Val.right_append]
      exact Prog.sem_inverse M.right hT p c
  leaf := by
    intro as bs o si h p
    cases o with
    | prim o =>
      obtain ⟨si₀, h₀, rfl⟩ := Option.map_eq_some_iff.1 (show (Sg.inv o).map SemiInv.prim = _ from h)
      exact M.leaf o si₀ h₀ p
    | fwd =>
      cases h
      simp only [linkSem, linkℓ, SemiInv.glue, Val.cast_rfl, Val.left_append, Val.right_append]
      exact P.density _ hT p
    | inv =>
      cases h
      simp only [linkSem, linkℓ, SemiInv.glue, Val.cast_rfl, Val.left_append, Val.right_append]
      exact (P.inverse ρ).density _ hTi p
  right := by
    intro as bs o si h b p
    cases o with
    | prim o =>
      obtain ⟨si₀, h₀, rfl⟩ := Option.map_eq_some_iff.1 (show (Sg.inv o).map SemiInv.prim = _ from h)
      exact M.right o si₀ h₀ b p
    | fwd =>
      cases h
      simp only [linkSem, SemiInv.glue, Val.cast_rfl, Val.left_append, Val.right_append]
      exact Prog.sem_inverse M.right hT p b
    | inv =>
      cases h
      simp only [linkSem, SemiInv.glue, Val.cast_rfl, Val.left_append, Val.right_append]
      exact Prog.inverse_sem _ hT p b

end Link

/-! ## The tower of definitions -/

/-- Definitions added one at a time over the base signature `Sg₀`, each a typed program (with a
typed inverse program) over the signature of the ones before it; `Sg` is the signature with all
of them declared. -/
inductive Linked (Sg₀ : Sig S) : Sig S → Type 1
  | base : Linked Sg₀ Sg₀
  | add {Sg : Sig S} (L : Linked Sg₀ Sg) (name : String) (P : Prog Sg) (ρ ρi : List Role)
      (hT : P.Typed ρ) (hTi : (P.inverse ρ).Typed ρi) (hi : P.ins ≠ []) (ho : P.outs ≠ []) :
      Linked Sg₀ (Sg.declare P hi ho)

/-- What an operation of the tower is: an operation of the base, or a call of a definition by
name, forward or inverse, with the number `k` of its leading arguments that are values (the rest
are its weights). -/
inductive Callee (Op : List S → List S → Type) (as bs : List S)
  | base : Op as bs → Callee Op as bs
  | call (name : String) (inv : Bool) (k : ℕ) : Callee Op as bs

namespace Linked
variable {Sg₀ : Sig S}

/-- An operation of the base as one of the tower (under `prim`s). -/
def lift : {Sg : Sig S} → Linked Sg₀ Sg → ∀ {as bs : List S}, Sg₀.Op as bs → Sg.Op as bs
  | _, .base, _, _, o => o
  | _, .add L .., _, _, o => DeclOp.prim (L.lift o)

/-- The forward operation of the definition `name` (the latest of that name). -/
def find : {Sg : Sig S} → Linked Sg₀ Sg → String → Option ((as : List S) × (bs : List S) × Sg.Op as bs)
  | _, .base, _ => none
  | _, .add L name' _ .., name =>
    if name = name' then some ⟨_, _, DeclOp.fwd⟩
    else (L.find name).map fun ⟨_, _, o⟩ => ⟨_, _, DeclOp.prim o⟩

/-- What an operation of the tower is (`Callee`). -/
def view : {Sg : Sig S} → Linked Sg₀ Sg → ∀ {as bs : List S}, Sg.Op as bs → Callee Sg₀.Op as bs
  | _, .base, _, _, o => .base o
  | _, .add L _ _ _ _ _ _ _ _, _, _, .prim o => L.view o
  | _, .add _ name P _ _ _ _ _ _, _, _, .fwd => .call name false P.ins.length
  | _, .add _ name P _ _ _ _ _ _, _, _, .inv => .call name true P.outs.length

/-- The model of the tower over a model of the base (`ModelInv.link` at each definition). -/
noncomputable def model [Sorts S] (M : ModelInv Sg₀) : {Sg : Sig S} → Linked Sg₀ Sg → ModelInv Sg
  | _, .base => M
  | _, .add L _ P _ _ hT hTi hi ho => (L.model M).link P hT hTi hi ho

variable [Sorts S] {Sg : Sig S} (L : Linked Sg₀ Sg) (M : ModelInv Sg₀) (P : Prog Sg) {ρ : List Role}
  (hT : P.Typed ρ)

/-- **Density** of a typed program with calls, under the model of the tower. -/
theorem density (θ : Val P.wts) :
    Cov (vol P.ins) (vol P.outs) (P.sem (L.model M).toModel hT.wf θ)
      fun x => expw (P.logw (L.model M).toModel hT.wf ρ θ x) :=
  P.density _ hT θ

/-- **Bijectivity** of a typed program with calls, with its inverse program as the inverse. -/
theorem bijective (θ : Val P.wts) : Function.Bijective (P.sem (L.model M).toModel hT.wf θ) :=
  Prog.bijective (L.model M).right hT θ

end Linked

end RoleFlow