-- VCVio/OracleComp/DistSemantics/EvalDist.lean

import VCVio.OracleComp.Traversal

import VCVio.OracleComp.SimSemantics.SimulateQ

import Mathlib.Probability.Distributions.Uniform

import ToMathlib.General

open OracleSpec Option ENNReal BigOperators

namespace OracleComp

variable {ι : Type u} {spec : OracleSpec ι} {ι' : Type v} {spec' : OracleSpec ι'}
  {α β γ : Type w} [spec.FiniteRange] [spec'.FiniteRange]

section evalDist

end evalDist

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

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

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

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

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

section bounds

variable {oa : OracleComp spec α} {x : α} {p : α → Prop}

end bounds

section support

variable (oa : OracleComp spec α) (x : α) (p q : α → Prop)

variable {oa x p q}

end support

section sums

variable (oa : OracleComp spec α) (p : α → Prop)

end sums

section pure

variable (x : α)

end pure

section bind

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

-- Repository definitions used:
content: notation "[⊥" "|" oa "]" => probFailure oa

content: notation "[=" x "|" oa "]" => probOutput oa 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: def OracleSpec (ι : Type u) : Type (max u (v + 1)) :=
  (i : ι) → Type v × Type v

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 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 -/

-- Repository lemmas used:
content: lemma tsum_option {α β : Type*} [AddCommMonoid α] [TopologicalSpace α]
    [ContinuousAdd α] [T2Space α]
    (f : Option β → α) (hf : Summable (Function.update f none 0)) :
    ∑' x : Option β, f x = f none + ∑' x : β, f (some x)

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

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

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

content: @[simp]
lemma probFailure_add_tsum_probOutput (oa : OracleComp spec α) :
    [⊥ | oa] + ∑' x, [= x | oa] = 1

content: @[simp]
lemma tsum_probOutput_add_probFailure (oa : OracleComp spec α) :
    ∑' x, [= x | oa] + [⊥ | oa] = 1

content: @[simp] lemma probFailure_le_one : [⊥ | oa] ≤ 1

content: @[simp] lemma tsum_probOutput_le_one : ∑' x : α, [= x | oa] ≤ 1

content: @[simp] lemma probOutput_ne_top : [= x | oa] ≠ ∞

content: @[simp] lemma probFailure_ne_top : [⊥ | oa] ≠ ∞

content: lemma tsum_probOutput_eq_sub (oa : OracleComp spec α) :
    ∑' x : α, [= x | oa] = 1 - [⊥ | oa]

content: lemma probFailure_eq_sub_tsum (oa : OracleComp spec α) :
    [⊥ | oa] = 1 - ∑' x : α, [= x | oa]

content: lemma probFailure_bind_eq_tsum :
    [⊥ | oa >>= ob] = [⊥ | oa] + ∑' x : α, [= x | oa] * [⊥ | ob x]

-- Target theorem and proof:
lemma probFailure_bind_eq_sub_mul {oa : OracleComp spec α} {ob : α → OracleComp spec β}
    (r : ℝ≥0∞) (h : ∀ x, [⊥ | ob x] = r) :
    [⊥ | oa >>= ob] = 1 - (1 - [⊥ | oa]) * (1 - r)
:= by
  rw [probFailure_bind_eq_tsum]
  rw [← tsum_probOutput_eq_sub]
  rw [← ENNReal.tsum_mul_right]
  have hl : ∀ x, [=x|oa] * [⊥|ob x] ≤ [=x|oa] :=
    λ x ↦ le_of_le_of_eq (mul_le_mul' le_rfl probFailure_le_one) (mul_one _)
  calc [⊥ | oa] + ∑' x, [= x | oa] * [⊥ | ob x]
    _ = 1 - (∑' x, [= x | oa]) + (∑' x, [= x | oa] * [⊥ | ob x]) := by
      rw [probFailure_eq_sub_tsum]
    _ = 1 - (∑' x, [= x | oa] - ∑' x, [= x | oa] * [⊥ | ob x]) := by
      exact Eq.symm (AddLECancellable.tsub_tsub_assoc
        (by simp) tsum_probOutput_le_one (ENNReal.tsum_le_tsum hl))
    _ = 1 - ∑' x, ([= x | oa] - [= x | oa] * [⊥ | ob x]) := by
      refine congr_arg (1 - ·) (ENNReal.eq_sub_of_add_eq ?_ ?_).symm
      · refine ne_top_of_le_ne_top one_ne_top ?_
        refine le_trans ?_ (@tsum_probOutput_le_one _ _ _ _ oa)
        refine ENNReal.tsum_le_tsum λ x ↦ ?_
        exact hl x
      rw [← ENNReal.tsum_add]
      refine tsum_congr λ x ↦ tsub_add_cancel_of_le (hl x)
    _ = 1 - ∑' x : α, [= x | oa] * (1 - r) := by
      refine congr_arg (1 - ·) (tsum_congr λ x ↦ ?_)
      rw [ENNReal.mul_sub (λ _ _ ↦ probOutput_ne_top), mul_one, ← h x]