-- VCVio/OracleComp/DistSemantics/Seq.lean

import VCVio.OracleComp.DistSemantics.Monad

namespace OracleComp

variable {ι : Type u} {spec : OracleSpec ι} {α β γ : Type v}

variable (oa : OracleComp spec α) (og : OracleComp spec (α → β))

section seq_map

variable (oa : OracleComp spec α) (ob : OracleComp spec β) (f : α → β → γ)

section swap

end swap

section injective2

end injective2

-- Repository definitions used:
content: notation "[=" x "|" oa "]" => probOutput oa x

content: notation "$[0.." n "]" => uniformFin n

content: notation:50 "$[" n "⋯" m "]" => uniformFin' n m

example : OracleComp probSpec ℕ := do
  let x ← $[314⋯31415]; let y ← $[0⋯x]
  return x + 2 * y

content: notation "[⊥" "|" oa "]" => probFailure oa

content: def OracleSpec (ι : Type u) : Type (max u (v + 1)) :=
  (i : ι) → Type v × Type v

content: noncomputable def probOutput (oa : OracleComp spec α) (x : α) : ℝ≥0∞ :=
  (evalDist oa).run (some x)

content: class HasEvalDist (m : Type u → Type v) [Monad m] where
  evalDist {α : Type u} (mx : m α) : SPMF α
  evalDist_pure {α : Type u} (x : α) : evalDist (pure x : m α) = pure x
  evalDist_bind {α β : Type u} (mx : m α) (my : α → m β) :
    evalDist (mx >>= my) = evalDist mx >>= fun x => evalDist (my x)

content: def OracleComp {ι : Type u} (spec : OracleSpec.{u,v} ι) :
    Type w → Type (max u (v + 1) w) :=
  OptionT (FreeMonad (OracleQuery.{u,v} spec))

content: inductive OracleQuery {ι : Type u} (spec : OracleSpec.{u,v} ι) : Type v → Type (max u v)
  | query (i : ι) (t : spec.domain i) : OracleQuery spec (spec.range i)

content: @[inline, reducible]
protected def domain (spec : OracleSpec ι) (i : ι) : Type v := (spec i).1

content: @[inline, reducible]
protected def range (spec : OracleSpec ι) (i : ι) : Type w := (spec i).2

content: inductive FreeMonad (f : Type u → Type v) (α : Type w) : Type (max (u + 1) v w)
  | protected pure (x : α) : FreeMonad f α
  | roll {β : Type u} (x : f β) (r : β → FreeMonad f α) : FreeMonad f α

content: def support (oa : OracleComp spec α) : Set α :=
  oa.supportWhen fun _ => Set.univ

content: def supportWhen (oa : OracleComp spec α)
    (possible_outputs : {α : Type v} → OracleQuery spec α → Set α) : Set α :=

content: @[elab_as_elim]
protected def induction {C : OracleComp spec α → Prop}
    (oa : OracleComp spec α) (pure : (a : α) → C (pure a))
    (query_bind : (i : ι) → (t : spec.domain i) →
      (oa : spec.range i → OracleComp spec α) → (∀ u, C (oa u)) → C (query i t >>= oa))
    (failure : C failure) : C oa :=
  FreeMonad.inductionOn (Option.rec failure pure) (λ (query i t) ↦ query_bind i t) oa

content: @[elab_as_elim]
protected def inductionOn {C : FreeMonad f α → Prop}
    (pure : ∀ x, C (pure x))
    (roll : ∀ {β} (x : f β), (r : β → FreeMonad f α) →
      (∀ y, C (r y)) → C (x >>= r)) :
    (oa : FreeMonad f α) → C oa
  | FreeMonad.pure x => pure x
  | FreeMonad.roll x r => roll x _ (λ u ↦
      FreeMonad.inductionOn pure roll (r u))

content: @[elab_as_elim]
protected def construct {C : OracleComp spec α → Type*}
    (pure : (a : α) → C (pure a))
    (query_bind : {β : Type v} → (q : OracleQuery spec β) →
      (oa : β → OracleComp spec α) → ((u : β) → C (oa u)) → C (q >>= oa))
    (failure : C failure) (oa : OracleComp spec α) : C oa :=
  FreeMonad.construct (Option.rec failure pure) query_bind oa

content: @[elab_as_elim]
protected def construct {C : FreeMonad f α → Type*}
    (pure : (x : α) → C (pure x))
    (roll : {β : Type u} → (x : f β) → (r : β → FreeMonad f α) →
      ((y : β) → C (r y)) → C (x >>= r)) :
    (oa : FreeMonad f α) → C oa
  | .pure x => pure x
  | .roll x r => roll x _ (λ u ↦ FreeMonad.construct pure roll (r u))

content: class FiniteRange (spec : OracleSpec ι) where
    range_inhabited' (i : ι) : Inhabited (spec.range i)
    range_fintype' (i : ι) : Fintype (spec.range i)

content: @[elab_as_elim]
protected def inductionOn {C : OracleComp spec α → Prop}
    (pure : (a : α) → C (pure a))
    (query_bind : (i : ι) → (t : spec.domain i) →
      (oa : spec.range i → OracleComp spec α) → (∀ u, C (oa u)) → C (query i t >>= oa))
    (failure : C failure) (oa : OracleComp spec α) : C oa :=
  FreeMonad.inductionOn (Option.rec failure pure) (λ (query i t) ↦ query_bind i t) oa

content: noncomputable def probEvent (oa : OracleComp spec α) (p : α → Prop) : ℝ≥0∞ :=
  (evalDist oa).run.toOuterMeasure (Option.some '' {x | p x})

content: def cases {α} {motive : {n : ℕ} → Vector α n → Sort*} (v_empty : motive #v[])
  (v_insert : {n : ℕ} → (hd : α) → (tl : Vector α n) → motive (tl.insertIdx 0 hd)) {m : ℕ} :
    (v : Vector α m) → motive v := match hm : m with
  | 0 => fun v => match v with | ⟨⟨[]⟩, rfl⟩ => v_empty
  | n + 1 => fun v => match hv : v with
    | ⟨⟨hd :: tl⟩, hSize⟩ => by admit /- proof elided -/

content: noncomputable def probFailure (oa : OracleComp spec α) : ℝ≥0∞ :=
  (evalDist oa).run none

content: @[reducible, inline] def uniformFin (n : ℕ) : ProbComp (Fin (n + 1)) :=
  unifSpec.query n ()

content: abbrev ProbComp : Type z → Type (max z 1) := OracleComp unifSpec

content: @[inline, reducible] def unifSpec : OracleSpec.{0,0} ℕ :=
  λ n ↦ (Unit, Fin (n + 1))

content: @[reducible, inline] def uniformFin' (n m : ℕ) : OracleComp probSpec (Fin (m + 1)) :=
  probSpec.query m n

content: def probSpec : OracleSpec.{0,0} ℕ :=
  fun n => (ℕ, Fin (n + 1))

content: notation "Pr[=" x "|" mx "]" => probOutput mx x

content: notation "Pr[" p "|" mx "]" => probEvent mx p

content: notation "Pr[⊥" "|" mx "]" => probFailure mx

content: notation "[=" x "|" oa "]" => probOutput oa x

content: notation "[⊥" "|" oa "]" => probFailure oa

content: notation "[" p "|" oa "]" => probEvent oa p

content: notation "$[0.." n "]" => uniformFin n

content: notation:50 "$[" n "⋯" m "]" => uniformFin' n m

example : OracleComp probSpec ℕ := do
  let x ← $[314⋯31415]; let y ← $[0⋯x]
  return x + 2 * y

-- Repository lemmas used:
content: lemma probOutput_bind_eq_tsum (y : β) :
    [= y | oa >>= ob] = ∑' x : α, [= x | oa] * [= y | ob x]

content: @[simp]
lemma probEvent_pure (p : α → Prop) [DecidablePred p] :
    [p | (pure x : OracleComp spec α)] = if p x then 1 else 0

content: @[simp]
lemma probOutput_pure_self (x : α) : [= x | (pure x : OracleComp spec α)] = 1

content: @[simp] lemma support_map (oa : OracleComp spec α) (f : α → β) :
    (f <$> oa).support = f '' oa.support

content: @[simp] lemma support_bind (oa : OracleComp spec α) (ob : α → OracleComp spec β) :
    (oa >>= ob).support = ⋃ x ∈ oa.support, (ob x).support

content: @[simp] lemma support_pure (x : α) :
    (pure x : OracleComp spec α).support = {x}

content: lemma probEvent_congr' {p q : α → Prop} {oa : OracleComp spec α} {oa' : OracleComp spec' α}
    (h1 : ∀ x, x ∈ oa.support → x ∈ oa'.support → (p x ↔ q x))
    (h2 : evalDist oa = evalDist oa') : [p | oa] = [q | oa']

content: lemma mem_support_iff_of_evalDist_eq {oa : OracleComp spec α} {oa' : OracleComp spec' α}
    (h : evalDist oa = evalDist oa') (x : α) : x ∈ oa.support ↔ x ∈ oa'.support

content: lemma mem_support_iff_probOutput_ne_zero : x ∈ oa.support ↔ [= x | oa] ≠ 0

content: @[simp low]
lemma probOutput_eq_zero_iff : [= x | oa] = 0 ↔ x ∉ oa.support

content: @[simp]
lemma mem_support_evalDist_iff (oa : OracleComp spec α) (x : α) :
    some x ∈ (evalDist oa).run.support ↔ x ∈ oa.support

content: lemma probOutput_def (oa : OracleComp spec α) (x : α) :
    [= x | oa] = (evalDist oa).run (some x)

content: lemma probEvent_eq_tsum_indicator (oa : OracleComp spec α) (p : α → Prop) :
    [p | oa] = ∑' x : α, {x | p x}.indicator ([= · | oa]) x

content: lemma probEvent_def (oa : OracleComp spec α) (p : α → Prop) :
    [p | oa] = (evalDist oa).run.toOuterMeasure (Option.some '' {x | p x})

content: lemma probOutput_congr {x y : α} {oa : OracleComp spec α} {oa' : OracleComp spec' α}
    (h1 : x = y) (h2 : evalDist oa = evalDist oa') : [= x | oa] = [= y | oa']

content: lemma probEvent_comp (q : β → Prop) : [q ∘ f | oa] = [q | f <$> oa]

content: @[simp]
lemma probEvent_map (q : β → Prop) : [q | f <$> oa] = [q ∘ f | oa]

content: @[simp]
lemma evalDist_map (oa : OracleComp spec α) (f : α → β) :
    evalDist (f <$> oa) = f <$> (evalDist oa)

content: lemma probEvent_eq_tsum_ite (oa : OracleComp spec α) (p : α → Prop) [DecidablePred p] :
    [p | oa] = ∑' x : α, if p x then [= x | oa] else 0

content: lemma probEvent_ext (h : ∀ x ∈ oa.support, p x ↔ q x) : [p | oa] = [q | oa]

content: lemma probEvent_mono (h : ∀ x ∈ oa.support, p x → q x) : [p | oa] ≤ [q | oa]

content: @[simp]
lemma evalDist_apply_eq_zero_iff (x : Option α) :
    (evalDist oa).run x = 0 ↔ x.rec ([⊥ | oa] = 0) (· ∉ oa.support)

content: lemma probFailure_def (oa : OracleComp spec α) :
    [⊥ | oa] = (evalDist oa).run none

content: lemma Prod.mk.injective2 {α β : Type*} :
    Function.Injective2 (Prod.mk : α → β → α × β)

content: lemma probEvent_bind_eq_tsum (q : β → Prop) :
    [q | oa >>= ob] = ∑' x : α, [= x | oa] * [q | ob x]

-- Local definitions used:

-- Local lemmas used:
content: @[simp low]
lemma support_seq : (og <*> oa).support = ⋃ g ∈ og.support, g '' oa.support

content: @[simp low + 1]
lemma support_seq_map_eq_image2 :
    (f <$> oa <*> ob).support = Set.image2 f oa.support ob.support

content: lemma probOutput_seq_map_eq_tsum [spec.FiniteRange]
    (z : γ) : [= z | f <$> oa <*> ob] = ∑' (x : α) (y : β),
      [= x | oa] * [= y | ob] * [= z | (pure (f x y) : OracleComp spec γ)]

content: lemma probEvent_seq_map_eq_probEvent_comp_uncurry [spec.FiniteRange]
    (p : γ → Prop) : [p | f <$> oa <*> ob] =
      [p ∘ f.uncurry | Prod.mk <$> oa <*> ob]

content: lemma probEvent_seq_map_eq_probEvent [spec.FiniteRange] (p : γ → Prop) :
    [p | f <$> oa <*> ob] = [λ z ↦ p (f z.1 z.2) | Prod.mk <$> oa <*> ob]

content: lemma probOutput_seq_map_eq_mul_of_injective2 [spec.FiniteRange]
    (hf : f.Injective2) (x : α) (y : β) :
    [= f x y | f <$> oa <*> ob] = [= x | oa] * [= y | ob]

-- Target theorem and proof:
lemma probEvent_seq_map_eq_mul {ι : Type u} {spec : OracleSpec ι}
    {α β γ : Type v} (f : α → β → γ) [spec.FiniteRange]
    (oa : OracleComp spec α) (ob : OracleComp spec β)
    (p : γ → Prop) (q1 : α → Prop) (q2 : β → Prop)
    (h : ∀ x ∈ oa.support, ∀ y ∈ ob.support, p (f x y) ↔ q1 x ∧ q2 y) :
    [p | f <$> oa <*> ob] = [q1 | oa] * [q2 | ob]
:= by
  have : DecidablePred q1 := Classical.decPred q1
  have : DecidablePred q2 := Classical.decPred q2
  rw [probEvent_seq_map_eq_probEvent]
  calc [λ z : α × β ↦ p (f z.1 z.2) | Prod.mk <$> oa <*> ob] =
      [λ z ↦ q1 z.1 ∧ q2 z.2 | Prod.mk <$> oa <*> ob] :=
        probEvent_ext <| by simpa using λ x y hx hy ↦ h x hx y hy
    _ = ∑' (x : α) (y : β), if q1 x ∧ q2 y then [= (x, y) | Prod.mk <$> oa <*> ob] else 0 := by
      rw [probEvent_eq_tsum_ite, ENNReal.tsum_prod']
    _ = ∑' (x : α) (y : β), if q1 x ∧ q2 y then [= x | oa] * [= y | ob] else 0 := by
      simp_rw [probOutput_seq_map_eq_mul_of_injective2 oa ob Prod.mk Prod.mk.injective2]
    _ = ∑' x : α, if q1 x then [= x | oa] * (∑' y : β, if q2 y then [= y | ob] else 0) else 0 :=
      tsum_congr (λ x ↦ by by_cases hx : q1 x <;> simp [hx, ← ENNReal.tsum_mul_left])
    _ = ∑' x : α, if q1 x then [= x | oa] * [q2 | ob] else 0 := by rw [probEvent_eq_tsum_ite]
    _ = [q1 | oa] * [q2 | ob] := by
      simp only [probEvent_eq_tsum_ite oa, ← ENNReal.tsum_mul_right, ite_mul, zero_mul]