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.Independence.Integration
import Mathlib.Probability.Kernel.Representation
import Mathlib.Probability.Kernel.Composition.MapComap
import Mathlib.Probability.IdentDistrib
import Mathlib.Probability.Independence.InfinitePi
import Mathlib.MeasureTheory.Function.FactorsThrough

/-! # Standalone extraction for `Bandits.ArrayModel.hasCondDistrib_reward_hist_action_pullCount`
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

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

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

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

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]

/-- The conditional distribution of the reward at time `n + 1`, given the history up to time `n`,
the action at time `n + 1`, and the number of times that action has been pulled before time `n + 1`,
is equal to the kernel `ν`. -/
lemma hasCondDistrib_reward_hist_action_pullCount
    (alg : Algorithm 𝓐 R) (ν : Kernel 𝓐 R) [IsMarkovKernel ν] (n : ℕ) :
    HasCondDistrib (reward alg (n + 1))
      (fun ω ↦ (hist alg ω n, action alg (n + 1) ω,
        pullCount (action alg) (action alg (n + 1) ω) (n + 1) ω))
      ((ν.prodMkRight _).prodMkLeft _) (arrayMeasure ν) := sorry

end Laws
end ArrayModel
end MeasureSpace
end Bandits
end