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.Basic
import Mathlib.Probability.Kernel.Composition.MapComap
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.etcAlgorithm`
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

-- ═══ 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]

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} {P : Measure Ω} [IsProbabilityMeasure P] {A : ℕ → Ω → 𝓐} {R' : ℕ → Ω → R} {a : 𝓐} {m n t : ℕ} {ω : Ω}
section PullCount

/-- 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 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 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.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 : 𝓐} {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.ETC ═══
section
open MeasureTheory ProbabilityTheory Finset Learning
open scoped ENNReal NNReal
namespace Bandits
variable {K : ℕ}
section AlgorithmDefinition

/-- Arm pulled by the ETC algorithm at time `n + 1`.
For `n < K * m - 1`, this is arm `(n + 1) % K`.
For `n = K * m - 1`, this is the arm with the highest empirical mean after the exploration phase.
For `n ≥ K * m`, this is the same arm as at time `n`. -/
noncomputable
def ETC.nextArm (hK : 0 < K) (m n : ℕ) (h : Iic n → Fin K × ℝ) : Fin K :=
  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
  if hn : n < K * m - 1 then RoundRobin.nextAction hK n
  else
    if hn_eq : n = K * m - 1 then argmax (empMean' n h)
    else (h ⟨n, sorry⟩).1

/-- The next arm pulled by ETC is chosen in a measurable way. -/
@[fun_prop]
lemma ETC.measurable_nextArm (hK : 0 < K) (m n : ℕ) : Measurable (nextArm hK m n) := sorry

/-- The Explore-Then-Commit algorithm: deterministic algorithm that chooses the next arm according
to `ETC.nextArm`. -/
noncomputable
def etcAlgorithm (hK : 0 < K) (m : ℕ) : Algorithm (Fin K) ℝ :=
  detAlgorithm (ETC.nextArm hK m) (sorry) ⟨0, hK⟩

end AlgorithmDefinition
end Bandits
end