import RoleFlow.Sig

/-! # Programs and their semantics (`docs/roleflow.tex`, `sec:prog`)

A program has a weight context `wts` (Θ), an input context `ins` (I), output sorts `outs` (O),
a list of lines `ys = op(args)` and the returned positions `ret`. Variables are named by their
position: weights, then inputs, then the results of each line in order. The body
`let ys = o(xs) in e` of the design is the list of lines followed by `ret`.

Well-formedness (`Prog.WF`, decidable) asks that every argument is in scope with the sort its
operation expects and that the returned variables have the output sorts. The trace semantics
extends the value of the context line by line (`trace`); the semantics `Prog.sem` picks the
returned positions of the final trace, and `Prog.logw` sums the log-weights of the lines whose
role is `spine`. All three are measurable jointly in the weights and the inputs. -/

open MeasureTheory

namespace RoleFlow

variable {S : Type} {Sg : Sig S}

/-- A line `ys = op(args)`; its results `ys` get the next free positions. -/
structure Line (Sg : Sig S) where
  {as bs : List S}
  op : Sg.Op as bs
  args : List ℕ

/-- A program `(Θ, I, e)` in A-normal form, with variables named by position. -/
structure Prog (Sg : Sig S) where
  wts : List S
  ins : List S
  outs : List S
  lines : List (Line Sg)
  ret : List ℕ

/-- The role of a line: a spine step or a conditioner. -/
inductive Role
  | spine
  | cond
  deriving DecidableEq, Repr

/-- The context after some lines. -/
def ctxAfter : List S → List (Line Sg) → List S
  | Γ, [] => Γ
  | Γ, l :: ls => ctxAfter (Γ ++ l.bs) ls

/-- Every line's arguments are in scope, with the sorts its operation expects. -/
def LinesWF : List S → List (Line Sg) → Prop
  | _, [] => True
  | Γ, l :: ls => Args Γ l.args l.as ∧ LinesWF (Γ ++ l.bs) ls

instance decLinesWF [DecidableEq S] : (Γ : List S) → (ls : List (Line Sg)) → Decidable (LinesWF Γ ls)
  | _, [] => inferInstanceAs (Decidable True)
  | Γ, l :: ls => have := decLinesWF (Γ ++ l.bs) ls; inferInstanceAs (Decidable (_ ∧ _))

theorem ctxAfter_append (Γ : List S) (ls : List (Line Sg)) : ∃ Δ, ctxAfter Γ ls = Γ ++ Δ := by
  induction ls generalizing Γ with
  | nil => exact ⟨[], by simp [ctxAfter]⟩
  | cons l ls ih =>
    obtain ⟨Δ, h⟩ := ih (Γ ++ l.bs)
    exact ⟨l.bs ++ Δ, by simp [ctxAfter, h]⟩

theorem length_ctxAfter : ∀ (Γ : List S) (ls : List (Line Sg)),
    (ctxAfter Γ ls).length = Γ.length + (ls.map fun l => l.bs.length).sum
  | _, [] => by simp [ctxAfter]
  | Γ, l :: ls => by simp only [ctxAfter, length_ctxAfter _ ls]; simp; omega

theorem getElem?_ctxAfter {Γ : List S} (ls : List (Line Sg)) {i : ℕ} (hi : i < Γ.length) :
    (ctxAfter Γ ls)[i]? = Γ[i]? := by
  obtain ⟨Δ, e⟩ := ctxAfter_append Γ ls; rw [e, List.getElem?_append_left hi]

/-- Arguments stay valid after more lines. -/
theorem Args.ctxAfter {Γ : List S} {xs : List ℕ} {as : List S} (h : Args Γ xs as)
    (ls : List (Line Sg)) : Args (ctxAfter Γ ls) xs as := by
  obtain ⟨Δ, e⟩ := ctxAfter_append Γ ls; rw [e]; exact h.append Δ

namespace Prog

/-- The initial context: weights, then inputs. -/
abbrev ctx₀ (P : Prog Sg) : List S := P.wts ++ P.ins

/-- The context of all variables. -/
abbrev ctx (P : Prog Sg) : List S := ctxAfter P.ctx₀ P.lines

/-- A well-formed program. -/
def WF (P : Prog Sg) : Prop := LinesWF P.ctx₀ P.lines ∧ Args P.ctx P.ret P.outs

instance [DecidableEq S] (P : Prog Sg) : Decidable P.WF := inferInstanceAs (Decidable (_ ∧ _))

/-- The number of weights. -/
abbrev nw (P : Prog Sg) : ℕ := P.wts.length

/-- The input positions. -/
def inputs (P : Prog Sg) : Finset ℕ := Finset.Ico P.wts.length (P.wts.length + P.ins.length)

end Prog

/-! ## Trace semantics -/

variable [Sorts S] (M : Model Sg)

/-- One line: extend the value of the context by the line's results. -/
def stepLine (l : Line Sg) {Γ : List S} (h : Args Γ l.args l.as) (v : Val Γ) : Val (Γ ++ l.bs) :=
  v.append (M.sem l.op (v.pick l.args h))

/-- The trace: run the lines from a value of the context they start in. -/
def trace : {Γ : List S} → (ls : List (Line Sg)) → LinesWF Γ ls → Val Γ → Val (ctxAfter Γ ls)
  | _, [], _, v => v
  | _, l :: ls, h, v => trace ls h.2 (stepLine M l h.1 v)

/-- The log-weight of the lines with role `spine`, from a value of the context they start in.
Missing roles count as conditioners. -/
noncomputable def logwLines : {Γ : List S} → (ls : List (Line Sg)) → LinesWF Γ ls → List Role →
    Val Γ → ℝ
  | _, [], _, _, _ => 0
  | _, _ :: _, _, [], _ => 0
  | _, l :: ls, h, r :: ρ, v =>
    (if r = .spine then M.ℓ l.op (v.pick l.args h.1) else 0) + logwLines ls h.2 ρ (stepLine M l h.1 v)

namespace Prog
variable (P : Prog Sg) (h : P.WF)

/-- `tr_θ(x)`: the value of every variable. -/
def tr (θ : Val P.wts) (x : Val P.ins) : Val P.ctx := trace M P.lines h.1 (θ.append x)

/-- `⟦P⟧_θ(x)`: the returned variables. -/
def sem (θ : Val P.wts) (x : Val P.ins) : Val P.outs := (P.tr M h θ x).pick P.ret h.2

/-- `L_θ(x)`: the sum of the log-weights of the spine steps under the roles `ρ`. -/
noncomputable def logw (ρ : List Role) (θ : Val P.wts) (x : Val P.ins) : ℝ :=
  logwLines M P.lines h.1 ρ (θ.append x)

end Prog

/-! ## Measurability -/

theorem measurable_stepLine (l : Line Sg) {Γ : List S} (h : Args Γ l.args l.as) :
    Measurable (stepLine M l h) := by
  unfold stepLine
  have := M.sem_meas l.op
  fun_prop

theorem measurable_trace : {Γ : List S} → (ls : List (Line Sg)) → (h : LinesWF Γ ls) →
    Measurable (trace M ls h)
  | _, [], _ => measurable_id
  | _, l :: ls, h => (measurable_trace ls h.2).comp (measurable_stepLine M l h.1)

theorem measurable_logwLines : {Γ : List S} → (ls : List (Line Sg)) → (h : LinesWF Γ ls) →
    (ρ : List Role) → Measurable (logwLines M ls h ρ)
  | _, [], _, _ => measurable_const
  | _, _ :: _, _, [] => measurable_const
  | _, l :: ls, h, r :: ρ => by
    have h₁ := (measurable_logwLines ls h.2 ρ).comp (measurable_stepLine M l h.1)
    have h₂ : Measurable fun v : Val _ => M.ℓ l.op (v.pick l.args h.1) :=
      (M.ℓ_meas l.op).comp (Val.measurable_pick h.1)
    simp only [logwLines]
    split_ifs
    · exact h₂.add h₁
    · exact measurable_const.add h₁

namespace Prog
variable (P : Prog Sg) (h : P.WF)

/-- The trace is measurable jointly in the weights and the inputs. -/
theorem measurable_tr : Measurable fun q : Val P.wts × Val P.ins => P.tr M h q.1 q.2 :=
  (measurable_trace M P.lines h.1).comp Val.measurable_append

/-- `(θ, x) ↦ ⟦P⟧_θ(x)` is measurable. -/
theorem measurable_sem : Measurable fun q : Val P.wts × Val P.ins => P.sem M h q.1 q.2 :=
  (Val.measurable_pick h.2).comp (P.measurable_tr M h)

/-- `(θ, x) ↦ L_θ(x)` is measurable. -/
theorem measurable_logw (ρ : List Role) :
    Measurable fun q : Val P.wts × Val P.ins => P.logw M h ρ q.1 q.2 :=
  (measurable_logwLines M P.lines h.1 ρ).comp Val.measurable_append

end Prog

/-! ## Traces extend values -/

theorem get_trace {Γ : List S} (ls : List (Line Sg)) (h : LinesWF Γ ls) (v : Val Γ) {i : ℕ}
    {a : S} (hi : (ctxAfter Γ ls)[i]? = some a) (hlt : i < Γ.length) :
    (trace M ls h v).get i hi = v.get i (by rwa [getElem?_ctxAfter ls hlt] at hi) := by
  induction ls generalizing Γ with
  | nil => rfl
  | cons l ls ih =>
    change (trace M ls h.2 (stepLine M l h.1 v)).get i hi = _
    rw [ih _ _ _ (by simp; omega), stepLine, Val.get_append_left _ _ hlt]

/-- A trace keeps the values it starts from. -/
theorem pick_trace {Γ : List S} (ls : List (Line Sg)) (h : LinesWF Γ ls) (v : Val Γ)
    {xs : List ℕ} {as : List S} (hx : Args Γ xs as) (hx' : Args (ctxAfter Γ ls) xs as) :
    (trace M ls h v).pick xs hx' = v.pick xs hx :=
  Val.pick_eq_of_get _ _ _ _ fun _ hx'' _ h₁ _ => get_trace M ls h v h₁ (hx.lt hx'')

/-- The results of a line, in a trace from before it. -/
theorem pick_trace_results {Γ : List S} (l : Line Sg) (ls : List (Line Sg))
    (h : LinesWF Γ (l :: ls)) (v : Val Γ)
    (hy : Args (ctxAfter Γ (l :: ls)) (List.range' Γ.length l.bs.length) l.bs) :
    (trace M (l :: ls) h v).pick _ hy = M.sem l.op (v.pick l.args h.1) :=
  (pick_trace M ls h.2 (stepLine M l h.1 v) (Args.range' Γ l.bs) hy).trans
    (Val.pick_range'_append _ _)

end RoleFlow