import Mathlib.Analysis.Normed.Ring.Basic
import Mathlib.MeasureTheory.Constructions.BorelSpace.Basic
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.Order.CompletePartialOrder
import Mathlib.Probability.Martingale.BorelCantelli

/-! # Standalone extraction for `Learning.IsAlgEnvSeq.stronglyAdapted_empMean_add_one`
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 Learning.IsAlgEnvSeq
end Learning.IsAlgEnvSeq

-- ═══ ForMathlib.MeasureTheory.Measurable ═══
section
open Finset
namespace MeasureTheory
variable {α β γ : Type*} {mα : MeasurableSpace α} {mβ : MeasurableSpace β} {mγ : MeasurableSpace γ} {μ : Measure α}

lemma measurable_comp_comap (f : α → β) {g : β → γ} (hg : Measurable g) :
    Measurable[mβ.comap f] (g ∘ f) := sorry

end MeasureTheory
end

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

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

@[fun_prop]
lemma measurable_history (hA : ∀ n, Measurable (A n))
    (hY : ∀ n, Measurable (Y n)) (n : ℕ) :
    Measurable (history A Y n) := sorry

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

section Filtration
namespace IsAlgEnvSeq

/-- Filtration generated by the history up to time `n`. -/
def filtration (h : IsAlgEnvSeq A Y alg env P) :
    Filtration ℕ mΩ where
  seq i := MeasurableSpace.comap (history A Y i) inferInstance
  mono' i j hij := by
    simp only
    rw [← measurable_iff_comap_le]
    have : history A Y i = (fun h k ↦ h ⟨k.1, by grind⟩) ∘ history A Y j := rfl
    rw [this]
    exact measurable_comp_comap _ (by fun_prop)
  le' i := by
    rw [← measurable_iff_comap_le]
    exact Learning.measurable_history h.measurable_action h.measurable_feedback i

end IsAlgEnvSeq
end Filtration
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 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))

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 obtained when pulling action `a` up to time `t` (exclusive). -/
noncomputable def sumRewards (A : ℕ → Ω → 𝓐) (R : ℕ → Ω → 𝓨) (a : 𝓐) (t : ℕ) (ω : Ω) : 𝓨 :=
  ∑ s ∈ range t, if A s ω = a then (R s) ω else 0

/-- Empirical mean reward obtained when pulling action `a` up to time `t` (exclusive). -/
noncomputable
def empMean (A : ℕ → Ω → 𝓐) (R : ℕ → Ω → ℝ) (a : 𝓐) (t : ℕ) (ω : Ω) : ℝ :=
  sumRewards A R a t ω / pullCount A a t ω

variable [MeasurableSingletonClass 𝓐]

lemma IsAlgEnvSeq.stronglyAdapted_empMean_add_one
    {R' : ℕ → Ω → ℝ} {alg : Algorithm 𝓐 ℝ} {env : Environment 𝓐 ℝ}
    (h : IsAlgEnvSeq A R' alg env P) (a : 𝓐) :
    StronglyAdapted h.filtration (fun n ↦ empMean A R' a (n + 1)) := sorry

end Learning
end