-- VCVio/OracleComp/Coercions/SubSpec.lean

import VCVio.OracleComp.SimSemantics.SimulateQ

import VCVio.OracleComp.Constructions.UniformSelect

open OracleSpec OracleComp BigOperators ENNReal

variable {ι : Type u} {τ : Type v}
  {spec : OracleSpec ι} {superSpec : OracleSpec τ} {α β γ : Type w}

namespace OracleSpec

infix : 50 " ⊂ₒ " => SubSpec

namespace SubSpec

variable [h : MonadLift (OracleQuery spec) (OracleQuery superSpec)]

end SubSpec

end OracleSpec

namespace OracleComp

section liftComp

def liftComp (oa : OracleComp spec α) (superSpec : OracleSpec τ)
      [h : MonadLift (OracleQuery spec) (OracleQuery superSpec)] :
      OracleComp superSpec α := simulateQ ⟨liftM⟩ oa

variable (superSpec : OracleSpec τ) [h : MonadLift (OracleQuery spec) (OracleQuery superSpec)]

-- Repository definitions used:
content: def simulateQ [AlternativeMonad m] (so : QueryImpl spec m) (oa : OracleComp spec α) : m α :=
  do Option.getM (← FreeMonad.mapM oa.run so.impl)

content: instance QueryImpl.Inhabited [∀ i, Inhabited (spec.range i)] [Pure m] :
  Inhabited (QueryImpl spec m) := ⟨{impl q := pure q.defaultOutput}⟩

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

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: @[inline, reducible]
protected def range (spec : OracleSpec ι) (i : ι) : Type w := (spec i).2

content: protected def mapM [Pure m] [Bind m] :
    (oa : FreeMonad f α) → (s : {α : Type u} → f α → m α) → m α
  | .pure x, _ => pure x
  | .roll x r, s => s x >>= λ u ↦ (r u).mapM s

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 OracleComp {ι : Type u} (spec : OracleSpec.{u,v} ι) :
    Type w → Type (max u (v + 1) w) :=
  OptionT (FreeMonad (OracleQuery.{u,v} spec))

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: @[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: @[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 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: class FiniteRange (spec : OracleSpec ι) where
    range_inhabited' (i : ι) : Inhabited (spec.range i)
    range_fintype' (i : ι) : Fintype (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 -/

-- Repository lemmas used:
content: @[simp] lemma simulateQ_bind (oa : OracleComp spec α) (ob : α → OracleComp spec β) :
    simulateQ so (oa >>= ob) = simulateQ so oa >>= (simulateQ so ∘ ob)

content: @[simp] lemma simulateQ_query_bind (q : OracleQuery spec α) (ob : α → OracleComp spec β) :
    simulateQ so (liftM q >>= ob) = so.impl q >>= (simulateQ so ∘ ob)

content: @[simp]
lemma evalDist_liftM [Nonempty α] [Fintype α] (q : OracleQuery spec α) :
    evalDist (q : OracleComp spec α) = OptionT.lift (PMF.uniformOfFintype α)

content: @[simp] lemma simulateQ_query (q : OracleQuery spec α) : simulateQ so q = so.impl q

-- Local definitions used:
content: class SubSpec (spec : OracleSpec.{u,w} ι) (superSpec : OracleSpec τ)
  extends MonadLift (OracleQuery spec) (OracleQuery superSpec) where

content: def liftComp (oa : OracleComp spec α) (superSpec : OracleSpec τ)
      [h : MonadLift (OracleQuery spec) (OracleQuery superSpec)] :
      OracleComp superSpec α := simulateQ ⟨liftM⟩ oa

-- Local lemmas used:
content: @[simp]
lemma liftComp_pure (x : α) : liftComp (pure x : OracleComp spec α) superSpec = pure x

-- Target theorem and proof:
@[simp]
lemma evalDist_liftComp [spec.FiniteRange] [superSpec.FiniteRange]
    (oa : OracleComp spec α) : evalDist (liftComp oa superSpec) = evalDist oa
:= by
  induction oa using OracleComp.inductionOn with
  | pure x => simp [liftComp_pure]
  | query_bind i t oa hoa =>
      simp only [liftComp, simulateQ_bind, simulateQ_query, StateT.run'_eq, StateT.run_bind,
        StateT.run_monadLift, SubSpec.liftM_query_eq_liftM_liftM, bind_pure_comp,
        Function.comp_apply, bind_map_left, map_bind, evalDist_bind, OracleComp.evalDist_liftM]
      refine congr_arg (_ >>= ·) (funext λ u ↦ ?_)
      simpa [StateT.run, liftComp] using hoa u
  | failure => simp