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.IdentDistrib
import Mathlib.Probability.Independence.InfinitePi
import Mathlib.MeasureTheory.Measure.AbsolutelyContinuous
import Mathlib.MeasureTheory.OuterMeasure.Basic
import Mathlib.Probability.Kernel.CompProdEqIff
import Mathlib.Probability.Kernel.Composition.MeasureComp
import Mathlib.Probability.Kernel.Composition.MeasureCompProd

/-! # Standalone extraction for `Bandits.TS.hasCondDistrib_action`
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 Learning.IsBayesAlgEnvSeq
end Learning.IsBayesAlgEnvSeq
namespace Learning.IT
end Learning.IT
namespace Learning.Algorithm
end Learning.Algorithm
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]

instance (alg : Algorithm 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (alg.policy n) := alg.h_policy n

instance (alg : Algorithm 𝓐 𝓨) : IsProbabilityMeasure alg.p0 := alg.hp0

/-- An algorithm with observations in `𝓧 × 𝓨` obtained from an algorithm with observations in `𝓨`
by ignoring the `𝓧` component of each observation. -/
@[simps]
def Algorithm.prodLeft (𝓧 : Type*) [MeasurableSpace 𝓧] (alg : Algorithm 𝓐 𝓨) :
    Algorithm 𝓐 (𝓧 × 𝓨) where
  policy n := (alg.policy n).comap (fun h i ↦ ((h i).1, (h i).2.2)) (sorry)
  p0 := alg.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]

instance (env : Environment 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (env.feedback n) := env.h_feedback n

instance (env : Environment 𝓐 𝓨) : IsMarkovKernel env.ν0 := env.hp0

/-- Kernel describing the distribution of the next action-feedback pair given the history
up to `n`. -/
noncomputable
def stepKernel (alg : Algorithm 𝓐 𝓨) (env : Environment 𝓐 𝓨) (n : ℕ) :
    Kernel (Iic n → 𝓐 × 𝓨) (𝓐 × 𝓨) :=
  alg.policy n ⊗ₖ env.feedback n

instance (alg : Algorithm 𝓐 𝓨) (env : Environment 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (stepKernel alg env n) := sorry

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

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

-- ═══ SequentialLearning.IonescuTulceaSpace ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset
open scoped ENNReal NNReal
namespace Learning
variable {𝓐 𝓨 Ω : Type*} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨} {mΩ : MeasurableSpace Ω}

/-- Measure on the sequence of actions and observations generated by the algorithm/environment. -/
noncomputable
def trajMeasure (alg : Algorithm 𝓐 𝓨) (env : Environment 𝓐 𝓨) :
    Measure (ℕ → 𝓐 × 𝓨) :=
  Kernel.trajMeasure (alg.p0 ⊗ₘ env.ν0) (stepKernel alg env)

instance (alg : Algorithm 𝓐 𝓨) (env : Environment 𝓐 𝓨) : IsProbabilityMeasure (trajMeasure alg env) := sorry

namespace IT

/-- `action n` is the action pulled at time `n`. This is a random variable on the measurable space
`ℕ → 𝓐 × 𝓨`. -/
def action (n : ℕ) (h : ℕ → 𝓐 × 𝓨) : 𝓐 := (h n).1

end IT
end Learning
end

-- ═══ SequentialLearning.BayesStationaryEnv ═══
section
open MeasureTheory ProbabilityTheory Finset
open scoped ENNReal NNReal
namespace Learning
variable {𝓔 𝓐 𝓨 Ω : Type*}
variable [MeasurableSpace 𝓔] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨] [MeasurableSpace Ω]

/-- `IsBayesAlgEnvSeq Q κ alg E A Y P` states that there is a measure `P : Measure Ω` such
  that the parameter `E : Ω → 𝓔` has law `Q` and that the sequences of actions `A : ℕ → Ω → 𝓐`
  and feedbacks `Y : ℕ → Ω → 𝓨` are generated by the algorithm `alg : Algorithm 𝓐 𝓨` interacting
  with an underlying environment that depends on `E` and `κ` (`stationaryEnv (κ.sectR (E ω))`). -/
structure IsBayesAlgEnvSeq
    (Q : Measure 𝓔) (κ : Kernel (𝓔 × 𝓐) 𝓨) (alg : Algorithm 𝓐 𝓨)
    (E : Ω → 𝓔) (A : ℕ → Ω → 𝓐) (Y : ℕ → Ω → 𝓨)
    (P : Measure Ω) [IsFiniteMeasure P] : Prop where
  measurable_param : Measurable E := sorry
  measurable_action n : Measurable (A n) := sorry
  measurable_feedback n : Measurable (Y n) := sorry
  hasLaw_env : HasLaw E Q P
  hasCondDistrib_action_zero : HasCondDistrib (A 0) E (Kernel.const _ alg.p0) P
  hasCondDistrib_feedback_zero : HasCondDistrib (Y 0) (fun ω ↦ (E ω, A 0 ω)) κ P
  hasCondDistrib_action n :
    HasCondDistrib (A (n + 1)) (fun ω ↦ (E ω, history A Y n ω))
      ((alg.policy n).prodMkLeft _) P
  hasCondDistrib_feedback n :
    HasCondDistrib (Y (n + 1)) (fun ω ↦ (history A Y n ω, E ω, A (n + 1) ω))
      (κ.prodMkLeft _) P

section IsAlgEnvSeq

/-- An environment with observations in `𝓔 × 𝓨`. The first element `e` of an observation is
sampled from `Q` once and remains constant. The second element of an observation is sampled from
`κ (e, a)`, where `a` is the corresponding action. -/
noncomputable
def bayesStationaryEnv (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨)
    [IsMarkovKernel κ] : Environment 𝓐 (𝓔 × 𝓨) where
  feedback n :=
    let g : (Iic n → 𝓐 × 𝓔 × 𝓨) × 𝓐 → 𝓔 × 𝓐 := fun (h, a) => ((h ⟨0, sorry⟩).2.1, a)
    (Kernel.deterministic (Prod.fst ∘ g) (sorry)) ×ₖ (κ.comap g (sorry))
  ν0 := (Kernel.const _ Q) ⊗ₖ κ.swapLeft

variable {Q : Measure 𝓔} [IsProbabilityMeasure Q] {κ : Kernel (𝓔 × 𝓐) 𝓨} [IsMarkovKernel κ]
variable {alg : Algorithm 𝓐 𝓨} {A : ℕ → Ω → 𝓐} {Y : ℕ → Ω → 𝓔 × 𝓨}
variable {P : Measure Ω} [IsProbabilityMeasure P]
end IsAlgEnvSeq
namespace IT

/-- A measure `P` on a measurable space that carries random variables `E`, `A`, and `Y` such that
`IsBayesAlgEnvSeq Q κ alg E A Y P`. -/
noncomputable
def bayesTrajMeasure (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨)
    [IsMarkovKernel κ] (alg : Algorithm 𝓐 𝓨) : Measure (ℕ → 𝓐 × 𝓔 × 𝓨) :=
  trajMeasure (alg.prodLeft 𝓔) (bayesStationaryEnv Q κ)

instance (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ] (alg : Algorithm 𝓐 𝓨) : IsProbabilityMeasure (bayesTrajMeasure Q κ alg) := sorry

/-- A kernel that represents the posterior over `E` given the history up to time `n`. -/
noncomputable
def bayesTrajMeasurePosterior [StandardBorelSpace 𝓔] [Nonempty 𝓔]
    (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ]
    (alg : Algorithm 𝓐 𝓨) (n : ℕ) : Kernel (Iic n → 𝓐 × 𝓨) 𝓔 :=
  condDistrib (fun ω ↦ (ω 0).2.1) (history action (fun n ω ↦ (ω n).2.2) n)
    (bayesTrajMeasure Q κ alg)

instance [StandardBorelSpace 𝓔] [Nonempty 𝓔] (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × 𝓐) 𝓨) [IsMarkovKernel κ] (alg : Algorithm 𝓐 𝓨) (n : ℕ) : IsMarkovKernel (bayesTrajMeasurePosterior Q κ alg n) := sorry

end IT
end Learning
end

-- ═══ SequentialLearning.Algorithms.RandomSampling.Basic ═══
section
open MeasureTheory ProbabilityTheory Learning Finset ENNReal Filter
open scoped Topology
namespace Learning
variable {𝓐 𝓨 Ω : Type*} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨} {mΩ : MeasurableSpace Ω} {μ : Measure 𝓐} [IsProbabilityMeasure μ] {P : Measure Ω} [IsProbabilityMeasure P]

open Set in
/-- The _Random Sampling_ algorithm, which samples from a fixed probability
measure at each iteration. -/
@[simps]
noncomputable def randomSampling (μ : Measure 𝓐) [IsProbabilityMeasure μ] : Algorithm 𝓐 𝓨 where
  policy _ := Kernel.const _ μ
  p0 := μ

end Learning
end

-- ═══ SequentialLearning.Algorithms.Uniform ═══
section
open MeasureTheory ProbabilityTheory Learning
open scoped Algorithm
namespace Learning
variable {𝓐 𝓨 : Type*} {m𝓐 : MeasurableSpace 𝓐} {m𝓨 : MeasurableSpace 𝓨}

/-- The Uniform algorithm: actions are chosen uniformly at random. -/
noncomputable
def uniformAlgorithm [Finite 𝓐] [Nonempty 𝓐] : Algorithm 𝓐 𝓨 := randomSampling (uniformOn Set.univ)

end Learning
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 Thompson sampling policy samples an action according to its probability of being optimal
under the posterior over environments given the history so far.
The posterior under a uniform algorithm is used to avoid a circular definition. -/
noncomputable
def TS.policy (hK : 0 < K) (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × Fin K) ℝ)
    [IsMarkovKernel κ] (n : ℕ) : Kernel (Iic n → (Fin K) × ℝ) (Fin K) :=
  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
  (IT.bayesTrajMeasurePosterior Q κ uniformAlgorithm n).map (bestAction κ id)

instance {hK : 0 < K} {Q : Measure 𝓔} [IsProbabilityMeasure Q] {κ : Kernel (𝓔 × Fin K) ℝ}
    [IsMarkovKernel κ] {n : ℕ} : IsMarkovKernel (TS.policy hK Q κ n) :=
  Kernel.IsMarkovKernel.map _ (sorry)

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

/-- The Thompson sampling algorithm with actions in `Fin K`, where `Q : Measure 𝓔` is a prior
  distribution over parameters, and `κ : Kernel (𝓔 × Fin K) ℝ` is a Markov kernel that defines the
  stationary environment `stationaryEnv (κ.sectR e)` that corresponds to a parameter `e : 𝓔`.

  At every time `n`, the Thompson sampling policy uses the posterior over the parameters given the
  history up to time `n` to derive the probability of each action being optimal. The action for time
  `n` is sampled according to these probabilities. -/
noncomputable
def tsAlgorithm (hK : 0 < K) (Q : Measure 𝓔) [IsProbabilityMeasure Q] (κ : Kernel (𝓔 × Fin K) ℝ)
    [IsMarkovKernel κ] : Algorithm (Fin K) ℝ where
  policy := TS.policy hK Q κ
  p0 := TS.initialPolicy hK Q κ

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

/-- If Thompson sampling has the correct prior over environments, then the conditional distribution
of the next action given the history so far is equal to the conditional distribution of the best
action given the history so far. -/
lemma TS.hasCondDistrib_action (hK : 0 < K) (h : IsBayesAlgEnvSeq Q κ (tsAlgorithm hK Q κ) E A R P)
    (n : ℕ) :
    HasCondDistrib (A (n + 1)) (history A R n)
      (condDistrib (bestAction κ E) (history A R n) P) P := sorry

end Bandits
end