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.Analysis.Normed.Ring.Basic
import Mathlib.MeasureTheory.Constructions.BorelSpace.Basic
import Mathlib.Order.CompletePartialOrder
import Mathlib.Probability.Martingale.BorelCantelli
import Mathlib.Probability.Kernel.Composition.MapComap
import Mathlib.Probability.Independence.Integration
import Mathlib.Probability.Kernel.Representation
import Mathlib.Probability.IdentDistrib
import Mathlib.Probability.Independence.InfinitePi
import Mathlib.MeasureTheory.Function.FactorsThrough

/-! # Standalone extraction for `Bandits.condDistrib_rewardByCount_stepsUntil`
Definitions are copied verbatim; theorem proofs are replaced by `sorry`.
Auto-generated by Referee. -/

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 Bandits
end Bandits

-- ═══ 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 Ω}

/-- A stochastic, sequential algorithm. -/
structure Algorithm (𝓐 𝓨 : Type*) [MeasurableSpace 𝓐] [MeasurableSpace 𝓨] where
  /-- Policy or sampling rule: distribution of the next action. -/
  policy : (n : ℕ) → Kernel (Iic n → 𝓐 × 𝓨) 𝓐
  /-- The policy is a Markov kernel. -/
  [h_policy : ∀ n, IsMarkovKernel (policy n)]
  /-- Distribution of the first action. -/
  p0 : Measure 𝓐
  /-- The first action distribution is a probability measure. -/
  [hp0 : IsProbabilityMeasure p0]

/-- A stochastic environment. -/
structure Environment (𝓐 𝓨 : Type*) [MeasurableSpace 𝓐] [MeasurableSpace 𝓨] where
  /-- Distribution of the next observation as function of the past history. -/
  feedback : (n : ℕ) → Kernel ((Iic n → 𝓐 × 𝓨) × 𝓐) 𝓨
  /-- The feedback kernels are Markov kernels. -/
  [h_feedback : ∀ n, IsMarkovKernel (feedback n)]
  /-- Distribution of the first observation given the first action. -/
  ν0 : Kernel 𝓐 𝓨
  /-- The initial observation kernel is a Markov kernel. -/
  [hp0 : IsMarkovKernel ν0]

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

/-- History of the algorithm-environment sequence up to time `n`. -/
def history (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨) (n : ℕ) (ω : Ω) : Iic n → 𝓐 × 𝓨 :=
  fun i ↦ (A i ω, Y i ω)

/-- An algorithm-environment sequence: a sequence of actions and feedbacks generated
by an algorithm interacting with an environment. -/
structure IsAlgEnvSeq
    (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨) (alg : Algorithm 𝓐 𝓨) (env : Environment 𝓐 𝓨)
    (P : Measure Ω) [IsFiniteMeasure P] : Prop where
  /-- 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 first action has the correct law. -/
  hasLaw_action_zero : HasLaw (fun ω ↦ (A 0 ω)) alg.p0 P
  /-- The first feedback has the correct conditional distribution. -/
  hasCondDistrib_feedback_zero : HasCondDistrib (Y 0) (A 0) env.ν0 P
  /-- The next action has the correct conditional distribution given the history. -/
  hasCondDistrib_action n :
    HasCondDistrib (A (n + 1)) (history A Y n) (alg.policy n) P
  /-- The next feedback has the correct conditional distribution given the history and
  next action. -/
  hasCondDistrib_feedback n :
    HasCondDistrib (Y (n + 1)) (fun ω ↦ (history A Y n ω, A (n + 1) ω))
      (env.feedback n) P

end IsAlgEnvSeq
end Learning
end

-- ═══ SequentialLearning.FiniteActions ═══
section
open MeasureTheory Finset Learning
namespace Learning
variable {𝓐 R Ω : Type*} {m𝓐 : MeasurableSpace 𝓐} {mR : MeasurableSpace R} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] {alg : Algorithm 𝓐 R} {env : Environment 𝓐 R} {P : Measure Ω} [IsProbabilityMeasure P] {A : ℕ → Ω → 𝓐} {R' : ℕ → Ω → R} {a : 𝓐} {m n t : ℕ} {ω : Ω}
section PullCount

/-- Number of times action `a` was chosen up to time `t` (excluding `t`). -/
noncomputable
def pullCount (A : ℕ → Ω → 𝓐) (a : 𝓐) (t : ℕ) (ω : Ω) : ℕ :=
  #(filter (fun s ↦ A s ω = a) (range t))

end PullCount
section StepsUntil

/-- Number of steps until action `a` was pulled exactly `m` times. -/
noncomputable
def stepsUntil (A : ℕ → Ω → 𝓐) (a : 𝓐) (m : ℕ) (ω : Ω) : ℕ∞ :=
  sInf ((↑) '' {s | pullCount A a (s + 1) ω = m})

end StepsUntil
section RewardByCount

/-- Reward obtained when pulling action `a` for the `m`-th time.
If it is never pulled `m` times, the reward is given by the second component of `ω`, which in
applications will be indepedent with same law. -/
noncomputable
def rewardByCount (A : ℕ → Ω → 𝓐) (R' : ℕ → Ω → R) (a : 𝓐) (m : ℕ) (ω : Ω × (ℕ → 𝓐 → R)) : R :=
  match (stepsUntil A a m ω.1) with
  | ⊤ => ω.2 m a
  | (n : ℕ) => R' n ω.1

variable {ω : Ω × (ℕ → 𝓐 → R)}
end RewardByCount
end Learning
end

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

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

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

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

-- ═══ Online.Bandit.ArrayProbSpace ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset Learning
open scoped ENNReal NNReal
namespace Bandits
variable {𝓐 R : Type*} {m𝓐 : MeasurableSpace 𝓐} {mR : MeasurableSpace R}
section MeasureSpace

/-- Measure of an infinite stream of rewards from each action. -/
noncomputable
def streamMeasure (ν : Kernel 𝓐 R) : Measure (ℕ → 𝓐 → R) :=
  Measure.infinitePi fun _ ↦ Measure.infinitePi ν

instance (ν : Kernel 𝓐 R) [IsMarkovKernel ν] : IsProbabilityMeasure (streamMeasure ν) := by
  unfold streamMeasure
  infer_instance

end MeasureSpace
end Bandits
end

-- ═══ Online.Bandit.RewardByCountMeasure ═══
section
open MeasureTheory ProbabilityTheory Finset Learning
open scoped ENNReal NNReal
namespace Bandits
variable {𝓐 Ω : Type*} {m𝓐 : MeasurableSpace 𝓐} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] {A : ℕ → Ω → 𝓐} {R : ℕ → Ω → ℝ} {P : Measure Ω} [IsProbabilityMeasure P] {alg : Algorithm 𝓐 ℝ} {ν : Kernel 𝓐 ℝ} [IsMarkovKernel ν] {h_inter : IsAlgEnvSeq A R alg (stationaryEnv ν) P}
local notation "𝔓" => P.prod (streamMeasure ν)
variable [StandardBorelSpace 𝓐]
variable [Nonempty 𝓐]

/-- The conditional distribution of the reward received at the `m`-th pull of action `a`
given the time at which number of pulls is `m` is the constant kernel with value `ν a`. -/
lemma condDistrib_rewardByCount_stepsUntil [StandardBorelSpace Ω] [Countable 𝓐]
    (h : IsAlgEnvSeq A R alg (stationaryEnv ν) P) (a : 𝓐) (m : ℕ) (hm : m ≠ 0) :
    condDistrib (rewardByCount A R a m) (fun ω ↦ stepsUntil A a m ω.1) 𝔓
      =ᵐ[(𝔓).map (fun ω ↦ stepsUntil A a m ω.1)] Kernel.const _ (ν a) := sorry

end Bandits
end