-- VCVio/EvalDist/Basic.lean

import Mathlib.Probability.ProbabilityMassFunction.Monad

import ToMathlib.General

open ENNReal

variable {α β γ : Type u} {m : Type u → Type v} [Monad m]

@[reducible] def SPMF : Type u → Type u := OptionT PMF

namespace SPMF

end SPMF

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)

def probOutput [HasEvalDist m] (mx : m α) (x : α) : ℝ≥0∞ := evalDist mx x

noncomputable def probEvent [HasEvalDist m] (mx : m α) (p : α → Prop) : ℝ≥0∞ :=
  (evalDist mx).run.toOuterMeasure (some '' {x | p x})

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

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

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

variable [HasEvalDist m]

section sums

end sums

section bounds

variable {mx : m α} {mxe : OptionT m α} {x : α} {p : α → Prop}

end bounds

section bind

variable (mx : m α) (my : α → m β)

end bind

namespace SPMF

instance hasEvalDist : HasEvalDist SPMF where
  evalDist := id
  evalDist_pure _ := rfl
  evalDist_bind _ _ := rfl

variable (p : SPMF α) (x : α)

end SPMF

namespace PMF

noncomputable instance hasEvalDist : HasEvalDist PMF where
  evalDist p := OptionT.mk p
  evalDist_pure _ := by admit /- proof elided -/

variable (p : PMF α) (x : α)

-- Repository definitions used:
content: notation "Pr[" p "|" mx "]" => probEvent mx p

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

content: macro_rules (kind := probEventBinding1)
  | `(Pr[$cond:term | $var:ident ← $src:term]) => `(Pr[fun $var => $cond | $src])

content: macro_rules (kind := probEventBinding2)
  
  | `(Pr{{$items*}}[$t]) => `(probOutput (do $items:doSeqItem* return $t:term) True)
  
  | `(Pr{$items*}[$t]) => `(probOutput (do $items:doSeqItem* return $t:term) True)

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

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

-- Repository lemmas used:
content: @[simp]
lemma PMF.monad_pure_eq_pure {α : Type u} (x : α) :
    (Pure.pure x : PMF α) = PMF.pure x

content: @[simp]
lemma PMF.monad_bind_eq_bind {α β : Type u}
      (p : PMF α) (q : α → PMF β) : p >>= q = p.bind q

-- Local definitions used:
content: @[reducible] def SPMF : Type u → Type u := OptionT PMF

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 probOutput [HasEvalDist m] (mx : m α) (x : α) : ℝ≥0∞ := evalDist mx x

content: noncomputable def probEvent [HasEvalDist m] (mx : m α) (p : α → Prop) : ℝ≥0∞ :=
  (evalDist mx).run.toOuterMeasure (some '' {x | p x})

content: instance hasEvalDist : HasEvalDist SPMF where
  evalDist := id
  evalDist_pure _ := rfl
  evalDist_bind _ _ := rfl

content: noncomputable instance hasEvalDist : HasEvalDist PMF where
  evalDist p := OptionT.mk p
  evalDist_pure _ := by admit /- proof elided -/

-- Local lemmas used:
content: lemma probOutput_def (mx : m α) (x : α) : Pr[= x | mx] = (evalDist mx).run (some x)

content: @[simp] lemma evalDist_eq : evalDist p = liftM p

-- Target theorem and proof:
@[simp] lemma probOutput_eq : probOutput p = p
:= by
  refine funext fun x => ?_
  simp only [probOutput_def, evalDist_eq, monad_pure_eq_pure, monad_bind_eq_bind, OptionT.run_mk,
    pure_apply, Option.some.injEq, mul_ite, mul_one, mul_zero]
  simp
  refine (PMF.map_apply _ _ _).trans ?_
  refine (tsum_eq_single x ?_).trans ?_
  · simp
    refine fun x h h' => ?_
    refine (h h'.symm).elim
  simp only [↓reduceIte]