import Mathlib.MeasureTheory.Order.Lattice
import Mathlib.Order.CompletePartialOrder
import Mathlib.Probability.Martingale.BorelCantelli
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.Probability.Kernel.Composition.MapComap
import Mathlib.CategoryTheory.Countable
import Mathlib.MeasureTheory.Constructions.Polish.Basic
import Mathlib.Probability.Independence.Integration
import Mathlib.Probability.Kernel.Representation
import Mathlib.Probability.IdentDistrib
import Mathlib.Probability.Independence.InfinitePi
import Mathlib.MeasureTheory.Function.FactorsThrough
import Mathlib.Probability.Moments.SubGaussian
import Mathlib.MeasureTheory.Measure.AbsolutelyContinuous
import Mathlib.MeasureTheory.OuterMeasure.Basic
import Mathlib.Probability.Kernel.CompProdEqIff
import Mathlib.Probability.Kernel.Composition.MeasureComp
import Mathlib.Probability.Kernel.Composition.MeasureCompProd

/-! # Standalone extraction for `Bandits.ClippedUCB.measurable_uncurry_ucb_comp`
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 Learning
end Learning
namespace ProbabilityTheory
end ProbabilityTheory
namespace Bandits
end Bandits
namespace Bandits.ClippedUCB
end Bandits.ClippedUCB

-- ═══ 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 𝓐} {mR : MeasurableSpace R} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] {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
end Learning
end

-- ═══ SequentialLearning.SumRewards ═══
section
open MeasureTheory Finset Learning
namespace Learning
variable {𝓐 𝓨 Ω : Type*} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] [AddCommGroup 𝓨] {P : Measure Ω} [IsProbabilityMeasure P] {A : ℕ → Ω → 𝓐} {R : ℕ → Ω → 𝓨} {a : 𝓐} {m n t : ℕ} {ω : Ω}

/-- Sum of rewards obtained when pulling action `a` up to time `t` (exclusive). -/
noncomputable def sumRewards (A : ℕ → Ω → 𝓐) (R : ℕ → Ω → 𝓨) (a : 𝓐) (t : ℕ) (ω : Ω) : 𝓨 :=
  ∑ s ∈ range t, if A s ω = a then (R s) ω else 0

/-- Empirical mean reward obtained when pulling action `a` up to time `t` (exclusive). -/
noncomputable
def empMean (A : ℕ → Ω → 𝓐) (R : ℕ → Ω → ℝ) (a : 𝓐) (t : ℕ) (ω : Ω) : ℝ :=
  sumRewards A R a t ω / pullCount A a t ω

variable [MeasurableSingletonClass 𝓐]
end Learning
end

-- ═══ Online.Bandit.Algorithms.Regret.BayesRegretTS ═══
section
open MeasureTheory ProbabilityTheory Finset Learning
namespace Bandits
namespace ClippedUCB
variable {K : ℕ} {l u σ2 δ : ℝ}
variable {Ω : Type*} {A : ℕ → Ω → Fin K} {R : ℕ → Ω → ℝ}

/-- Clipped upper confidence bound used in the regret analysis of Thompson sampling. -/
noncomputable
def ucb (A : ℕ → Ω → Fin K) (R : ℕ → Ω → ℝ) (l u σ2 δ : ℝ) (a : Fin K) (n : ℕ) (ω : Ω) : ℝ :=
  if pullCount A a n ω = 0 then u
  else max l (min u (empMean A R a n ω + √(2 * σ2 * Real.log (1 / δ) / (pullCount A a n ω))))

@[fun_prop]
lemma measurable_uncurry_ucb_comp [MeasurableSpace Ω] (hA : ∀ t, Measurable (A t))
    (hR : ∀ t, Measurable (R t)) {f : Ω → Fin K} (hf : Measurable f) {g : Ω → ℕ}
    (hg : Measurable g) : Measurable (fun ω ↦ ucb A R l u σ2 δ (f ω) (g ω) ω) := sorry

variable [Nonempty (Fin K)]
variable [MeasurableSpace Ω]
variable {𝓔 : Type*} [MeasurableSpace 𝓔]
variable {E : Ω → 𝓔}
variable {Q : Measure 𝓔} {κ : Kernel (𝓔 × Fin K) ℝ} [IsMarkovKernel κ]
variable {P : Measure Ω} [IsProbabilityMeasure P]
end ClippedUCB
end Bandits
end