-- 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 probFailure_ne_top : [⊥ | oa] ≠ ∞

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

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

-- Target theorem and proof:
lemma probFailure_bind_of_const [Nonempty α] (r : ℝ≥0∞) (h : ∀ x, [⊥ | ob x] = r) :
    [⊥ | oa >>= ob] = [⊥ | oa] + r - [⊥ | oa] * r
:= by
  have : r ≠ ⊤ := λ hr ↦ probFailure_ne_top ((h (Classical.arbitrary α)).trans hr)
  simp [probFailure_bind_eq_tsum, h, ENNReal.tsum_mul_right, tsum_probOutput_eq_sub]
  rw [ENNReal.sub_mul λ _ _ ↦ this, one_mul]
  refine symm (AddLECancellable.add_tsub_assoc_of_le ?_ ?_ _)
  · refine ENNReal.addLECancellable_iff_ne.2 (ENNReal.mul_ne_top probFailure_ne_top this)
  · by_cases hr : r = 0
    · simp only [hr, mul_zero, le_refl]
    refine mul_le_of_le_div (le_of_le_of_eq probFailure_le_one ?_)
    refine symm (ENNReal.div_self hr this)