import RoleFlow.Front
import RoleFlow.Inverse

/-! # The swap preserves the meaning (module `FrontSem`)

In a model where the swappable operations commute, value and log-weight alike (`Swap.Lawful`;
for `add`, whose log-weight is `0`), a program and any program with some swappable lines swapped
(`SwapLine`) have the same trace values, semantics and log-weight. The theorems about a typed
swap therefore hold for the program as written: density, concordance with the swap's inverse
program, and bijectivity. `swapArgs` emits a swap (`swapArgs_swap`). -/

namespace RoleFlow

open MeasureTheory

variable {S : Type} {Sg : Sig S} [Swap Sg]

/-- `l'` is `l`, or `l = o(a, b)` and `l' = o(b, a)` for a swappable `o`. -/
inductive SwapLine : Line Sg → Line Sg → Prop
  | refl (l : Line Sg) : SwapLine l l
  | swap {as bs : List S} (o : Sg.Op as bs) (hs : Swap.swappable o) (a b : ℕ) :
      SwapLine ⟨o, [a, b]⟩ ⟨o, [b, a]⟩

theorem SwapLine.symm {l l' : Line Sg} : SwapLine l l' → SwapLine l' l
  | .refl l => .refl l
  | .swap o hs a b => .swap o hs b a

theorem swapLine_swapIf (s : Bool) (l : Line Sg) : SwapLine l (swapIf s l) := by
  unfold swapIf; split
  · rename_i h
    rcases l with ⟨op, args⟩
    simp only [Bool.and_eq_true, beq_iff_eq] at h
    obtain ⟨⟨-, hop⟩, hlen⟩ := h
    match args, hlen with
    | [a, b], _ => exact .swap op hop a b
  · exact .refl _

theorem swapLines_swap (sw : ℕ → Bool) :
    ∀ (i : ℕ) (ls : List (Line Sg)), List.Forall₂ SwapLine ls (swapLines sw i ls)
  | _, [] => .nil
  | i, l :: ls => .cons (swapLine_swapIf _ l) (swapLines_swap sw (i + 1) ls)

/-- **The front end emits a swap.** -/
theorem swapArgs_swap (P : Prog Sg) : List.Forall₂ SwapLine P.lines (swapArgs P).lines :=
  swapLines_swap _ 0 P.lines

section Lines
variable {Γ : List S} {ls ls' : List (Line Sg)}

theorem ctxAfter_swap (hs : List.Forall₂ SwapLine ls ls') (Γ : List S) :
    ctxAfter Γ ls' = ctxAfter Γ ls := by
  induction hs generalizing Γ with
  | nil => rfl
  | cons hl _ ih => cases hl <;> exact ih _

theorem SwapLine.args {l l' : Line Sg} (hl : SwapLine l l') (h : Args Γ l.args l.as) :
    Args Γ l'.args l'.as := by
  cases hl with
  | refl => exact h
  | swap o hs a b =>
    obtain ⟨c, rfl⟩ := Swap.sorts o hs
    cases h with | cons ha h => cases h with | cons hb h => exact .cons hb (.cons ha h)

theorem linesWF_swap (hs : List.Forall₂ SwapLine ls ls') (h : LinesWF Γ ls) : LinesWF Γ ls' := by
  induction hs generalizing Γ with
  | nil => trivial
  | cons hl _ ih => cases hl with
    | refl => exact ⟨h.1, ih h.2⟩
    | swap o hso a b => exact ⟨(SwapLine.swap o hso a b).args h.1, ih h.2⟩

end Lines

/-- A model is lawful for the swaps if every swappable operation commutes, in its value and its
log-weight. -/
def Swap.Lawful [Sorts S] (M : Model Sg) : Prop :=
  ∀ {as bs : List S} (o : Sg.Op as bs), Swap.swappable o → ∀ {Γ : List S} (a b : ℕ)
    (h : Args Γ [a, b] as) (h' : Args Γ [b, a] as) (v : Val Γ),
    M.sem o (v.pick [b, a] h') = M.sem o (v.pick [a, b] h) ∧
      M.ℓ o (v.pick [b, a] h') = M.ℓ o (v.pick [a, b] h)

section Lines
variable [Sorts S] {M : Model Sg} (hM : Swap.Lawful M) {Γ : List S} {ls ls' : List (Line Sg)}
include hM

/-- `o(b, a)` steps to the same value as `o(a, b)`. -/
theorem stepLine_swap {as bs : List S} (o : Sg.Op as bs) (hs : Swap.swappable o) (a b : ℕ)
    (h : Args Γ [a, b] as) (h' : Args Γ [b, a] as) (v : Val Γ) :
    stepLine M ⟨o, [b, a]⟩ h' v = stepLine M ⟨o, [a, b]⟩ h v := by
  simp only [stepLine, (hM o hs a b h h' v).1]

theorem trace_swap (hs : List.Forall₂ SwapLine ls ls') (h : LinesWF Γ ls) (h' : LinesWF Γ ls')
    (v : Val Γ) {i : ℕ} {s : S} (hi : (ctxAfter Γ ls)[i]? = some s)
    (hi' : (ctxAfter Γ ls')[i]? = some s) :
    (trace M ls' h' v).get i hi' = (trace M ls h v).get i hi := by
  induction hs generalizing Γ with
  | nil => rfl
  | cons hl _ ih =>
    cases hl with
    | refl => exact ih h.2 h'.2 _ hi hi'
    | swap o hso a b =>
      show (trace M _ h'.2 (stepLine M _ h'.1 v)).get i hi' =
        (trace M _ h.2 (stepLine M _ h.1 v)).get i hi
      rw [stepLine_swap hM o hso a b h.1 h'.1]
      exact ih h.2 h'.2 _ hi hi'

theorem logwLines_swap (hs : List.Forall₂ SwapLine ls ls') (h : LinesWF Γ ls)
    (h' : LinesWF Γ ls') (ρ : List Role) (v : Val Γ) :
    logwLines M ls' h' ρ v = logwLines M ls h ρ v := by
  induction hs generalizing Γ ρ with
  | nil => cases ρ <;> rfl
  | cons hl _ ih =>
    cases ρ with
    | nil => rfl
    | cons r ρ =>
      cases hl with
      | refl => simp only [logwLines]; rw [ih h.2 h'.2]
      | swap o hso a b =>
        simp only [logwLines]
        rw [stepLine_swap hM o hso a b h.1 h'.1, ih h.2 h'.2, (hM o hso a b h.1 h'.1 v).2]

end Lines

/-! ## Programs -/

namespace Prog

variable {P : Prog Sg} {ls : List (Line Sg)} {ρ : List Role}

/-- A swap of `P`. -/
abbrev swap (P : Prog Sg) (ls : List (Line Sg)) : Prog Sg := { P with lines := ls }

theorem wf_swap (hs : List.Forall₂ SwapLine P.lines ls) (h : P.WF) : (P.swap ls).WF :=
  ⟨linesWF_swap hs h.1, by
    show Args (ctxAfter P.ctx₀ ls) P.ret P.outs
    rw [ctxAfter_swap hs]; exact h.2⟩

/-- The written program is well formed if a swap is. -/
theorem wf_of_swap (hs : List.Forall₂ SwapLine P.lines ls) (h : (P.swap ls).WF) : P.WF :=
  wf_swap (P := P.swap ls) (hs.flip.imp fun _ _ h => h.symm) h

variable [Sorts S] {M : Model Sg} (hM : Swap.Lawful M)
include hM

theorem sem_swap (hs : List.Forall₂ SwapLine P.lines ls) (h : P.WF) (h' : (P.swap ls).WF)
    (θ : Val P.wts) (x : Val P.ins) : (P.swap ls).sem M h' θ x = P.sem M h θ x :=
  Val.pick_eq_of_get _ _ _ _ fun _ _ _ hx hx' => trace_swap hM hs h.1 h'.1 _ hx' hx

theorem logw_swap (hs : List.Forall₂ SwapLine P.lines ls) (h : P.WF) (h' : (P.swap ls).WF)
    (θ : Val P.wts) (x : Val P.ins) : (P.swap ls).logw M h' ρ θ x = P.logw M h ρ θ x :=
  logwLines_swap hM hs h.1 h'.1 ρ _

variable (hs : List.Forall₂ SwapLine P.lines ls) (hT : (P.swap ls).Typed ρ)
include hs hT

/-- **Density** of the written program, from a typed swap. -/
theorem density_swap (θ : Val P.wts) :
    Cov (vol P.ins) (vol P.outs) (P.sem M (wf_of_swap hs hT.wf) θ)
      fun x => expw (P.logw M (wf_of_swap hs hT.wf) ρ θ x) := by
  have e₁ := funext (sem_swap hM hs (wf_of_swap hs hT.wf) hT.wf θ)
  have e₂ : (fun x => expw ((P.swap ls).logw M hT.wf ρ θ x)) =
      fun x => expw (P.logw M (wf_of_swap hs hT.wf) ρ θ x) :=
    funext fun x => by rw [logw_swap hM hs (wf_of_swap hs hT.wf) hT.wf]
  rw [← e₁, ← e₂]; exact Prog.density M _ hT θ

/-- **Concordance**: the swap's inverse program undoes the written program. -/
theorem inverse_sem_swap (θ : Val P.wts) (x : Val P.ins) :
    ((P.swap ls).inverse ρ).sem M (inverse_wf hT) θ (P.sem M (wf_of_swap hs hT.wf) θ x) = x := by
  rw [← sem_swap hM hs (wf_of_swap hs hT.wf) hT.wf]; exact inverse_sem M hT θ x

/-- **Bijectivity** of the written program, from a typed swap, under the right law. -/
theorem bijective_swap (hR : M.RightLaw) (θ : Val P.wts) :
    Function.Bijective (P.sem M (wf_of_swap hs hT.wf) θ) := by
  rw [← funext (sem_swap hM hs (wf_of_swap hs hT.wf) hT.wf θ)]; exact bijective hR hT θ

end Prog

end RoleFlow