import Mathlib.MeasureTheory.Order.Lattice
import Mathlib.Probability.Kernel.Basic
import Mathlib.Probability.Kernel.IonescuTulcea.Traj
import Mathlib.Probability.Process.FiniteDimensionalLaws
import Mathlib.Probability.HasCondDistrib
import Mathlib.MeasureTheory.Measure.ProbabilityMeasure
import Mathlib.Probability.Independence.Basic
import Mathlib.Probability.Independence.Conditional
import Mathlib.MeasureTheory.Measure.SubFinite
import Mathlib.Probability.Kernel.RadonNikodym
import Mathlib.MeasureTheory.MeasurableSpace.Embedding
import Mathlib.Order.Restriction
import Mathlib.Probability.Kernel.IonescuTulcea.Maps
import Mathlib.Analysis.Normed.Ring.Basic
import Mathlib.MeasureTheory.Constructions.BorelSpace.Basic
import Mathlib.Probability.Kernel.Composition.MapComap

/-! # Standalone extraction for `Learning.IsBayesAlgEnvSeq.condDistrib_param_history_zero`
Definitions are copied verbatim; theorem proofs are replaced by `sorry`.
Auto-generated by ChallengeGen. -/

set_option quotPrecheck false

-- Namespace stubs (so later `open`s resolve).
namespace Finset
end Finset
namespace MeasureTheory
end MeasureTheory
namespace ProbabilityTheory
end ProbabilityTheory
namespace ENNReal
end ENNReal
namespace Learning
end Learning
namespace Learning.Round
end Learning.Round
namespace Learning.IsBayesAlgEnvSeq
end Learning.IsBayesAlgEnvSeq

-- ═══ ForMathlib.MeasureTheory.Order.Lattice ═══
section
open Finset
variable {α δ : Type*} [MeasurableSpace δ] [SemilatticeInf α] {m : MeasurableSpace α} [MeasurableInf₂ α]
attribute [to_dual existing] MeasurableInf₂
end

-- ═══ SequentialLearning.Algorithm ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset
open scoped ENNReal NNReal
namespace Learning
variable {𝓞 𝓐 𝓨 Ω : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨} {mΩ : MeasurableSpace Ω}

/-- One round of interaction: an observation, then an action, then a feedback. -/
abbrev Round (𝓞 𝓐 𝓨 : Type*) := 𝓞 × 𝓐 × 𝓨

/-- The observation of a round. -/
def Round.obs (r : Round 𝓞 𝓐 𝓨) : 𝓞 := r.1

/-- The action of a round. -/
def Round.action (r : Round 𝓞 𝓐 𝓨) : 𝓐 := r.2.1

/-- The feedback of a round. -/
def Round.feedback (r : Round 𝓞 𝓐 𝓨) : 𝓨 := r.2.2

namespace Round

@[fun_prop]
lemma measurable_obs : Measurable (Round.obs (𝓞 := 𝓞) (𝓐 := 𝓐) (𝓨 := 𝓨)) := sorry

end Round

/-- History of `n` complete rounds; `n = 0` is the empty history. -/
abbrev Hist (𝓞 𝓐 𝓨 : Type*) (n : ℕ) := Fin n → Round 𝓞 𝓐 𝓨

/-- A stochastic, sequential algorithm.
At each round, it sees an observation in `𝓞`, then takes an action in `𝓐`, and finally receives
feedback in `𝓨`. The action is a random function of the past rounds and the current observation. -/
@[ext]
structure Algorithm (𝓞 𝓐 𝓨 : Type*) [MeasurableSpace 𝓞] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨]
    where
  /-- Law of the action of round `n` given the past rounds and the current observation. -/
  policy : (n : ℕ) → Kernel (Hist 𝓞 𝓐 𝓨 n × 𝓞) 𝓐
  /-- The policy is a Markov kernel. -/
  [isMarkovKernel_policy : ∀ n, IsMarkovKernel (policy n)]

instance (alg : Algorithm 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (alg.policy n) :=
  alg.isMarkovKernel_policy n

/-- A stochastic environment.
At each round, an observation is drawn prior to the algorithm taking an action. Then the environment
provides feedback based on the observation and the action. -/
@[ext]
structure Environment (𝓞 𝓐 𝓨 : Type*) [MeasurableSpace 𝓞] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨]
    where
  /-- Law of the observation of round `n` given the past rounds. -/
  obs : (n : ℕ) → Kernel (Hist 𝓞 𝓐 𝓨 n) 𝓞
  /-- Law of the feedback of round `n` given the past rounds, the observation and the action. -/
  feedback : (n : ℕ) → Kernel ((Hist 𝓞 𝓐 𝓨 n × 𝓞) × 𝓐) 𝓨
  /-- The observation kernel is a Markov kernel. -/
  [isMarkovKernel_obs : ∀ n, IsMarkovKernel (obs n)]
  /-- The feedback kernel is a Markov kernel. -/
  [isMarkovKernel_feedback : ∀ n, IsMarkovKernel (feedback n)]

section IsAlgEnvSeq
variable {O : ℕ → Ω → 𝓞} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓨} {alg : Algorithm 𝓞 𝓐 𝓨} {env : Environment 𝓞 𝓐 𝓨} {P : Measure Ω} [IsFiniteMeasure P] {N : ℕ}

/-- History of the algorithm-environment sequence before time `n`: the rounds at
times `0, ..., n - 1`. -/
def history (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨) (n : ℕ) (ω : Ω) : Hist 𝓞 𝓐 𝓨 n :=
  fun i ↦ (O i ω, A i ω, Y i ω)

/-- An algorithm-environment sequence: a sequence of observations, actions and feedbacks generated
by an algorithm interacting with an environment. -/
structure IsAlgEnvSeq
    (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨)
    (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨)
    (P : Measure Ω) [IsFiniteMeasure P] : Prop where
  /-- The observation sequence is measurable. -/
  measurable_obs n : Measurable (O n) := sorry
  /-- The action sequence is measurable. -/
  measurable_action n : Measurable (A n) := sorry
  /-- The feedback sequence is measurable. -/
  measurable_feedback n : Measurable (Y n) := sorry
  /-- The observation at time `n` has the correct conditional distribution given the history. -/
  hasCondDistrib_obs n :
    HasCondDistrib (O n) (history O A Y n) (env.obs n) P
  /-- The action at time `n` has the correct conditional distribution given the history and the
  observation at time `n`. -/
  hasCondDistrib_action n :
    HasCondDistrib (A n) (fun ω ↦ (history O A Y n ω, O n ω)) (alg.policy n) P
  /-- The feedback at time `n` has the correct conditional distribution given the history, the
  observation and the action at time `n`. -/
  hasCondDistrib_feedback n :
    HasCondDistrib (Y n) (fun ω ↦ ((history O A Y n ω, O n ω), A n ω)) (env.feedback n) P

end IsAlgEnvSeq
section NoObservation

/-- The observation process of an algorithm-environment sequence without observations. -/
def noObs (Ω : Type*) : ℕ → Ω → Unit := fun _ _ ↦ ()

end NoObservation
end Learning
end

-- ═══ SequentialLearning.Comap ═══
section
open MeasureTheory ProbabilityTheory
namespace Learning
variable {𝓞 𝓞' 𝓞'' 𝓐 𝓐' 𝓐'' 𝓨 𝓨' 𝓨'' Ω : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓞' : MeasurableSpace 𝓞'} {m𝓞'' : MeasurableSpace 𝓞''} {m𝓐 : MeasurableSpace 𝓐} {m𝓐' : MeasurableSpace 𝓐'} {m𝓐'' : MeasurableSpace 𝓐''} {m𝓨 : MeasurableSpace 𝓨} {m𝓨' : MeasurableSpace 𝓨'} {m𝓨'' : MeasurableSpace 𝓨''} {mΩ : MeasurableSpace Ω} {fo : 𝓞 → 𝓞'} {fa : 𝓐 → 𝓐'} {fy : 𝓨 → 𝓨'} {go : 𝓞' → 𝓞''} {ga : 𝓐' → 𝓐''} {gy : 𝓨' → 𝓨''}
section Map

/-- Transport a round along maps of the observation, the action and the feedback. -/
def Round.map (fo : 𝓞 → 𝓞') (fa : 𝓐 → 𝓐') (fy : 𝓨 → 𝓨') (r : Round 𝓞 𝓐 𝓨) : Round 𝓞' 𝓐' 𝓨' :=
  (fo r.obs, fa r.action, fy r.feedback)

/-- Transport a history round-wise. -/
def Hist.map (fo : 𝓞 → 𝓞') (fa : 𝓐 → 𝓐') (fy : 𝓨 → 𝓨') {n : ℕ} (h : Hist 𝓞 𝓐 𝓨 n) :
    Hist 𝓞' 𝓐' 𝓨' n :=
  fun i ↦ Round.map fo fa fy (h i)

@[fun_prop]
lemma Hist.measurable_map (hfo : Measurable fo) (hfa : Measurable fa) (hfy : Measurable fy)
    (n : ℕ) :
    Measurable (Hist.map fo fa fy (n := n)) := sorry

/-- Transport the observations of a history. -/
abbrev Hist.mapObs (f : 𝓞 → 𝓞') {n : ℕ} (h : Hist 𝓞 𝓐 𝓨 n) : Hist 𝓞' 𝓐 𝓨 n :=
  Hist.map f id id h

end Map
section Comap

/-- The algorithm with observations in `𝓞'` and feedbacks in `𝓨'` obtained from
`alg : Algorithm 𝓞 𝓐 𝓨` by transforming the pair (past rounds, current observation) by `F n` at
each round `n` before applying the policy of `alg`.

This is the primitive transport operation on algorithms: `Algorithm.comapObs` and
`Algorithm.comapFeedback` are the special cases in which `F n` is a round-wise map of the
observation and of the feedback. -/
def Algorithm.comap (alg : Algorithm 𝓞 𝓐 𝓨)
    (F : (n : ℕ) → Hist 𝓞' 𝓐 𝓨' n × 𝓞' → Hist 𝓞 𝓐 𝓨 n × 𝓞) (hF : ∀ n, Measurable (F n)) :
    Algorithm 𝓞' 𝓐 𝓨' where
  policy n := (alg.policy n).comap (F n) (hF n)

section ComapObs
variable {f : 𝓞' → 𝓞}

/-- The algorithm that sees `f o` when the observation is `o`, both in the current round and in the
past rounds. -/
def Algorithm.comapObs (alg : Algorithm 𝓞 𝓐 𝓨) (f : 𝓞' → 𝓞)
    (hf : Measurable f := by fun_prop) : Algorithm 𝓞' 𝓐 𝓨 :=
  alg.comap (fun _ p ↦ (Hist.mapObs f p.1, f p.2)) fun _ ↦ sorry

end ComapObs
end Comap
end Learning
end

-- ═══ SequentialLearning.BayesStationaryEnv ═══
section
open MeasureTheory ProbabilityTheory Finset
open scoped ENNReal NNReal
namespace Learning
variable {𝓔 𝓐 𝓨 Ω : Type*}
variable [MeasurableSpace 𝓔] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨] [MeasurableSpace Ω]
section BayesEnv

/-- The environment that draws a parameter `e : 𝓔` from `Q` before the first round, announces it as
the observation of every round, and returns a feedback drawn from `κ (e, a)` when the action is
`a`. The algorithm is meant to ignore the announced parameter, that is, to run through
`Algorithm.comapObs (fun _ ↦ ())`. -/
noncomputable
def bayesStationaryEnv (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨)
    [IsMarkovKernel κ] : Environment 𝓔 𝓐 𝓨 where
  obs
    | 0 => Kernel.const _ Q
    | _ + 1 => Kernel.deterministic (fun h ↦ (h 0).obs) (sorry)
  feedback _ := κ.comap (fun p ↦ (p.1.2, p.2)) (sorry)
  isMarkovKernel_obs n := by cases n <;> infer_instance

variable {Q : Measure 𝓔} [IsProbabilityMeasure Q] {κ : Kernel (𝓔 × 𝓐) 𝓨} [IsMarkovKernel κ]
end BayesEnv

/-- `IsBayesAlgEnvSeq Q κ alg E A Y P` states that there is a measure `P : Measure Ω` such
that the parameter `E : Ω → 𝓔` has law `Q` and that the sequences of actions `A : ℕ → Ω → 𝓐`
and feedbacks `Y : ℕ → Ω → 𝓨` are generated by the algorithm `alg : Algorithm Unit 𝓐 𝓨`
interacting with an underlying environment that depends on `E` and `κ`
(`stationaryEnv (κ.sectR (E ω))`).

This is `IsAlgEnvSeq` for the announcing environment `bayesStationaryEnv Q κ` and the algorithm
`alg.comapObs (fun _ ↦ ())` that ignores the announced parameter: the observation at every round
is `E ω`. -/
def IsBayesAlgEnvSeq (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨)
    [IsMarkovKernel κ] (alg : Algorithm Unit 𝓐 𝓨)
    (E : Ω → 𝓔) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨)
    (P : Measure Ω) [IsProbabilityMeasure P] : Prop :=
  IsAlgEnvSeq (fun _ ↦ E) A Y (alg.comapObs (fun _ : 𝓔 ↦ ())) (bayesStationaryEnv Q κ) P

namespace IsBayesAlgEnvSeq
variable {Q : Measure 𝓔} [IsProbabilityMeasure Q] {κ : Kernel (𝓔 × 𝓐) 𝓨} [IsMarkovKernel κ] {alg : Algorithm Unit 𝓐 𝓨} {P : Measure Ω} [IsProbabilityMeasure P] {E : Ω → 𝓔} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓨}

/-- The posterior over the parameter given the empty history is the prior. -/
lemma condDistrib_param_history_zero [StandardBorelSpace 𝓔] [Nonempty 𝓔]
    (h : IsBayesAlgEnvSeq Q κ alg E A Y P) :
    condDistrib E (history (noObs Ω) A Y 0) P = Kernel.const _ Q := sorry

variable [StandardBorelSpace 𝓐] [Nonempty 𝓐] [StandardBorelSpace 𝓨] [Nonempty 𝓨]
end IsBayesAlgEnvSeq
end Learning
end