-- VCVio/OracleComp/Constructions/UniformSelect.lean

import VCVio.OracleComp.DistSemantics.Prod

import Batteries.Control.OptionT

open OracleSpec BigOperators ENNReal

namespace OracleComp

section uniformSelect

prefix : 50 "$" => uniformSelect _

variable {cont β : Type} [h : HasUniformSelect cont β]

end uniformSelect

section uniformSelectList

variable {α : Type}

end uniformSelectList

section uniformSelectVector

end uniformSelectVector

section uniformSelectFinset

variable {α : Type}

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

content: notation "[]ₒ" => emptySpec

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

content: def OracleSpec (ι : Type u) : Type (max u (v + 1)) :=
  (i : ι) → Type v × Type v

content: @[inline, reducible]
protected def domain (spec : OracleSpec ι) (i : ι) : Type v := (spec i).1

content: @[inline, reducible]
protected def range (spec : OracleSpec ι) (i : ι) : Type w := (spec i).2

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

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: 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: noncomputable def probOutput (oa : OracleComp spec α) (x : α) : ℝ≥0∞ :=
  (evalDist oa).run (some x)

content: instance isEmpty : IsEmpty (OracleQuery []ₒ α) where false | query i t => i.elim

content: @[inline, reducible]
def emptySpec : OracleSpec PEmpty := λ _ ↦ (PUnit, PUnit)

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: @[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: @[elab_as_elim]
def induction {α} {motive : {n : ℕ} → Vector α n → Sort*} (v_empty : motive #v[])
  (v_insert : {n : ℕ} → (hd : α) → (tl : Vector α n) → motive tl → motive (tl.insertIdx 0 hd))
    {m : ℕ} : (v : Vector α m) → motive v :=

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

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

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

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

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 "[]ₒ" => emptySpec

-- Repository lemmas used:
content: lemma probOutput_map_eq_sum_fintype_ite [Fintype α] [DecidableEq β] (y : β) :
    [= y | f <$> oa] = ∑ x : α, if y = f x then [= x | oa] else 0

content: lemma probOutput_map_eq_tsum_ite [DecidableEq β] (y : β) :
    [= y | f <$> oa] = ∑' x : α, if y = f x then [= x | oa] else 0

content: lemma probOutput_bind_eq_tsum (y : β) :
    [= y | oa >>= ob] = ∑' x : α, [= x | oa] * [= y | ob x]

content: @[simp]
lemma probOutput_pure [DecidableEq α] (y : α) :
    [= y | (pure x : OracleComp spec α)] = if y = x then 1 else 0

content: lemma List.countP_eq_sum_fin_ite {α : Type*} (xs : List α) (p : α → Bool) :
    (∑ i : Fin xs.length, if p xs[i] then 1 else 0) = xs.countP p

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:

-- Local lemmas used:
content: lemma uniformSelectList_cons (x : α) (xs : List α) :
    ($ x :: xs : ProbComp α) = ((x :: xs)[·]) <$> $[0..xs.length]

content: @[simp] lemma probOutput_uniformSelectList [DecidableEq α] (xs : List α) (x : α) :
    [= x | $ xs] = if xs.isEmpty then 0 else (xs.count x : ℝ≥0∞) / xs.length

content: lemma uniformSelectFinset_def {α : Type} [DecidableEq α] (s : Finset α) :
    ($ s) = ($ s.toList)

content: @[simp] lemma probOutput_uniformSelectFinset [DecidableEq α] (s : Finset α) (x : α) :
    [= x | $ s] = if x ∈ s then (s.card : ℝ≥0∞)⁻¹ else 0

content: @[simp] lemma probFailure_uniformSelectFinset [DecidableEq α] (s : Finset α) :
    [⊥ | $ s] = if s.Nonempty then 0 else 1

-- Target theorem and proof:
@[simp] lemma evalDist_uniformSelectFinset [DecidableEq α] (s : Finset α) :
    evalDist ($ s) = if hs : s.Nonempty then
      OptionT.lift (PMF.uniformOfFinset s hs) else failure
:= by
  refine PMF.ext λ x ↦ ?_
  by_cases hs : s.Nonempty
  · cases x with
    | none =>
        refine (probFailure_uniformSelectFinset _).trans ?_
        simp [hs, OptionT.lift, OptionT.mk]
    | some x =>
        simp only [hs, ↓reduceDIte]
        refine (probOutput_uniformSelectFinset _ _).trans ?_
        simp only [OptionT.lift, OptionT.mk, PMF.monad_pure_eq_pure, PMF.monad_bind_eq_bind,
          PMF.bind_apply, PMF.uniformOfFinset_apply, PMF.pure_apply, Option.some.injEq, mul_ite,
          mul_one, mul_zero]
        refine symm <| (tsum_eq_single x ?_).trans ?_
        · simp only [ne_eq, @eq_comm _ x, ite_eq_right_iff, ENNReal.inv_eq_zero,
            natCast_ne_top, imp_false]
          intros
          tauto
        · simp only [↓reduceIte, ite_eq_ite]
  · simp only [Finset.not_nonempty_iff_eq_empty] at hs
    simp [hs, uniformSelectFinset_def]