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
import Mathlib.InformationTheory.KullbackLeibler.Basic
import Mathlib.MeasureTheory.Integral.Indicator
import Mathlib.Probability.Martingale.Convergence
import Mathlib.InformationTheory.KullbackLeibler.DataProcessing
import Mathlib.MeasureTheory.Function.ConditionalExpectation.RadonNikodym
import Mathlib.InformationTheory.KullbackLeibler.ChainRule
import Mathlib.Probability.Kernel.Composition.RadonNikodym
import Mathlib.Probability.Kernel.Composition.AbsolutelyContinuous
import Mathlib.Probability.Kernel.Composition.Lemmas

/-! # Standalone extraction for `Learning.IsAlgEnvSeq.klDiv_map_trajectory`
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 InformationTheory
end InformationTheory

-- ═══ 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*) := 𝓞 × 𝓐 × 𝓨

/-- 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)]

/-- 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 : ℕ}

/-- A random variable that gives the sequence of rounds. -/
def trajectory (O : ℕ → Ω → 𝓞) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨) (ω : Ω) : ℕ → Round 𝓞 𝓐 𝓨 :=
  fun n ↦ (O n ω, A n ω, Y 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
end Learning
end

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

/-- An oblivious environment without observations, in which the distribution of the next feedback
depends only on the last action, but in a possibly time-dependent manner. -/
@[simps]
noncomputable
def obliviousEnv (ν : ℕ → Kernel 𝓐 𝓨) [∀ n, IsMarkovKernel (ν n)] : Environment Unit 𝓐 𝓨 where
  obs _ := Kernel.const _ (Measure.dirac ())
  feedback n := (ν n).prodMkLeft _

/-- A stationary environment without observations, in which the distribution of the next feedback
depends only on the last action. -/
noncomputable
def stationaryEnv (ν : Kernel 𝓐 𝓨) [IsMarkovKernel ν] : Environment Unit 𝓐 𝓨 :=
  obliviousEnv fun _ ↦ ν

variable {Ω : Type*} {mΩ : MeasurableSpace Ω} {alg : Algorithm Unit 𝓐 𝓨} {ν : Kernel 𝓐 𝓨} [IsMarkovKernel ν] {P : Measure Ω} [IsProbabilityMeasure P] {O : ℕ → Ω → Unit} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓨}
end Learning
end

-- ═══ SequentialLearning.DivergenceDecomposition ═══
section
open MeasureTheory ProbabilityTheory InformationTheory Finset
open scoped ENNReal RealInnerProductSpace ENat
namespace Learning
variable {𝓞 𝓐 𝓨 : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨} {Ω Ω' : Type*} {mΩ : MeasurableSpace Ω} {mΩ' : MeasurableSpace Ω'} {P : Measure Ω} {P' : Measure Ω'} [IsProbabilityMeasure P] [IsProbabilityMeasure P'] {O : ℕ → Ω → 𝓞} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓨} {O' : ℕ → Ω' → 𝓞} {A' : ℕ → Ω' → 𝓐} {Y' : ℕ → Ω' → 𝓨} {alg alg' : Algorithm 𝓞 𝓐 𝓨} {env env' : Environment 𝓞 𝓐 𝓨}
section StationaryEnv
variable {O : ℕ → Ω → Unit} {O' : ℕ → Ω' → Unit} {alg : Algorithm Unit 𝓐 𝓨} {κ κ' : Kernel 𝓐 𝓨} [IsMarkovKernel κ] [IsMarkovKernel κ']

/-- Chain rule for trajectories of a single algorithm versus two stationary environments. -/
lemma IsAlgEnvSeq.klDiv_map_trajectory [MeasurableSpace.CountablyGenerated 𝓨]
    (h : IsAlgEnvSeq O A Y alg (stationaryEnv κ) P)
    (h' : IsAlgEnvSeq O' A' Y' alg (stationaryEnv κ') P') :
    klDiv (P.map (trajectory O A Y)) (P'.map (trajectory O' A' Y')) =
      ∑' t : ℕ, ∫⁻ ω, klDiv (κ (A t ω)) (κ' (A t ω)) ∂P := sorry

end StationaryEnv
end Learning
end