import Mathlib.MeasureTheory.Order.Lattice
import Mathlib.CategoryTheory.Countable
import Mathlib.MeasureTheory.Constructions.Polish.Basic
import Mathlib.Order.CompletePartialOrder
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.Martingale.BorelCantelli
import Mathlib.Probability.Kernel.Composition.MapComap
import Mathlib.Probability.Kernel.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

/-! # Standalone extraction for `Bandits.UCB.some_sum_eq_zero`
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 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 Ω}

/-- 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))

/-- Number of pulls of arm `a` up to (and including) time `n`.
This is the number of entries in `h` in which the arm is `a`. -/
noncomputable
def pullCount' (n : ℕ) (h : Iic n → 𝓐 × R) (a : 𝓐) := #{s | (h s).1 = a}

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

/-- Sum of rewards of arm `a` up to (and including) time `n`. -/
noncomputable
def sumRewards' (n : ℕ) (h : Iic n → 𝓐 × 𝓨) (a : 𝓐) :=
  ∑ s, if (h s).1 = a then (h s).2 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` at time `n`. -/
noncomputable
def empMean' (n : ℕ) (h : Iic n → 𝓐 × ℝ) (a : 𝓐) :=
  sumRewards' n h a / pullCount' n h a

variable [MeasurableSingletonClass 𝓐]
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.Regret ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset Learning
open scoped ENNReal NNReal
namespace Bandits
variable {𝓐 Ω : Type*} [DecidableEq 𝓐] {m𝓐 : MeasurableSpace 𝓐} {mΩ : MeasurableSpace Ω} {ν : Kernel 𝓐 ℝ} {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 𝓨}
variable {nextA : (n : ℕ) → (Iic n → 𝓐 × 𝓨) → 𝓐} {h_next : ∀ n, Measurable (nextA n)} {action0 : 𝓐} {env : Environment 𝓐 𝓨} {f0 : 𝓐 → 𝓨} {hf0 : Measurable f0} {f : (n : ℕ) → ((Iic n → 𝓐 × 𝓨) × 𝓐) → 𝓨} {hf : ∀ n, Measurable (f n)}

/-- A deterministic algorithm, which chooses the action given by the function `nextAction`. -/
@[simps]
noncomputable
def detAlgorithm (nextA : (n : ℕ) → (Iic n → 𝓐 × 𝓨) → 𝓐)
    (h_next : ∀ n, Measurable (nextA n)) (action0 : 𝓐) :
    Algorithm 𝓐 𝓨 where
  policy n := Kernel.deterministic (nextA n) (h_next n)
  p0 := Measure.dirac action0

end Learning
end

-- ═══ SequentialLearning.Algorithms.RoundRobin ═══
section
open MeasureTheory ProbabilityTheory Finset Learning
open scoped ENNReal NNReal
namespace Learning
variable {𝓨 : Type*} {m𝓨 : MeasurableSpace 𝓨} {K : ℕ}
section AlgorithmDefinition

/-- Action chosen by the Round-Robin algorithm at time `n + 1`. This is action `(n + 1) % K`. -/
noncomputable
def RoundRobin.nextAction (hK : 0 < K) (n : ℕ) : Fin K := ⟨(n + 1) % K, Nat.mod_lt _ hK⟩

end AlgorithmDefinition
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 ucbWidth' (c : ℝ) (n : ℕ) (h : Iic n → Fin K × ℝ) (a : Fin K) : ℝ :=
  √(2 * c * log (n + 2) / pullCount' n h a)

open Classical in
/-- Arm pulled by the UCB algorithm at time `n + 1`. -/
noncomputable
def UCB.nextArm (hK : 0 < K) (c : ℝ) (n : ℕ) (h : Iic n → Fin K × ℝ) : Fin K :=
  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
  if n < K - 1 then RoundRobin.nextAction hK n else
  argmax (fun a ↦ empMean' n h a + ucbWidth' c n h a)

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

/-- The UCB algorithm. -/
noncomputable
def ucbAlgorithm (hK : 0 < K) (c : ℝ) : Algorithm (Fin K) ℝ :=
  detAlgorithm (UCB.nextArm hK c) (sorry) ⟨0, hK⟩

end Algorithm
namespace UCB
variable {hK : 0 < K} {c : ℝ} {ν : Kernel (Fin K) ℝ} [IsMarkovKernel ν] {Ω : Type*} {mΩ : MeasurableSpace Ω} {P : Measure Ω} [IsProbabilityMeasure P] {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

lemma some_sum_eq_zero [Nonempty (Fin K)]
    (h : IsAlgEnvSeq A R (ucbAlgorithm hK (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