import Mathlib.MeasureTheory.Order.Lattice
import Mathlib.Order.CompletePartialOrder
import Mathlib.Probability.Martingale.BorelCantelli
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.Representation
import Mathlib.Probability.Kernel.Composition.MapComap
import Mathlib.Probability.IdentDistrib
import Mathlib.Probability.Independence.InfinitePi
import Mathlib.MeasureTheory.Function.FactorsThrough
import Mathlib.MeasureTheory.Constructions.Cylinders
import Mathlib.MeasureTheory.Integral.Indicator
import Mathlib.Probability.ConditionalProbability
import Mathlib.Probability.HasLaw
import Mathlib.MeasureTheory.Measure.Prod

/-! # Standalone extraction for `Bandits.hasLaw_rewardByCountUntil_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 Learning
end Learning
namespace ProbabilityTheory
end ProbabilityTheory
namespace ENNReal
end ENNReal
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.FiniteActions ═══
section
open MeasureTheory Finset Learning
namespace Learning
variable {𝓞 𝓐 R Ω : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {mR : MeasurableSpace R} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] {P : Measure Ω} [IsProbabilityMeasure P] {O : ℕ → Ω → 𝓞} {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 RewardByCount
variable {ω : Ω × (ℕ → 𝓐 → R)}

/-- Array of rewards by count, truncated at time `t`: the entry `(a, m)` is the reward received at
the `(m + 1)`-th pull of action `a` if that pull happened before time `t`, and the entry
`(m + 1, a)` of the auxiliary array `ω.2` otherwise.

This is an auxiliary definition used to prove results about the distribution of `rewardByCount`.

It is defined recursively: at time `t`, the
entry `(A t, pullCount A (A t) t)` is replaced by the reward `R' t`.
See `rewardByCountUntil_apply_of_lt_pullCount` and `rewardByCountUntil_apply_of_pullCount_le`. -/
noncomputable
def rewardByCountUntil (A : ℕ → Ω → 𝓐) (R' : ℕ → Ω → R) : ℕ → Ω × (ℕ → 𝓐 → R) → 𝓐 × ℕ → R
  | 0, ω => fun p ↦ ω.2 (p.2 + 1) p.1
  | t + 1, ω => Function.update (rewardByCountUntil A R' t ω)
      (A t ω.1, pullCount A (A t ω.1) t ω.1) (R' t ω.1)

end RewardByCount
end Learning
end

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

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

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 𝓐] {O : ℕ → Ω → Unit} {A : ℕ → Ω → 𝓐} {R : ℕ → Ω → ℝ} {P : Measure Ω} [IsProbabilityMeasure P] {ν : Kernel 𝓐 ℝ} [IsMarkovKernel ν]
local notation "𝔓" => P.prod (streamMeasure ν)
section Independence

/-- The law of the array `rewardByCountUntil A R 0`, which is a sub-array of the auxiliary array,
is the product measure `⨂ (a, m), ν a`. -/
lemma hasLaw_rewardByCountUntil_zero (μ : Measure Ω) [IsProbabilityMeasure μ] :
    HasLaw (rewardByCountUntil A R 0) (Measure.infinitePi fun p : 𝓐 × ℕ ↦ ν p.1)
      (μ.prod (streamMeasure ν)) := sorry

variable [MeasurableSingletonClass 𝓐]
variable [Countable 𝓐]
end Independence
end Bandits
end