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.CategoryTheory.Countable
import Mathlib.MeasureTheory.Constructions.Polish.Basic
import Mathlib.Probability.Kernel.Composition.MapComap
import Mathlib.Analysis.SpecialFunctions.Integrals.Basic
import Mathlib.Analysis.SumIntegralComparisons
import Mathlib.Probability.Kernel.Representation
import Mathlib.Probability.IdentDistrib
import Mathlib.Probability.Independence.InfinitePi
import Mathlib.MeasureTheory.Function.FactorsThrough
import Mathlib.Probability.Moments.SubGaussian

/-! # Standalone extraction for `Bandits.UCB.gap_le_two_mul_ucbWidth`
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
namespace Bandits.UCB
end Bandits.UCB

-- ═══ 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
end Learning
end

-- ═══ Online.Bandit.Regret ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset Learning
open scoped ENNReal NNReal
namespace Bandits
variable {𝓞 𝓐 Ω : Type*} [DecidableEq 𝓐] {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {mΩ : MeasurableSpace Ω} {ν : Kernel 𝓐 ℝ} {O : ℕ → Ω → 𝓞} {A : ℕ → Ω → 𝓐} {R : ℕ → Ω → ℝ} {ω : Ω} {m n t : ℕ} {a : 𝓐}

/-- Gap of an action `a`: difference between the highest mean of the actions and the mean of `a`. -/
noncomputable
def gap (ν : Kernel 𝓐 ℝ) (a : 𝓐) : ℝ := (⨆ i, (ν i)[id]) - (ν a)[id]

section bestArm
variable [Fintype 𝓐] [Nonempty 𝓐]

/-- action with the highest mean. -/
noncomputable def bestArm (ν : Kernel 𝓐 ℝ) : 𝓐 :=
  (exists_max_image univ (fun a ↦ (ν a)[id]) (univ_nonempty_iff.mpr inferInstance)).choose

end bestArm
end Bandits
end

-- ═══ SequentialLearning.SumRewards ═══
section
open MeasureTheory Finset Learning
namespace Learning
variable {𝓞 𝓐 𝓨 Ω : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] [AddCommGroup 𝓨] {P : Measure Ω} [IsProbabilityMeasure P] {O : ℕ → Ω → 𝓞} {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.UCB ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset Learning
open scoped ENNReal NNReal
namespace Bandits
variable {K : ℕ}
namespace UCB
variable [NeZero K] {c : ℝ} {ν : Kernel (Fin K) ℝ} [IsMarkovKernel ν] {Ω : Type*} {mΩ : MeasurableSpace Ω} {P : Measure Ω} [IsProbabilityMeasure P] {O : ℕ → Ω → Unit} {A : ℕ → Ω → Fin K} {R : ℕ → Ω → ℝ} {σ2 : ℝ≥0} {n : ℕ} {ω : Ω}
section AlgorithmBehavior

/-- The exploration bonus of the UCB algorithm, which corresponds to the width of
a confidence interval. -/
noncomputable def ucbWidth (A : ℕ → Ω → Fin K) (c : ℝ) (a : Fin K) (n : ℕ) (ω : Ω) : ℝ :=
  √(2 * c * log (n + 1) / pullCount A a n ω)

end AlgorithmBehavior
end UCB
end Bandits
end

-- ═══ Online.Bandit.Algorithms.Regret.UCB ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset Learning
open scoped ENNReal NNReal
namespace Bandits
namespace UCB
variable {K : ℕ} [NeZero K] {c : ℝ} {ν : Kernel (Fin K) ℝ} [IsMarkovKernel ν] {Ω : Type*} {mΩ : MeasurableSpace Ω} {P : Measure Ω} [IsProbabilityMeasure P] {O : ℕ → Ω → Unit} {A : ℕ → Ω → Fin K} {R : ℕ → Ω → ℝ} {σ2 : ℝ≥0} {n : ℕ} {ω : Ω}

omit [IsMarkovKernel ν] in
/-- If the means of the best arm and of arm `b` lie in their confidence intervals and the UCB index
of `b` is at least that of the best arm, then the gap of `b` is at most twice its confidence
width. -/
lemma gap_le_two_mul_ucbWidth {b : Fin K}
    (h_best : (ν (bestArm ν))[id] ≤ empMean A R (bestArm ν) n ω + ucbWidth A c (bestArm ν) n ω)
    (h_arm : empMean A R b n ω - ucbWidth A c b n ω ≤ (ν b)[id])
    (h_le : empMean A R (bestArm ν) n ω + ucbWidth A c (bestArm ν) n ω ≤
      empMean A R b n ω + ucbWidth A c b n ω) :
    gap ν b ≤ 2 * ucbWidth A c b n ω := sorry

end UCB
end Bandits
end