import Mathlib.MeasureTheory.Order.Lattice
import Mathlib.CategoryTheory.Countable
import Mathlib.MeasureTheory.Constructions.Polish.Basic
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.MeasureTheory.Measure.AbsolutelyContinuous
import Mathlib.MeasureTheory.OuterMeasure.Basic
import Mathlib.Probability.IdentDistrib
import Mathlib.Probability.Independence.InfinitePi
import Mathlib.Probability.Kernel.CompProdEqIff
import Mathlib.Probability.Kernel.Composition.MeasureComp
import Mathlib.Probability.Kernel.Composition.MeasureCompProd
import Mathlib.Probability.Kernel.Composition.MapComap

/-! # Standalone extraction for `Bandits.instIsProbabilityMeasureFinInitialPolicy`
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 Learning.IsBayesAlgEnvSeq
end Learning.IsBayesAlgEnvSeq
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

-- ═══ Online.Bandit.BayesRegret ═══
section
open MeasureTheory ProbabilityTheory Finset
namespace Learning.IsBayesAlgEnvSeq
variable {𝓔 𝓐 𝓨 Ω : Type*}
variable [MeasurableSpace 𝓔] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨] [MeasurableSpace Ω]

/-- A random variable that gives the mean feedback of action `a`. -/
noncomputable
def actionMean (κ : Kernel (𝓔 × 𝓐) ℝ) (E : Ω → 𝓔) (a : 𝓐) (ω : Ω) : ℝ := (κ (E ω, a))[id]

/-- A random variable that gives the action with the highest mean feedback. -/
noncomputable
def bestAction [Nonempty 𝓐] [Fintype 𝓐] (κ : Kernel (𝓔 × 𝓐) ℝ) (E : Ω → 𝓔) (ω : Ω) : 𝓐 :=
  argmax (fun a ↦ actionMean κ E a ω)

@[fun_prop]
lemma measurable_bestAction [Nonempty 𝓐] [Fintype 𝓐] {κ : Kernel (𝓔 × 𝓐) ℝ} {E : Ω → 𝓔}
    (hE : Measurable E) : Measurable (bestAction κ E) := sorry

end Learning.IsBayesAlgEnvSeq
end

-- ═══ Online.Bandit.Algorithms.TS ═══
section
open MeasureTheory ProbabilityTheory Finset Learning
open IsBayesAlgEnvSeq (bestAction)
namespace Bandits
section Algorithm
variable {K : ℕ}
variable {𝓔 : Type*} [MeasurableSpace 𝓔] [StandardBorelSpace 𝓔] [Nonempty 𝓔]

/-- The initial action is sampled according to its probability of being optimal under the prior over
environments. -/
noncomputable
def TS.initialPolicy (hK : 0 < K) (Q : Measure 𝓔) (κ : Kernel (𝓔 × Fin K) ℝ) : Measure (Fin K) :=
  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
  Q.map (bestAction κ id)

instance {hK : 0 < K} {Q : Measure 𝓔} [IsProbabilityMeasure Q] {κ : Kernel (𝓔 × Fin K) ℝ} :
    IsProbabilityMeasure (TS.initialPolicy hK Q κ) :=
  Measure.isProbabilityMeasure_map (sorry)

end Algorithm
variable {K : ℕ}
variable {Ω : Type*} [MeasurableSpace Ω]
variable {𝓔 : Type*} [MeasurableSpace 𝓔] [StandardBorelSpace 𝓔] [Nonempty 𝓔]
variable {E : Ω → 𝓔} {A : ℕ → Ω → Fin K} {R : ℕ → Ω → ℝ}
variable {Q : Measure 𝓔} [IsProbabilityMeasure Q] [IsMarkovKernel κ]
variable {P : Measure Ω} [IsProbabilityMeasure P]
end Bandits
end