import Mathlib.MeasureTheory.Order.Lattice
import Mathlib.Probability.Kernel.Basic
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.MeasureTheory.MeasurableSpace.Embedding
import Mathlib.Order.Restriction
import Mathlib.Probability.Kernel.IonescuTulcea.Maps
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.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.preimage_action_pullCount_eq`
Definitions are copied verbatim; theorem proofs are replaced by `sorry`.
Auto-generated by ChallengeGen. -/

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 𝓨} {mΩ : MeasurableSpace Ω}

/-- One round of interaction: an observation, then an action, then a feedback. -/
abbrev Round (𝓞 𝓐 𝓨 : Type*) := 𝓞 × 𝓐 × 𝓨

/-- The action of a round. -/
def Round.action (r : Round 𝓞 𝓐 𝓨) : 𝓐 := r.2.1

/-- History of `n` complete rounds; `n = 0` is the empty history. -/
abbrev Hist (𝓞 𝓐 𝓨 : Type*) (n : ℕ) := Fin n → Round 𝓞 𝓐 𝓨

/-- A stochastic, sequential algorithm.
At each round, it sees an observation in `𝓞`, then takes an action in `𝓐`, and finally receives
feedback in `𝓨`. The action is a random function of the past rounds and the current observation. -/
@[ext]
structure Algorithm (𝓞 𝓐 𝓨 : Type*) [MeasurableSpace 𝓞] [MeasurableSpace 𝓐] [MeasurableSpace 𝓨]
    where
  /-- Law of the action of round `n` given the past rounds and the current observation. -/
  policy : (n : ℕ) → Kernel (Hist 𝓞 𝓐 𝓨 n × 𝓞) 𝓐
  /-- The policy is a Markov kernel. -/
  [isMarkovKernel_policy : ∀ n, IsMarkovKernel (policy n)]

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

end Learning
end

-- ═══ SequentialLearning.FiniteActions ═══
section
open MeasureTheory Finset Learning
namespace Learning
variable {𝓞 𝓐 R Ω : Type*} {m𝓞 : MeasurableSpace 𝓞} {m𝓐 : MeasurableSpace 𝓐} {mR : MeasurableSpace R} {mΩ : MeasurableSpace Ω} [DecidableEq 𝓐] {alg : Algorithm 𝓞 𝓐 R} {P : Measure Ω} [IsProbabilityMeasure P] {O : ℕ → Ω → 𝓞} {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` in the history before time `n`.
This is the number of entries in `h` in which the arm is `a`. -/
noncomputable
def pullCount' (n : ℕ) (h : Hist 𝓞 𝓐 R n) (a : 𝓐) := #{s | (h s).action = a}

end PullCount
end Learning
end

-- ═══ Online.Bandit.ArrayProbSpace ═══
section
open MeasureTheory ProbabilityTheory Filter Finset Learning
open scoped ENNReal NNReal
namespace Bandits
variable {𝓐 𝓡 : Type*} {m𝓐 : MeasurableSpace 𝓐} {m𝓡 : MeasurableSpace 𝓡}
section MeasureSpace
namespace ArrayModel
open unitInterval
section ProbabilitySpace

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

variable [Nonempty 𝓐] [StandardBorelSpace 𝓐]

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

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

/-- History of actions and rewards before time `n` in the array model. -/
noncomputable
def hist [DecidableEq 𝓐] (alg : Algorithm Unit 𝓐 𝓡) (ω : probSpace 𝓐 𝓡) :
    (n : ℕ) → Hist Unit 𝓐 𝓡 n
| 0 => default
| n + 1 =>
  let hn : Hist Unit 𝓐 𝓡 n := hist alg ω n
  let a : 𝓐 := algFunction alg n hn (ω.1 n)
  Fin.snoc hn ((), a, ω.2 (pullCount' n hn a) a)

/-- Action taken at time `n` in the array model. -/
noncomputable
def action [DecidableEq 𝓐] (alg : Algorithm Unit 𝓐 𝓡) (n : ℕ) (ω : probSpace 𝓐 𝓡) : 𝓐 :=
  algFunction alg n (hist alg ω n) (ω.1 n)

end HistoryActionReward
variable [DecidableEq 𝓐]
section MeasurabilityAdvanced

lemma preimage_action_pullCount_eq (alg : Algorithm Unit 𝓐 𝓡) (a : 𝓐) (n m : ℕ) :
    (fun ω ↦ (action alg n ω, pullCount (action alg) (action alg n ω) n ω)) ⁻¹' {(a, m)} =
      {ω | action alg n ω = a ∧ pullCount (action alg) a n ω = m}.indicator (fun _ ↦ 1)
        ⁻¹' {1} := sorry

end MeasurabilityAdvanced
end ArrayModel
end MeasureSpace
end Bandits
end