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

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

-- Repository definitions used:
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 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: @[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: 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)

-- 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_indicator (oa : OracleComp spec α) (p : α → Prop) :
    [p | oa] = ∑' x : α, {x | p x}.indicator ([= · | oa]) x

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

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

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

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 probOutput_congr {x y : α} {oa : OracleComp spec α} {oa' : OracleComp spec' α}
    (h1 : x = y) (h2 : evalDist oa = evalDist oa') : [= x | oa] = [= y | oa']

-- Target theorem and proof:
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']
:= by
  have h : ∀ x, x ∈ oa.support ↔ x ∈ oa'.support := mem_support_iff_of_evalDist_eq h2
  have h' : ∀ x, [= x | oa] = [= x | oa'] := λ x ↦ probOutput_congr rfl h2
  rw [probEvent_eq_tsum_indicator, probEvent_eq_tsum_indicator]
  refine tsum_congr λ x ↦ ?_
  simp [Set.indicator, h']
  by_cases hp : p x
  · by_cases hq : q x
    · simp [hp, hq]
    · simp [hp, hq, h]
      refine λ hoa ↦ hq ?_
      refine (h1 _ ?_ hoa).1 hp
      refine (h _).2 hoa
  · by_cases hq : q x
    · simp [hp, hq]
      simp [h] at h1
      intro hoa
      specialize h1 _ hoa
      tauto
    · rw [if_neg hp, if_neg hq]