import Mathlib.MeasureTheory.Order.Lattice
import Mathlib.CategoryTheory.Countable
import Mathlib.MeasureTheory.Constructions.Polish.Basic
import Mathlib.Order.CompletePartialOrder
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.Martingale.BorelCantelli
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.sum_indicator_good_event_eq_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 Function
end Function
namespace MeasureTheory
end MeasureTheory
namespace ProbabilityTheory
end ProbabilityTheory
namespace ENNReal
end ENNReal
namespace Learning
end Learning
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

-- ═══ ForMathlib.MeasureTheory.Order.MeasurableArg ═══
section
open Finset
variable {ι α : Type*} [LinearOrder α] [Fintype ι] [Nonempty ι] (f : ι → α)
namespace Function

/-- The maximum value of a tuple. -/
@[to_dual /-- The minimum value of a tuple. -/]
abbrev max : α := univ.sup' univ_nonempty f

end Function
section Argmax

@[to_dual]
lemma exists_argmax : ∃ i, f i = f.max := sorry

/-- The index of the maximum value of a tuple. -/
@[to_dual /-- The index of the minimum value of a tuple. -/]
noncomputable def argmax := (exists_argmax f).choose

variable [MeasurableSpace α]
end Argmax
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*) := 𝓞 × 𝓐 × 𝓨

/-- The action of a round. -/
def Round.action (r : Round 𝓞 𝓐 𝓨) : 𝓐 := r.2.1

/-- The feedback of a round. -/
def Round.feedback (r : Round 𝓞 𝓐 𝓨) : 𝓨 := r.2.2

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

/-- 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.FiniteActions ═══
section
open MeasureTheory Finset Learning
namespace Learning
variable {𝓞 𝓐 R Ω : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {mR : MeasurableSpace R} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] {alg : Algorithm 𝓞 𝓐 R} {env : Environment 𝓞 𝓐 R} {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))

/-- Number of pulls of arm `a` in the history before time `n`.
This is the number of entries in `h` in which the arm is `a`. -/
noncomputable
def pullCount' (n : ℕ) (h : Hist 𝓞 𝓐 R n) (a : 𝓐) := #{s | (h s).action = a}

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.Deterministic ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset
open scoped ENNReal NNReal
namespace Learning
variable {𝓞 𝓐 𝓨 : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨}
variable {nextA : (n : ℕ) → (Hist 𝓞 𝓐 𝓨 n × 𝓞) → 𝓐} {h_next : ∀ n, Measurable (nextA n)} {env : Environment 𝓞 𝓐 𝓨} {obs : (n : ℕ) → Kernel (Hist 𝓞 𝓐 𝓨 n) 𝓞} [∀ n, IsMarkovKernel (obs n)] {f : (n : ℕ) → ((Hist 𝓞 𝓐 𝓨 n × 𝓞) × 𝓐) → 𝓨} {hf : ∀ n, Measurable (f n)}

/-- A deterministic algorithm, which chooses the action given by the function `nextA`.
The initial action is `fun o ↦ nextA 0 (default, o)`. -/
@[simps]
noncomputable
def detAlgorithm (nextA : (n : ℕ) → (Hist 𝓞 𝓐 𝓨 n × 𝓞) → 𝓐)
    (h_next : ∀ n, Measurable (nextA n)) :
    Algorithm 𝓞 𝓐 𝓨 where
  policy n := Kernel.deterministic (nextA n) (h_next n)

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.Algorithms.RoundRobin ═══
section
open MeasureTheory ProbabilityTheory Finset Learning
open scoped ENNReal NNReal
namespace Learning
variable {𝓞 𝓨 : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓨 : MeasurableSpace 𝓨} {K : ℕ}
section AlgorithmDefinition

variable (K) in
/-- Action chosen by the Round-Robin algorithm at time `n`. This is action `n % K`. -/
noncomputable
def RoundRobin.nextAction [NeZero K] (n : ℕ) : Fin K :=
  ⟨n % K, Nat.mod_lt _ (Nat.pos_of_neZero K)⟩

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

/-- Sum of rewards of arm `a` in the history before time `n`. -/
noncomputable
def sumRewards' (n : ℕ) (h : Hist 𝓞 𝓐 𝓨 n) (a : 𝓐) :=
  ∑ s, if (h s).action = a then (h s).feedback 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 ω

/-- Empirical mean of arm `a` in the history before time `n`. -/
noncomputable
def empMean' (n : ℕ) (h : Hist 𝓞 𝓐 ℝ n) (a : 𝓐) :=
  sumRewards' n h a / pullCount' n h a

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 : ℕ}
section Algorithm

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

/-- Arm pulled by the UCB algorithm at time `n`, as a function of the history before `n`. -/
noncomputable
def UCB.nextArm (K : ℕ) [NeZero K] (c : ℝ) (n : ℕ) (h : Hist Unit (Fin K) ℝ n) : Fin K :=
  if n < K then RoundRobin.nextAction K n else
  argmax (fun a ↦ empMean' n h a + ucbWidth' c n h a)

@[fun_prop]
lemma UCB.measurable_nextArm [NeZero K] (c : ℝ) (n : ℕ) : Measurable (nextArm K c n) := sorry

variable (K) in
/-- The UCB algorithm. -/
noncomputable
def ucbAlgorithm [NeZero K] (c : ℝ) : Algorithm Unit (Fin K) ℝ :=
  detAlgorithm (fun n p ↦ UCB.nextArm K c n p.1) (sorry)

end Algorithm
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 : ℕ} {ω : Ω}

/-- If `C ≥ 8 * c * σ2 * log (n + 1) / gap ν a ^ 2`, then arm `a` is never pulled at a time `s < n`
at which it already has more than `C` pulls and the means of the best arm and of `a` lie in
their confidence intervals. -/
lemma sum_indicator_good_event_eq_zero
    (h : IsAlgEnvSeq O A R (ucbAlgorithm K (c * σ2)) (stationaryEnv ν) P)
    (hc : 0 ≤ c) (a : Fin K) (h_gap : 0 < gap ν a) (n C : ℕ)
    (hC : C ≠ 0) (hC' : 8 * c * σ2 * log (n + 1) / gap ν a ^ 2 ≤ C) :
    ∀ᵐ ω ∂P,
    ∑ s ∈ range n, {s | A s ω = a ∧ C < pullCount A a s ω ∧
      (ν (bestArm ν))[id] ≤ empMean A R (bestArm ν) s ω + ucbWidth A (c * σ2) (bestArm ν) s ω ∧
      empMean A R (A s ω) s ω - ucbWidth A (c * σ2) (A s ω) s ω
        ≤ (ν (A s ω))[id]}.indicator 1 s = 0 := sorry

end UCB
end Bandits
end