-- 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 β)

end bind

section mul_le_probEvent_bind

end mul_le_probEvent_bind

section bind_const

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

end bind_const

section query

variable (i : ι) (t : spec.domain i)

end query

section failure

end failure

section map

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

end map

section neverFails

end neverFails

section unit

end unit

section bool

end bool

section eqRec

variable (oa : OracleComp spec α) (h : α = β)

end eqRec

section ite

variable (p : Prop) [Decidable p] (oa oa' : OracleComp spec α)

end ite

section coin

end coin

section uniformFin

variable (n : ℕ)

-- 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: @[reducible, inline] def uniformFin (n : ℕ) : ProbComp (Fin (n + 1)) :=
  unifSpec.query n ()

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

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: @[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: 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: protected class DecidableEq (spec : OracleSpec ι) where
  domain_decidableEq' (i : ι) : DecidableEq (spec.domain i)
  range_decidableEq' (i : ι) : DecidableEq (spec.range i)

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: 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: def finSupport [∀ i, Fintype (spec.range i)] [DecidableEq α] (oa : OracleComp spec α) : Finset α :=
  oa.finSupportWhen fun | query _ _ => Finset.univ

content: def finSupportWhen [DecidableEq α] (oa : OracleComp spec α)
    (possible_outputs : {α : Type v} → OracleQuery spec α → Finset α) : Finset α :=

content: @[reducible, inline] def coin : OracleComp coinSpec Bool :=
  coinSpec.query () ()

content: @[inline, reducible] def coinSpec : OracleSpec.{0,0} Unit := Unit →ₒ Bool

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

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

content: @[simp] lemma finSupport_uniformFin (n : ℕ) : ($[0..n]).finSupport = Finset.univ

content: lemma finSupport_query [spec.FiniteRange] [DecidableEq α] (q : OracleQuery spec α) :
    (q : OracleComp spec _).finSupport = match q with

content: @[simp] lemma finSupportWhen_query [DecidableEq α] (q : OracleQuery spec α) :
    (q : OracleComp spec α).finSupportWhen fin_poss = fin_poss q

content: lemma finSupport_def [∀ i, Fintype (spec.range i)] [DecidableEq α] (oa : OracleComp spec α) :
    oa.finSupport = oa.finSupportWhen fun

-- 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: lemma probEvent_def (oa : OracleComp spec α) (p : α → Prop) :
    [p | oa] = (evalDist oa).run.toOuterMeasure (Option.some '' {x | p x})

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_eq_sum_filter_finSupport [spec.DecidableEq] [DecidablePred p] [DecidableEq α] :
    [p | oa] = ∑ x ∈ oa.finSupport.filter p, [= x | oa]

content: @[simp]
lemma probOutput_liftM [Fintype α] (q : OracleQuery spec α) (u : α) :
    [= u | (q : OracleComp spec α)] = (Fintype.card α : ℝ≥0∞)⁻¹

content: lemma probOutput_query (u : spec.range i) :
    [= u | (query i t : OracleComp spec _)] = (Fintype.card (spec.range i) : ℝ≥0∞)⁻¹

content: @[simp]
lemma probOutput_uniformFin (x : Fin (n + 1)) : [= x | $[0..n]] = ((n : ℝ≥0∞) + 1)⁻¹

-- Target theorem and proof:
@[simp]
lemma probEvent_uniformFin (p : Fin (n + 1) → Prop) [DecidablePred p] :
    [p | $[0..n]] = (Finset.univ.filter p).card * (n + 1 : ℝ≥0∞)⁻¹
:= by
  simp only [probEvent_eq_sum_filter_finSupport, finSupport_uniformFin, probOutput_uniformFin,
    Finset.sum_const, nsmul_eq_mul]