import Mathlib.MeasureTheory.MeasurableSpace.Embedding
import Mathlib.Order.Restriction
import Mathlib.Probability.Kernel.IonescuTulcea.Maps
import Mathlib.MeasureTheory.Order.Lattice
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.Probability.Kernel.Basic
import Mathlib.Analysis.Normed.Ring.Basic
import Mathlib.MeasureTheory.Constructions.BorelSpace.Basic
import Mathlib.Probability.Kernel.Composition.MapComap

/-! # Standalone extraction for `Learning.IT.bayesTrajMeasurePosterior_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 MeasurableEquiv
end MeasurableEquiv
namespace Function
end Function
namespace MeasurableSpace
end MeasurableSpace
namespace MeasureTheory
end MeasureTheory
namespace ProbabilityTheory
end ProbabilityTheory
namespace ProbabilityTheory.Kernel
end ProbabilityTheory.Kernel
namespace ENNReal
end ENNReal
namespace Learning
end Learning
namespace Learning.Round
end Learning.Round
namespace Learning.IT
end Learning.IT

-- ═══ ForMathlib.MeasureTheory.MeasurableSpace.Embedding ═══
section
open Finset Preorder
namespace MeasurableEquiv
section Fin
variable {X : ℕ → Type*} [∀ n, MeasurableSpace (X n)]

/-- Measurable equivalence between `Π i : Fin (n + 1), X i` and `Π i : Iic n, X i`. -/
def finSuccPiIic (X : ℕ → Type*) [∀ n, MeasurableSpace (X n)] (n : ℕ) :
    (Π i : Fin (n + 1), X i) ≃ᵐ (Π i : Iic n, X i) where
  toFun h i := h ⟨i.1, Nat.lt_succ_of_le (mem_Iic.mp i.2)⟩
  invFun h i := h ⟨i.1, mem_Iic.mpr (Nat.le_of_lt_succ i.2)⟩
  left_inv _ := rfl
  right_inv _ := rfl
  measurable_toFun := .of_eval fun _ ↦ measurable_pi_apply _
  measurable_invFun := .of_eval fun _ ↦ measurable_pi_apply _

end Fin
end MeasurableEquiv
end

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

-- ═══ ForMathlib.Probability.Kernel.IonescuTulcea.Traj ═══
section
open Filter Finset Function MeasurableEquiv MeasurableSpace MeasureTheory Preorder ProbabilityTheory
variable {Ω : Type*} {mΩ : MeasurableSpace Ω} {P : Measure Ω} {X : ℕ → Type*} [∀ n, MeasurableSpace (X n)] {κ : (n : ℕ) → Kernel (Π i : Iic n, X i) (X (n + 1))} [∀ n, IsMarkovKernel (κ n)] {μ₀ : Measure (X 0)} [IsProbabilityMeasure μ₀]
namespace ProbabilityTheory.Kernel
section FinTraj
variable {κ' : (n : ℕ) → Kernel (Π i : Fin n, X i) (X n)} [∀ n, IsMarkovKernel (κ' n)]

/-- Kernels indexed by `Iic n` (as needed for `Kernel.traj`), obtained from kernels indexed by
`Fin n`: the kernel `κ' (n + 1)` on `Π i : Fin (n + 1), X i` is seen as a kernel on
`Π i : Iic n, X i`. -/
noncomputable
def iicOfFin (κ' : (n : ℕ) → Kernel (Π i : Fin n, X i) (X n)) (n : ℕ) :
    Kernel (Π i : Iic n, X i) (X (n + 1)) :=
  (κ' (n + 1)).comap (MeasurableEquiv.finSuccPiIic X n).symm (sorry)

instance (n : ℕ) : IsMarkovKernel (iicOfFin κ' n) := by unfold iicOfFin; infer_instance

/-- Measure on trajectories `Π n, X n` built from kernels `κ' n : Kernel (Π i : Fin n, X i) (X n)`
describing the law of the coordinate `n` given the `n` previous coordinates.
The initial measure is `κ' 0 default`. -/
noncomputable
def trajMeasureFin (κ' : (n : ℕ) → Kernel (Π i : Fin n, X i) (X n)) [∀ n, IsMarkovKernel (κ' n)] :
    Measure (Π n, X n) :=
  trajMeasure (κ' 0 default) (iicOfFin κ')

instance (κ' : (n : ℕ) → Kernel (Π i : Fin n, X i) (X n)) [∀ n, IsMarkovKernel (κ' n)] : IsProbabilityMeasure (trajMeasureFin κ') := sorry

end FinTraj
end ProbabilityTheory.Kernel
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)]

instance (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (env.obs n) := env.isMarkovKernel_obs n

instance (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (env.feedback n) :=
  env.isMarkovKernel_feedback n

/-- Kernel describing the distribution of the round at time `n` given the history before `n`. -/
noncomputable
def stepKernel (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) :
    Kernel (Hist 𝓞 𝓐 𝓨 n) (Round 𝓞 𝓐 𝓨) :=
  env.obs n ⊗ₖ (alg.policy n ⊗ₖ env.feedback n)

instance (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (stepKernel alg env n) := sorry

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

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.IonescuTulceaSpace ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset
open scoped ENNReal NNReal
namespace Learning
variable {𝓞 𝓐 𝓨 Ω : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨} {mΩ : MeasurableSpace Ω}

/-- Measure on the sequence of rounds generated by the algorithm/environment. -/
noncomputable
def trajMeasure (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨) :
    Measure (ℕ → Round 𝓞 𝓐 𝓨) :=
  Kernel.trajMeasureFin (stepKernel alg env)

instance (alg : Algorithm 𝓞 𝓐 𝓨) (env : Environment 𝓞 𝓐 𝓨) : IsProbabilityMeasure (trajMeasure alg env) := sorry

namespace IT

/-- `obs n` is the observation at time `n`. This is a random variable on the measurable space
`ℕ → Round 𝓞 𝓐 𝓨`. -/
def obs (n : ℕ) (h : ℕ → Round 𝓞 𝓐 𝓨) : 𝓞 := (h n).obs

/-- `action n` is the action pulled at time `n`. This is a random variable on the measurable space
`ℕ → Round 𝓞 𝓐 𝓨`. -/
def action (n : ℕ) (h : ℕ → Round 𝓞 𝓐 𝓨) : 𝓐 := (h n).action

/-- `feedback n` is the feedback at time `n`. This is a random variable on the measurable space
`ℕ → Round 𝓞 𝓐 𝓨`. -/
def feedback (n : ℕ) (h : ℕ → Round 𝓞 𝓐 𝓨) : 𝓨 := (h n).feedback

end IT
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
namespace IT

/-- A measure `P` on a measurable space that carries random variables `E`, `A`, and `Y` such that
`IsBayesAlgEnvSeq Q κ alg E A Y P`. The parameter is the observation of the first round,
`IT.obs 0`. -/
noncomputable
def bayesTrajMeasure (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨)
    [IsMarkovKernel κ] (alg : Algorithm Unit 𝓐 𝓨) : Measure (ℕ → Round 𝓔 𝓐 𝓨) :=
  trajMeasure (alg.comapObs (fun _ : 𝓔 ↦ ())) (bayesStationaryEnv Q κ)

instance (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ] (alg : Algorithm Unit 𝓐 𝓨) : IsProbabilityMeasure (bayesTrajMeasure Q κ alg) := sorry

/-- A kernel that represents the posterior over `E` given the history before time `n`. -/
noncomputable
def bayesTrajMeasurePosterior [StandardBorelSpace 𝓔] [Nonempty 𝓔]
    (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ]
    (alg : Algorithm Unit 𝓐 𝓨) (n : ℕ) : Kernel (Hist Unit 𝓐 𝓨 n) 𝓔 :=
  condDistrib (obs 0) (history (noObs _) action feedback n) (bayesTrajMeasure Q κ alg)

instance [StandardBorelSpace 𝓔] [Nonempty 𝓔] (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ] (alg : Algorithm Unit 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (bayesTrajMeasurePosterior Q κ alg n) := sorry

/-- The posterior given the empty history is the prior. -/
lemma bayesTrajMeasurePosterior_zero [StandardBorelSpace 𝓔] [Nonempty 𝓔]
    (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ]
    (alg : Algorithm Unit 𝓐 𝓨) :
    bayesTrajMeasurePosterior Q κ alg 0 = Kernel.const _ Q := sorry

end IT
end Learning
end