import Mathlib.MeasureTheory.Order.Lattice
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

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

-- ═══ ForMathlib.MeasureTheory.Order.Lattice ═══
section
open Finset
variable {α δ : Type*} [MeasurableSpace δ] [SemilatticeInf α] {m : MeasurableSpace α} [MeasurableInf₂ α]
attribute [to_dual existing] MeasurableInf₂
end

-- ═══ SequentialLearning.FiniteActions ═══
section
open MeasureTheory Finset Learning
namespace Learning
variable {𝓐 R Ω : Type*} {m𝓐 : MeasurableSpace 𝓐} {mR : MeasurableSpace R} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] {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
section StepsUntil

/-- Number of steps until action `a` was pulled exactly `m` times. -/
noncomputable
def stepsUntil (A : ℕ → Ω → 𝓐) (a : 𝓐) (m : ℕ) (ω : Ω) : ℕ∞ :=
  sInf ((↑) '' {s | pullCount A a (s + 1) ω = m})

end StepsUntil
section RewardByCount

/-- Reward obtained when pulling action `a` for the `m`-th time.
If it is never pulled `m` times, the reward is given by the second component of `ω`, which in
applications will be indepedent with same law. -/
noncomputable
def rewardByCount (A : ℕ → Ω → 𝓐) (R' : ℕ → Ω → R) (a : 𝓐) (m : ℕ) (ω : Ω × (ℕ → 𝓐 → R)) : R :=
  match (stepsUntil A a m ω.1) with
  | ⊤ => ω.2 m a
  | (n : ℕ) => R' n ω.1

variable {ω : Ω × (ℕ → 𝓐 → R)}

lemma rewardByCount_eq_ite (a : 𝓐) (m : ℕ) (ω : Ω × (ℕ → 𝓐 → R)) :
    rewardByCount A R' a m ω =
      if (stepsUntil A a m ω.1) = ⊤ then ω.2 m a else R' (stepsUntil A a m ω.1).toNat ω.1 := sorry

end RewardByCount
end Learning
end