import RoleFlow.Density

/-! # A discrete instance (`TODO.md`, discrete sorts)

The generic development needs a σ-finite reference measure per sort and nothing else, so
discrete sorts under counting measure fit as they are: there a density is a probability mass
function. This module is a minimal instance with two sorts, `int` (`ℤ`) and `bool`, both under
counting measure:

| operation | `inv = false` | `inv = true` | log-weight |
|---|---|---|---|
| `add` | `x + q` | `x - q` | `0` |
| `xor` | `x ^^ q` | (itself) | `0` |

and the conditioners `mul` and `parity`, which have no semi-inverse. Every leaf is a bijection
of a countable type, which preserves counting measure (`measurePreserving_count`). The examples
at the end are integer and binary additive couplings (as in integer discrete flows); their
densities are pmfs by `Prog.density`. -/

open MeasureTheory

namespace RoleFlow.Discrete

/-! ## Sorts -/

/-- The sorts: integers and booleans, each under counting measure. -/
inductive DS
  | int
  | bool
  deriving DecidableEq, Repr

/-- The space of a sort. -/
abbrev DS.X : DS → Type
  | .int => ℤ
  | .bool => Bool

instance (s : DS) : MeasurableSpace s.X := ⊤

instance (s : DS) : MeasurableSingletonClass s.X := ⟨fun _ => MeasurableSpace.measurableSet_top⟩

instance (s : DS) : Countable s.X := by cases s <;> (show Countable (DS.X _); infer_instance)

noncomputable instance : Sorts DS where
  X := DS.X
  meas _ := inferInstance
  μ _ := Measure.count
  sigmaFinite _ := inferInstance

instance (s : DS) : MeasurableSingletonClass (X s) := inferInstanceAs (MeasurableSingletonClass s.X)

instance (s : DS) : Countable (X s) := inferInstanceAs (Countable s.X)

open DS

/-- A bijection of a countable type preserves counting measure. -/
theorem measurePreserving_count {α : Type*} [MeasurableSpace α] [MeasurableSingletonClass α]
    (e : α ≃ α) (he : Measurable e) : MeasurePreserving e Measure.count Measure.count := by
  refine ⟨he, Measure.ext fun s hs => ?_⟩
  rw [Measure.map_apply he hs, Measure.count_apply hs, Measure.count_apply (he hs),
    Set.encard_preimage_of_bijective e.bijective]

/-! ## Operations -/

inductive DOp : List DS → List DS → Type
  /-- `x + q`; `x - q` if `inv`. -/
  | add (inv := false) : DOp [int, int] [int]
  /-- `x ^^ q`, its own semi-inverse. -/
  | xor : DOp [bool, bool] [bool]
  | mul : DOp [int, int] [int]
  | parity : DOp [int] [bool]

/-- A semi-inverse solving the first argument from the result and the second argument. -/
abbrev si1 {a b c : DS} (o : DOp [c, b] [a]) : SemiInv DOp [a, b] [c] :=
  ⟨[a], [b], rfl, by simp, o⟩

def DOp.inv : {as bs : List DS} → DOp as bs → Option (SemiInv DOp as bs)
  | _, _, .add inv => some (si1 (.add !inv))
  | _, _, .xor => some (si1 .xor)
  | _, _, _ => none

def discSig : Sig DS where
  Op := DOp
  bs_ne := fun o => by cases o <;> simp
  inv := DOp.inv

/-! ## The model -/

noncomputable section

abbrev one {a : DS} (x : X a) : Val [a] := Val.cons x Val.nil

/-- The integer operations, at the type of the sort. -/
def toZ (x : X int) : ℤ := x
def toB (x : X bool) : Bool := x
def zadd (x q : X int) : X int := toZ x + toZ q
def zsub (x q : X int) : X int := toZ x - toZ q
def zmul (x q : X int) : X int := toZ x * toZ q
def bxor (x q : X bool) : X bool := Bool.xor (toB x) (toB q)
def zpar (x : X int) : X bool := decide (toZ x % 2 = 1)

def semOp : {as bs : List DS} → DOp as bs → Val as → Val bs
  | _, _, .add inv => fun v => one (a := int)
      (bif inv then zsub v.head v.tail.head else zadd v.head v.tail.head)
  | _, _, .xor => fun v => one (a := bool) (bxor v.head v.tail.head)
  | _, _, .mul => fun v => one (a := int) (zmul v.head v.tail.head)
  | _, _, .parity => fun v => one (a := bool) (zpar v.head)

/-- On a discrete sort, a measurable map is any map. -/
theorem measurable_val {Γ : List DS} {Y : Type*} [MeasurableSpace Y] (f : Val Γ → Y) :
    Measurable f :=
  have : Countable (Val Γ) := inferInstanceAs (Countable (∀ i, X (Γ.get i)))
  have : MeasurableSingletonClass (Val Γ) :=
    inferInstanceAs (MeasurableSingletonClass (∀ i, X (Γ.get i)))
  measurable_of_countable f

theorem eq_one {a : DS} (v : Val [a]) : one v.head = v := by
  conv_rhs => rw [← Val.cons_head_tail v, Val.eq_nil v.tail]

/-- The left law of an operation that consumes its first argument. -/
theorem left1 {a b c : DS} {o : DOp [a, b] [c]} {o' : DOp [c, b] [a]} (F : X a → X b → X c)
    (G : X c → X b → X a) (hF : ∀ v, semOp o v = one (F v.head v.tail.head))
    (hG : ∀ v, semOp o' v = one (G v.head v.tail.head)) (h : ∀ x q, G (F x q) q = x) :
    ∀ c p, semOp (si1 o').op ((semOp o ((si1 o').glue c p)).append p) = c :=
  fun (x : Val [a]) (p : Val [b]) => by
    show semOp o' ((semOp o (x.append p)).append p) = x
    rw [hF, hG]; show one (G (F x.head p.head) p.head) = x; rw [h]; exact eq_one x

/-- A leaf of weight `1` from a bijection of the consumed argument at each parameter. -/
theorem leaf1 {a b : DS} {o : DOp [a, b] [a]} (F : X a → X b → X a)
    (hF : ∀ v, semOp o v = one (F v.head v.tail.head)) (e : X b → X a ≃ X a)
    (he : ∀ q x, e q x = F x q) (p : Val [b]) :
    Cov (vol [a]) (vol [a]) (fun c => semOp o ((c.append p).cast rfl)) (fun _ => expw 0) := by
  have h : MeasurePreserving (fun c : Val [a] => one (e p.head c.head)) (vol [a]) (vol [a]) :=
    Val.measurePreserving_single.comp
      ((measurePreserving_count (e p.head) (measurable_of_countable _)).comp
        Val.measurePreserving_head)
  refine (Cov.pres h).congr (fun c => ?_) (fun _ => by simp)
  show _ = semOp o (c.append p); rw [hF, he]; rfl

/-- The left law of `o` with the semi-inverse `si`. -/
abbrev LeftLaw {as bs : List DS} (o : DOp as bs) (si : SemiInv DOp as bs) : Prop :=
  ∀ c p, semOp si.op ((semOp o (si.glue c p)).append p) = c

theorem left_add (inv : Bool) : LeftLaw (.add inv) (si1 (.add !inv)) := by
  cases inv
  exacts [left1 zadd zsub (fun _ => rfl) (fun _ => rfl) (by simp [zadd, zsub, toZ]),
    left1 zsub zadd (fun _ => rfl) (fun _ => rfl) (by simp [zadd, zsub, toZ])]

theorem left_xor : LeftLaw .xor (si1 .xor) :=
  left1 bxor bxor (fun _ => rfl) (fun _ => rfl) (by simp [bxor, toB])

/-- `· ^^ q` as a bijection. -/
def xorEquiv (q : X bool) : X bool ≃ X bool :=
  ⟨(bxor · q), (bxor · q), fun _ => by simp [bxor, toB], fun _ => by simp [bxor, toB]⟩

def discModel : ModelInv discSig where
  sem := semOp
  ℓ := fun _ _ => 0
  sem_meas := fun _ => measurable_val _
  ℓ_meas := fun _ => measurable_const
  left := by
    intro as bs o si h
    cases o <;> cases h
    case add inv => exact left_add inv
    case xor => exact left_xor
  leaf := by
    intro as bs o si h p
    cases o <;> cases h
    case add inv =>
      cases inv
      exacts [leaf1 (o := .add) zadd (fun _ => rfl)
          (fun q => show ℤ ≃ ℤ from Equiv.addRight (toZ q)) (fun _ _ => rfl) p,
        leaf1 (o := .add true) zsub (fun _ => rfl)
          (fun q => show ℤ ≃ ℤ from Equiv.subRight (toZ q)) (fun _ _ => rfl) p]
    case xor => exact leaf1 (o := .xor) bxor (fun _ => rfl) xorEquiv (fun _ _ => rfl) p
  right := by
    intro as bs o si h b p
    cases o <;> cases h
    case add inv => cases inv; exacts [left_add true b p, left_add false b p]
    case xor => exact left_xor b p

end

/-! ## Examples -/

def ln {as bs : List DS} (o : DOp as bs) (args : List ℕ) : Line discSig := ⟨o, args⟩

/-- Integer additive coupling: `θ : int` (0), `x₁ x₂ : int` (1, 2); `t = mul(θ, x₁)` (3),
`y = add(x₂, t)` (4); returns `(x₁, y)`. -/
def intCoupling : Prog discSig where
  wts := [int]
  ins := [int, int]
  outs := [int, int]
  lines := [ln .mul [0, 1], ln .add [2, 3]]
  ret := [1, 4]

/-- Binary coupling across sorts: `x : int` (0), `b : bool` (1); `p = parity(x)` (2),
`c = xor(b, p)` (3); returns `(x, c)`. -/
def boolCoupling : Prog discSig where
  wts := []
  ins := [int, bool]
  outs := [int, bool]
  lines := [ln .parity [0], ln .xor [1, 2]]
  ret := [0, 3]

set_option linter.style.nativeDecide false in
theorem intCoupling_typed : intCoupling.Typed [.cond, .spine] := by native_decide

set_option linter.style.nativeDecide false in
theorem boolCoupling_typed : boolCoupling.Typed [.cond, .spine] := by native_decide

/-- **Density** of the integer coupling under counting measure: a pmf. -/
example (θ : Val intCoupling.wts) :
    Cov (vol intCoupling.ins) (vol intCoupling.outs)
      (intCoupling.sem discModel.toModel intCoupling_typed.wf θ)
      fun x => expw (intCoupling.logw discModel.toModel intCoupling_typed.wf [.cond, .spine] θ x) :=
  intCoupling.density _ intCoupling_typed θ

example : Cov (vol boolCoupling.ins) (vol boolCoupling.outs)
      (boolCoupling.sem discModel.toModel boolCoupling_typed.wf Val.nil)
      fun x => expw (boolCoupling.logw discModel.toModel boolCoupling_typed.wf [.cond, .spine]
        Val.nil x) :=
  boolCoupling.density _ boolCoupling_typed Val.nil

end RoleFlow.Discrete