import Mathlib.MeasureTheory.Order.Lattice
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.Order.CompletePartialOrder
import Mathlib.Probability.Martingale.BorelCantelli
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

/-! # Standalone extraction for `Bandits.ArrayModel.isAlgEnvSeq_arrayMeasure`
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 MeasureTheory
end MeasureTheory
namespace ProbabilityTheory
end ProbabilityTheory
namespace ENNReal
end ENNReal
namespace Learning
end Learning
namespace Bandits
end Bandits
namespace Bandits.ArrayModel
end Bandits.ArrayModel

-- ═══ ForMathlib.MeasureTheory.Order.Lattice ═══
section
open Finset
variable {α δ : Type*} [MeasurableSpace δ] [SemilatticeInf α] {m : MeasurableSpace α} [MeasurableInf₂ α]
attribute [to_dual existing] MeasurableInf₂
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

/-- 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 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.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.ArrayProbSpace ═══
section
open MeasureTheory ProbabilityTheory Filter Real Finset Learning
open scoped ENNReal NNReal
namespace Bandits
variable {𝓐 R : Type*} {m𝓐 : MeasurableSpace 𝓐} {mR : MeasurableSpace R}
section MeasureSpace

/-- Measure of an infinite stream of rewards from each action. -/
noncomputable
def streamMeasure (ν : Kernel 𝓐 R) : Measure (ℕ → 𝓐 → R) :=
  Measure.infinitePi fun _ ↦ Measure.infinitePi ν

instance (ν : Kernel 𝓐 R) [IsMarkovKernel ν] : IsProbabilityMeasure (streamMeasure ν) := by
  unfold streamMeasure
  infer_instance

namespace ArrayModel
open unitInterval
section ProbabilitySpace

variable (𝓐 R) in
/-- Probability space for the array model of stochastic bandits. -/
def probSpace : Type _ := (ℕ → I) × (ℕ → 𝓐 → R)

instance {𝓐 R : Type*} [MeasurableSpace R] : MeasurableSpace (probSpace 𝓐 R) :=
  inferInstanceAs (MeasurableSpace ((ℕ → I) × (ℕ → 𝓐 → R)))

/-- Probability measure for the array model of stochastic bandits. -/
noncomputable
def arrayMeasure (ν : Kernel 𝓐 R) : Measure (probSpace 𝓐 R) :=
  (Measure.infinitePi fun _ ↦ volume).prod (streamMeasure ν)

instance (ν : Kernel 𝓐 R) [IsMarkovKernel ν] : IsProbabilityMeasure (arrayMeasure ν) :=
  Measure.prod.instIsProbabilityMeasure _ _

variable [Nonempty 𝓐] [StandardBorelSpace 𝓐]

/-- The initial action is the image of a uniform random variable by this function. -/
noncomputable
def initAlgFunction (alg : Algorithm 𝓐 R) : I → 𝓐 :=
  (Measure.exists_measurable_map_eq alg.p0).choose

/-- The next action is the image of the history and a uniform random variable by this function. -/
noncomputable
def algFunction (alg : Algorithm 𝓐 R) (n : ℕ) :
    (Iic n → 𝓐 × R) → I → 𝓐 :=
  (Kernel.exists_measurable_map_eq_unitInterval (alg.policy n)).choose

end ProbabilitySpace
variable [Nonempty 𝓐] [StandardBorelSpace 𝓐]
section HistoryActionReward

/-- History of actions and rewards up to time `n` in the array model. -/
noncomputable
def hist [DecidableEq 𝓐] (alg : Algorithm 𝓐 R) (ω : probSpace 𝓐 R) : (n : ℕ) → Iic n → 𝓐 × R
| 0 => fun _ ↦ (initAlgFunction alg (ω.1 0), ω.2 0 (initAlgFunction alg (ω.1 0)))
| n + 1 =>
  let hn : Iic n → 𝓐 × R := hist alg ω n
  let a : 𝓐 := algFunction alg n hn (ω.1 (n + 1))
  fun i ↦ if hin : i ≤ n then hn ⟨i, sorry⟩ else (a, ω.2 (pullCount' n hn a) a)

/-- Action taken at time `n` in the array model. -/
noncomputable
def action [DecidableEq 𝓐] (alg : Algorithm 𝓐 R) (n : ℕ) (ω : probSpace 𝓐 R) : 𝓐 :=
  (hist alg ω n ⟨n, sorry⟩).1

/-- Reward received at time `n` in the array model. -/
noncomputable
def reward [DecidableEq 𝓐] (alg : Algorithm 𝓐 R) (n : ℕ) (ω : probSpace 𝓐 R) : R :=
  (hist alg ω n ⟨n, sorry⟩).2

end HistoryActionReward
variable [DecidableEq 𝓐]
section Laws
variable [Countable 𝓐]
variable [StandardBorelSpace R] [Nonempty R]

lemma isAlgEnvSeq_arrayMeasure (alg : Algorithm 𝓐 R) (ν : Kernel 𝓐 R) [IsMarkovKernel ν] :
    IsAlgEnvSeq (action alg) (reward alg) alg (stationaryEnv ν) (arrayMeasure ν) := sorry

end Laws
end ArrayModel
end MeasureSpace
end Bandits
end