From 0eb829a7eef0c0bd1833d6360738685eac5c4e78 Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Fri, 24 Oct 2025 17:12:56 +0200
Subject: [PATCH 1/6] delete a lemma

---
 LeanBandits/RewardByCountMeasure.lean | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/LeanBandits/RewardByCountMeasure.lean b/LeanBandits/RewardByCountMeasure.lean
index 60df2295..4fa6f6b2 100644
--- a/LeanBandits/RewardByCountMeasure.lean
+++ b/LeanBandits/RewardByCountMeasure.lean
@@ -350,22 +350,18 @@ lemma indepFun_rewardByCount_Iic (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [
       fun ω (i : Iic n) ↦ rewardByCount a i ω.1 ω.2 := by
   sorry
 
-lemma iIndepFun_rewardByCount' (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [IsMarkovKernel ν] (a : α) :
+lemma iIndepFun_rewardByCount (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [IsMarkovKernel ν] (a : α) :
     iIndepFun (fun n ω ↦ rewardByCount a n ω.1 ω.2) (Bandit.measure alg ν) := by
   rw [iIndepFun_nat_iff_forall_indepFun (by fun_prop)]
   exact indepFun_rewardByCount_Iic alg ν a
 
-lemma iIndepFun_rewardByCount (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [IsMarkovKernel ν] :
-    iIndepFun (fun (p : α × ℕ) ω ↦ rewardByCount p.1 p.2 ω.1 ω.2) (Bandit.measure alg ν) := by
-  sorry
-
 lemma identDistrib_rewardByCount_stream' [Countable α] [StandardBorelSpace α] [Nonempty α]
     (a : α) :
     IdentDistrib (fun ω n ↦ rewardByCount a (n + 1) ω.1 ω.2) (fun ω n ↦ ω n a)
       (Bandit.measure alg ν) (Bandit.streamMeasure ν) := by
   refine IdentDistrib.pi (fun n ↦ ?_) ?_ ?_
   · refine identDistrib_rewardByCount_eval a (n + 1) n (by simp) (ν := ν)
-  · have h_indep := iIndepFun_rewardByCount' alg ν a
+  · have h_indep := iIndepFun_rewardByCount alg ν a
     exact iIndepFun.precomp (g := fun n ↦ n + 1) (fun i j hij ↦ by grind) h_indep
   · exact iIndepFun_eval_streamMeasure'' ν a
 

From 32e5d30a23f2fae0f383944990ea00dc0ba7acbd Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Fri, 24 Oct 2025 18:32:01 +0200
Subject: [PATCH 2/6] undo

---
 LeanBandits/RewardByCountMeasure.lean | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/LeanBandits/RewardByCountMeasure.lean b/LeanBandits/RewardByCountMeasure.lean
index 4fa6f6b2..60df2295 100644
--- a/LeanBandits/RewardByCountMeasure.lean
+++ b/LeanBandits/RewardByCountMeasure.lean
@@ -350,18 +350,22 @@ lemma indepFun_rewardByCount_Iic (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [
       fun ω (i : Iic n) ↦ rewardByCount a i ω.1 ω.2 := by
   sorry
 
-lemma iIndepFun_rewardByCount (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [IsMarkovKernel ν] (a : α) :
+lemma iIndepFun_rewardByCount' (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [IsMarkovKernel ν] (a : α) :
     iIndepFun (fun n ω ↦ rewardByCount a n ω.1 ω.2) (Bandit.measure alg ν) := by
   rw [iIndepFun_nat_iff_forall_indepFun (by fun_prop)]
   exact indepFun_rewardByCount_Iic alg ν a
 
+lemma iIndepFun_rewardByCount (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [IsMarkovKernel ν] :
+    iIndepFun (fun (p : α × ℕ) ω ↦ rewardByCount p.1 p.2 ω.1 ω.2) (Bandit.measure alg ν) := by
+  sorry
+
 lemma identDistrib_rewardByCount_stream' [Countable α] [StandardBorelSpace α] [Nonempty α]
     (a : α) :
     IdentDistrib (fun ω n ↦ rewardByCount a (n + 1) ω.1 ω.2) (fun ω n ↦ ω n a)
       (Bandit.measure alg ν) (Bandit.streamMeasure ν) := by
   refine IdentDistrib.pi (fun n ↦ ?_) ?_ ?_
   · refine identDistrib_rewardByCount_eval a (n + 1) n (by simp) (ν := ν)
-  · have h_indep := iIndepFun_rewardByCount alg ν a
+  · have h_indep := iIndepFun_rewardByCount' alg ν a
     exact iIndepFun.precomp (g := fun n ↦ n + 1) (fun i j hij ↦ by grind) h_indep
   · exact iIndepFun_eval_streamMeasure'' ν a
 

From c26f63b1eac4efbcdbeeeb6aff47a994b1a8b6f7 Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Fri, 24 Oct 2025 18:57:04 +0200
Subject: [PATCH 3/6] create folders

---
 LeanBandits.lean                                    | 10 +++++-----
 LeanBandits/AlgorithmAndRandomVariables.lean        |  2 +-
 LeanBandits/AlgorithmBuilding.lean                  |  2 +-
 LeanBandits/{ => Bandit}/Bandit.lean                |  2 +-
 LeanBandits/{ => Bandit}/Regret.lean                |  2 +-
 LeanBandits/{ => BanditAlgorithms}/ETC.lean         |  2 +-
 LeanBandits/{ => BanditAlgorithms}/UCB.lean         |  2 +-
 LeanBandits/RewardByCountMeasure.lean               |  4 ++--
 LeanBandits/{ => SequentialLearning}/Algorithm.lean |  0
 9 files changed, 13 insertions(+), 13 deletions(-)
 rename LeanBandits/{ => Bandit}/Bandit.lean (99%)
 rename LeanBandits/{ => Bandit}/Regret.lean (99%)
 rename LeanBandits/{ => BanditAlgorithms}/ETC.lean (99%)
 rename LeanBandits/{ => BanditAlgorithms}/UCB.lean (98%)
 rename LeanBandits/{ => SequentialLearning}/Algorithm.lean (100%)

diff --git a/LeanBandits.lean b/LeanBandits.lean
index aca14ced..9d11fbcd 100644
--- a/LeanBandits.lean
+++ b/LeanBandits.lean
@@ -1,8 +1,9 @@
-import LeanBandits.Algorithm
 import LeanBandits.AlgorithmAndRandomVariables
 import LeanBandits.AlgorithmBuilding
-import LeanBandits.Bandit
-import LeanBandits.ETC
+import LeanBandits.Bandit.Bandit
+import LeanBandits.Bandit.Regret
+import LeanBandits.BanditAlgorithms.ETC
+import LeanBandits.BanditAlgorithms.UCB
 import LeanBandits.ForMathlib.CondDistrib
 import LeanBandits.ForMathlib.IdentDistrib
 import LeanBandits.ForMathlib.IndepFun
@@ -10,6 +11,5 @@ import LeanBandits.ForMathlib.KernelSub
 import LeanBandits.ForMathlib.Measurable
 import LeanBandits.ForMathlib.SubGaussian
 import LeanBandits.ForMathlib.Traj
-import LeanBandits.Regret
 import LeanBandits.RewardByCountMeasure
-import LeanBandits.UCB
+import LeanBandits.SequentialLearning.Algorithm
diff --git a/LeanBandits/AlgorithmAndRandomVariables.lean b/LeanBandits/AlgorithmAndRandomVariables.lean
index 03c20be2..04c305ff 100644
--- a/LeanBandits/AlgorithmAndRandomVariables.lean
+++ b/LeanBandits/AlgorithmAndRandomVariables.lean
@@ -3,7 +3,7 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
-import LeanBandits.Regret
+import LeanBandits.Bandit.Regret
 import LeanBandits.AlgorithmBuilding
 
 /-!
diff --git a/LeanBandits/AlgorithmBuilding.lean b/LeanBandits/AlgorithmBuilding.lean
index 6ece4b8e..622755cd 100644
--- a/LeanBandits/AlgorithmBuilding.lean
+++ b/LeanBandits/AlgorithmBuilding.lean
@@ -3,7 +3,7 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
-import LeanBandits.Bandit
+import LeanBandits.Bandit.Bandit
 
 /-! # Tools to build bandit algorithms
 
diff --git a/LeanBandits/Bandit.lean b/LeanBandits/Bandit/Bandit.lean
similarity index 99%
rename from LeanBandits/Bandit.lean
rename to LeanBandits/Bandit/Bandit.lean
index 0fe0c865..82a0b15f 100644
--- a/LeanBandits/Bandit.lean
+++ b/LeanBandits/Bandit/Bandit.lean
@@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne, Paulo Rauber
 -/
 import Mathlib
-import LeanBandits.Algorithm
+import LeanBandits.SequentialLearning.Algorithm
 import LeanBandits.ForMathlib.CondDistrib
 import LeanBandits.ForMathlib.Traj
 
diff --git a/LeanBandits/Regret.lean b/LeanBandits/Bandit/Regret.lean
similarity index 99%
rename from LeanBandits/Regret.lean
rename to LeanBandits/Bandit/Regret.lean
index 99596a6d..0ada931b 100644
--- a/LeanBandits/Regret.lean
+++ b/LeanBandits/Bandit/Regret.lean
@@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne, Paulo Rauber
 -/
 import Mathlib
-import LeanBandits.Bandit
+import LeanBandits.Bandit.Bandit
 
 /-!
 # Regret
diff --git a/LeanBandits/ETC.lean b/LeanBandits/BanditAlgorithms/ETC.lean
similarity index 99%
rename from LeanBandits/ETC.lean
rename to LeanBandits/BanditAlgorithms/ETC.lean
index 75d7b53f..d212581d 100644
--- a/LeanBandits/ETC.lean
+++ b/LeanBandits/BanditAlgorithms/ETC.lean
@@ -6,8 +6,8 @@ Authors: Rémy Degenne
 import Mathlib.Probability.Moments.SubGaussian
 import LeanBandits.AlgorithmAndRandomVariables
 import LeanBandits.AlgorithmBuilding
+import LeanBandits.Bandit.Regret
 import LeanBandits.ForMathlib.SubGaussian
-import LeanBandits.Regret
 import LeanBandits.RewardByCountMeasure
 
 /-! # The Explore-Then-Commit Algorithm
diff --git a/LeanBandits/UCB.lean b/LeanBandits/BanditAlgorithms/UCB.lean
similarity index 98%
rename from LeanBandits/UCB.lean
rename to LeanBandits/BanditAlgorithms/UCB.lean
index 29d2bc4b..548bcd53 100644
--- a/LeanBandits/UCB.lean
+++ b/LeanBandits/BanditAlgorithms/UCB.lean
@@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
 import LeanBandits.AlgorithmBuilding
-import LeanBandits.Regret
+import LeanBandits.Bandit.Regret
 
 /-!
 # UCB algorithm
diff --git a/LeanBandits/RewardByCountMeasure.lean b/LeanBandits/RewardByCountMeasure.lean
index 60df2295..cf073f5e 100644
--- a/LeanBandits/RewardByCountMeasure.lean
+++ b/LeanBandits/RewardByCountMeasure.lean
@@ -3,10 +3,10 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
-import LeanBandits.Bandit
+import LeanBandits.Bandit.Bandit
+import LeanBandits.Bandit.Regret
 import LeanBandits.ForMathlib.IdentDistrib
 import LeanBandits.ForMathlib.IndepFun
-import LeanBandits.Regret
 
 /-! # Laws of `stepsUntil` and `rewardByCount`
 -/
diff --git a/LeanBandits/Algorithm.lean b/LeanBandits/SequentialLearning/Algorithm.lean
similarity index 100%
rename from LeanBandits/Algorithm.lean
rename to LeanBandits/SequentialLearning/Algorithm.lean

From 3b19a6e2c151050927415b5fa4cef261b1f88be5 Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Fri, 24 Oct 2025 18:59:46 +0200
Subject: [PATCH 4/6] split file

---
 LeanBandits.lean                              |  1 +
 LeanBandits/Bandit/Bandit.lean                |  2 +-
 LeanBandits/SequentialLearning/Algorithm.lean | 42 +-------------
 .../SequentialLearning/Deterministic.lean     | 56 +++++++++++++++++++
 4 files changed, 59 insertions(+), 42 deletions(-)
 create mode 100644 LeanBandits/SequentialLearning/Deterministic.lean

diff --git a/LeanBandits.lean b/LeanBandits.lean
index 9d11fbcd..51c2b3a9 100644
--- a/LeanBandits.lean
+++ b/LeanBandits.lean
@@ -13,3 +13,4 @@ import LeanBandits.ForMathlib.SubGaussian
 import LeanBandits.ForMathlib.Traj
 import LeanBandits.RewardByCountMeasure
 import LeanBandits.SequentialLearning.Algorithm
+import LeanBandits.SequentialLearning.Deterministic
diff --git a/LeanBandits/Bandit/Bandit.lean b/LeanBandits/Bandit/Bandit.lean
index 82a0b15f..b15e4196 100644
--- a/LeanBandits/Bandit/Bandit.lean
+++ b/LeanBandits/Bandit/Bandit.lean
@@ -4,7 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne, Paulo Rauber
 -/
 import Mathlib
-import LeanBandits.SequentialLearning.Algorithm
+import LeanBandits.SequentialLearning.Deterministic
 import LeanBandits.ForMathlib.CondDistrib
 import LeanBandits.ForMathlib.Traj
 
diff --git a/LeanBandits/SequentialLearning/Algorithm.lean b/LeanBandits/SequentialLearning/Algorithm.lean
index b108495a..e561b57b 100644
--- a/LeanBandits/SequentialLearning/Algorithm.lean
+++ b/LeanBandits/SequentialLearning/Algorithm.lean
@@ -9,7 +9,7 @@ import LeanBandits.ForMathlib.Measurable
 import LeanBandits.ForMathlib.Traj
 
 /-!
-# Bandit
+# Algorithms
 -/
 
 open MeasureTheory ProbabilityTheory Filter Real Finset
@@ -235,46 +235,6 @@ lemma condDistrib_reward_zero [StandardBorelSpace R] [Nonempty R]
   have h_action := (hasLaw_action_zero alg env).map_eq
   rwa [condDistrib_ae_eq_iff_measure_eq_compProd _ (by fun_prop), h_action]
 
-section DetAlgorithm
-
-/-- A deterministic algorithm. -/
-@[simps]
-noncomputable
-def detAlgorithm (nextaction : (n : ℕ) → (Iic n → α × R) → α)
-    (h_next : ∀ n, Measurable (nextaction n)) (action0 : α) :
-    Algorithm α R where
-  policy n := Kernel.deterministic (nextaction n) (h_next n)
-  p0 := Measure.dirac action0
-
-variable {nextaction : (n : ℕ) → (Iic n → α × R) → α} {h_next : ∀ n, Measurable (nextaction n)}
-  {action0 : α} {env : Environment α R}
-
-local notation "𝔓" => trajMeasure (detAlgorithm nextaction h_next action0) env
-
-lemma HasLaw_action_zero_detAlgorithm : HasLaw (action 0) (Measure.dirac action0) 𝔓 where
-  map_eq := (hasLaw_action_zero _ _).map_eq
-
-lemma action_zero_detAlgorithm [MeasurableSingletonClass α] : action 0 =ᵐ[𝔓] fun _ ↦ action0 := by
-  have h_eq : ∀ᵐ x ∂((𝔓).map (action 0)), x = action0 := by
-    rw [(hasLaw_action_zero _ _).map_eq]
-    simp [detAlgorithm]
-  exact ae_of_ae_map (by fun_prop) h_eq
-
-lemma action_detAlgorithm_ae_eq
-    [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R]
-    (n : ℕ) :
-    action (n + 1) =ᵐ[𝔓] fun h ↦ nextaction n (fun i ↦ h i) := by
-  have h := condDistrib_action (detAlgorithm nextaction h_next action0) env n
-  simp only [detAlgorithm_policy] at h
-  sorry
-
-example [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R] :
-    ∀ᵐ h ∂𝔓, action 0 h = action0 ∧ ∀ n, action (n + 1) h = nextaction n (fun i ↦ h i) := by
-  rw [eventually_and, ae_all_iff]
-  exact ⟨action_zero_detAlgorithm, action_detAlgorithm_ae_eq⟩
-
-end DetAlgorithm
-
 section stationaryEnv
 
 /-- A stationary environment, in which the distribution of the next reward depends only on the last
diff --git a/LeanBandits/SequentialLearning/Deterministic.lean b/LeanBandits/SequentialLearning/Deterministic.lean
new file mode 100644
index 00000000..5a56ffe8
--- /dev/null
+++ b/LeanBandits/SequentialLearning/Deterministic.lean
@@ -0,0 +1,56 @@
+/-
+Copyright (c) 2025 Rémy Degenne. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Rémy Degenne
+-/
+import LeanBandits.SequentialLearning.Algorithm
+
+/-!
+# Deterministic algorithms
+-/
+
+open MeasureTheory ProbabilityTheory Filter Real Finset
+
+open scoped ENNReal NNReal
+
+namespace Learning
+
+variable {α R : Type*} {mα : MeasurableSpace α} {mR : MeasurableSpace R}
+
+/-- A deterministic algorithm. -/
+@[simps]
+noncomputable
+def detAlgorithm (nextaction : (n : ℕ) → (Iic n → α × R) → α)
+    (h_next : ∀ n, Measurable (nextaction n)) (action0 : α) :
+    Algorithm α R where
+  policy n := Kernel.deterministic (nextaction n) (h_next n)
+  p0 := Measure.dirac action0
+
+variable {nextaction : (n : ℕ) → (Iic n → α × R) → α} {h_next : ∀ n, Measurable (nextaction n)}
+  {action0 : α} {env : Environment α R}
+
+local notation "𝔓" => trajMeasure (detAlgorithm nextaction h_next action0) env
+
+lemma HasLaw_action_zero_detAlgorithm : HasLaw (action 0) (Measure.dirac action0) 𝔓 where
+  map_eq := (hasLaw_action_zero _ _).map_eq
+
+lemma action_zero_detAlgorithm [MeasurableSingletonClass α] : action 0 =ᵐ[𝔓] fun _ ↦ action0 := by
+  have h_eq : ∀ᵐ x ∂((𝔓).map (action 0)), x = action0 := by
+    rw [(hasLaw_action_zero _ _).map_eq]
+    simp [detAlgorithm]
+  exact ae_of_ae_map (by fun_prop) h_eq
+
+lemma action_detAlgorithm_ae_eq
+    [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R]
+    (n : ℕ) :
+    action (n + 1) =ᵐ[𝔓] fun h ↦ nextaction n (fun i ↦ h i) := by
+  have h := condDistrib_action (detAlgorithm nextaction h_next action0) env n
+  simp only [detAlgorithm_policy] at h
+  sorry
+
+example [StandardBorelSpace α] [Nonempty α] [StandardBorelSpace R] [Nonempty R] :
+    ∀ᵐ h ∂𝔓, action 0 h = action0 ∧ ∀ n, action (n + 1) h = nextaction n (fun i ↦ h i) := by
+  rw [eventually_and, ae_all_iff]
+  exact ⟨action_zero_detAlgorithm, action_detAlgorithm_ae_eq⟩
+
+end Learning

From 486fca110ddcd7914bc8b99d00c092942fffdde7 Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Fri, 24 Oct 2025 19:03:10 +0200
Subject: [PATCH 5/6] split file

---
 LeanBandits.lean                             |  1 +
 LeanBandits/AlgorithmBuilding.lean           | 72 +----------------
 LeanBandits/ForMathlib/MeasurableArgMax.lean | 84 ++++++++++++++++++++
 3 files changed, 86 insertions(+), 71 deletions(-)
 create mode 100644 LeanBandits/ForMathlib/MeasurableArgMax.lean

diff --git a/LeanBandits.lean b/LeanBandits.lean
index 51c2b3a9..69143345 100644
--- a/LeanBandits.lean
+++ b/LeanBandits.lean
@@ -9,6 +9,7 @@ import LeanBandits.ForMathlib.IdentDistrib
 import LeanBandits.ForMathlib.IndepFun
 import LeanBandits.ForMathlib.KernelSub
 import LeanBandits.ForMathlib.Measurable
+import LeanBandits.ForMathlib.MeasurableArgMax
 import LeanBandits.ForMathlib.SubGaussian
 import LeanBandits.ForMathlib.Traj
 import LeanBandits.RewardByCountMeasure
diff --git a/LeanBandits/AlgorithmBuilding.lean b/LeanBandits/AlgorithmBuilding.lean
index 622755cd..1755f1b5 100644
--- a/LeanBandits/AlgorithmBuilding.lean
+++ b/LeanBandits/AlgorithmBuilding.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
 import LeanBandits.Bandit.Bandit
+import LeanBandits.ForMathlib.MeasurableArgMax
 
 /-! # Tools to build bandit algorithms
 
@@ -12,77 +13,6 @@ import LeanBandits.Bandit.Bandit
 open MeasureTheory ProbabilityTheory Finset
 open scoped ENNReal NNReal
 
-section MeasurableArgmax -- copied from PR #27579 (and changed from argmin to argmax)
-
-lemma measurable_encode {α : Type*} {_ : MeasurableSpace α} [Encodable α]
-    [MeasurableSingletonClass α] :
-    Measurable (Encodable.encode (α := α)) := by
-  refine measurable_to_nat fun a ↦ ?_
-  have : Encodable.encode ⁻¹' {Encodable.encode a} = {a} := by ext; simp
-  rw [this]
-  exact measurableSet_singleton _
-
-lemma measurableEmbedding_encode (α : Type*) {_ : MeasurableSpace α} [Encodable α]
-    [MeasurableSingletonClass α] :
-    MeasurableEmbedding (Encodable.encode (α := α)) where
-  injective := Encodable.encode_injective
-  measurable := measurable_encode
-  measurableSet_image' _ _ := .of_discrete
-
-section Finite
-
-variable {𝓧 𝓨 α : Type*} {m𝓧 : MeasurableSpace 𝓧} {m𝓨 : MeasurableSpace 𝓨}
-  {mα : MeasurableSpace α} [TopologicalSpace α] [LinearOrder α]
-  [OpensMeasurableSpace α] [OrderClosedTopology α] [SecondCountableTopology α]
-
-lemma measurableSet_isMax [Countable 𝓨]
-    {f : 𝓧 → 𝓨 → α} (hf : ∀ y, Measurable (fun x ↦ f x y)) (y : 𝓨) :
-    MeasurableSet {x | ∀ z, f x z ≤ f x y} := by
-  rw [show {x | ∀ y', f x y' ≤ f x y} = ⋂ y', {x | f x y' ≤ f x y} by ext; simp]
-  exact MeasurableSet.iInter fun z ↦ measurableSet_le (by fun_prop) (by fun_prop)
-
-lemma exists_isMaxOn' {α : Type*} [LinearOrder α]
-    [Nonempty 𝓨] [Finite 𝓨] [Encodable 𝓨] (f : 𝓧 → 𝓨 → α) (x : 𝓧) :
-    ∃ n : ℕ, ∃ y, n = Encodable.encode y ∧ ∀ z, f x z ≤ f x y := by
-  obtain ⟨y, h⟩ := Finite.exists_max (f x)
-  exact ⟨Encodable.encode y, y, rfl, h⟩
-
-/-- A measurable argmax function. -/
-noncomputable
-def measurableArgmax [Nonempty 𝓨] [Finite 𝓨] [Encodable 𝓨] [MeasurableSingletonClass 𝓨]
-    (f : 𝓧 → 𝓨 → α)
-    [∀ x, DecidablePred fun n ↦ ∃ y, n = Encodable.encode y ∧ ∀ (z : 𝓨), f x z ≤ f x y]
-    (x : 𝓧) :
-    𝓨 :=
-  (measurableEmbedding_encode 𝓨).invFun (Nat.find (exists_isMaxOn' f x))
-
-lemma measurable_measurableArgmax [Nonempty 𝓨] [Finite 𝓨] [Encodable 𝓨] [MeasurableSingletonClass 𝓨]
-    {f : 𝓧 → 𝓨 → α}
-    [∀ x, DecidablePred fun n ↦ ∃ y, n = Encodable.encode y ∧ ∀ (z : 𝓨), f x z ≤ f x y]
-    (hf : ∀ y, Measurable (fun x ↦ f x y)) :
-    Measurable (measurableArgmax f) := by
-  refine (MeasurableEmbedding.measurable_invFun (measurableEmbedding_encode 𝓨)).comp ?_
-  refine measurable_find _ fun n ↦ ?_
-  have : {x | ∃ y, n = Encodable.encode y ∧ ∀ (z : 𝓨), f x z ≤ f x y}
-      = ⋃ y, ({x | n = Encodable.encode y} ∩ {x | ∀ z, f x z ≤ f x y}) := by ext; simp
-  rw [this]
-  refine MeasurableSet.iUnion fun y ↦ (MeasurableSet.inter (by simp) ?_)
-  exact measurableSet_isMax (by fun_prop) y
-
-lemma isMaxOn_measurableArgmax {α : Type*} [LinearOrder α]
-    [Nonempty 𝓨] [Finite 𝓨] [Encodable 𝓨] [MeasurableSingletonClass 𝓨]
-    (f : 𝓧 → 𝓨 → α)
-    [∀ x, DecidablePred fun n ↦ ∃ y, n = Encodable.encode y ∧ ∀ (z : 𝓨), f x z ≤ f x y]
-    (x : 𝓧) (z : 𝓨) :
-    f x z ≤ f x (measurableArgmax f x) := by
-  obtain ⟨y, h_eq, h_le⟩ := Nat.find_spec (exists_isMaxOn' f x)
-  refine le_trans (h_le z) (le_of_eq ?_)
-  rw [measurableArgmax, h_eq,
-    MeasurableEmbedding.leftInverse_invFun (measurableEmbedding_encode 𝓨) y]
-
-end Finite
-end MeasurableArgmax
-
 namespace Bandits
 
 variable {α : Type*} [DecidableEq α] [MeasurableSpace α]
diff --git a/LeanBandits/ForMathlib/MeasurableArgMax.lean b/LeanBandits/ForMathlib/MeasurableArgMax.lean
new file mode 100644
index 00000000..0769a25d
--- /dev/null
+++ b/LeanBandits/ForMathlib/MeasurableArgMax.lean
@@ -0,0 +1,84 @@
+/-
+Copyright (c) 2025 Rémy Degenne. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Rémy Degenne
+-/
+import Mathlib.MeasureTheory.Constructions.BorelSpace.Order
+
+/-! # Measurable argmax function
+
+-/
+
+open MeasureTheory Finset
+open scoped ENNReal NNReal
+
+section MeasurableArgmax -- copied from PR #27579 (and changed from argmin to argmax)
+
+lemma measurable_encode {α : Type*} {_ : MeasurableSpace α} [Encodable α]
+    [MeasurableSingletonClass α] :
+    Measurable (Encodable.encode (α := α)) := by
+  refine measurable_to_nat fun a ↦ ?_
+  have : Encodable.encode ⁻¹' {Encodable.encode a} = {a} := by ext; simp
+  rw [this]
+  exact measurableSet_singleton _
+
+lemma measurableEmbedding_encode (α : Type*) {_ : MeasurableSpace α} [Encodable α]
+    [MeasurableSingletonClass α] :
+    MeasurableEmbedding (Encodable.encode (α := α)) where
+  injective := Encodable.encode_injective
+  measurable := measurable_encode
+  measurableSet_image' _ _ := .of_discrete
+
+section Finite
+
+variable {𝓧 𝓨 α : Type*} {m𝓧 : MeasurableSpace 𝓧} {m𝓨 : MeasurableSpace 𝓨}
+  {mα : MeasurableSpace α} [TopologicalSpace α] [LinearOrder α]
+  [OpensMeasurableSpace α] [OrderClosedTopology α] [SecondCountableTopology α]
+
+lemma measurableSet_isMax [Countable 𝓨]
+    {f : 𝓧 → 𝓨 → α} (hf : ∀ y, Measurable (fun x ↦ f x y)) (y : 𝓨) :
+    MeasurableSet {x | ∀ z, f x z ≤ f x y} := by
+  rw [show {x | ∀ y', f x y' ≤ f x y} = ⋂ y', {x | f x y' ≤ f x y} by ext; simp]
+  exact MeasurableSet.iInter fun z ↦ measurableSet_le (by fun_prop) (by fun_prop)
+
+lemma exists_isMaxOn' {α : Type*} [LinearOrder α]
+    [Nonempty 𝓨] [Finite 𝓨] [Encodable 𝓨] (f : 𝓧 → 𝓨 → α) (x : 𝓧) :
+    ∃ n : ℕ, ∃ y, n = Encodable.encode y ∧ ∀ z, f x z ≤ f x y := by
+  obtain ⟨y, h⟩ := Finite.exists_max (f x)
+  exact ⟨Encodable.encode y, y, rfl, h⟩
+
+/-- A measurable argmax function. -/
+noncomputable
+def measurableArgmax [Nonempty 𝓨] [Finite 𝓨] [Encodable 𝓨] [MeasurableSingletonClass 𝓨]
+    (f : 𝓧 → 𝓨 → α)
+    [∀ x, DecidablePred fun n ↦ ∃ y, n = Encodable.encode y ∧ ∀ (z : 𝓨), f x z ≤ f x y]
+    (x : 𝓧) :
+    𝓨 :=
+  (measurableEmbedding_encode 𝓨).invFun (Nat.find (exists_isMaxOn' f x))
+
+lemma measurable_measurableArgmax [Nonempty 𝓨] [Finite 𝓨] [Encodable 𝓨] [MeasurableSingletonClass 𝓨]
+    {f : 𝓧 → 𝓨 → α}
+    [∀ x, DecidablePred fun n ↦ ∃ y, n = Encodable.encode y ∧ ∀ (z : 𝓨), f x z ≤ f x y]
+    (hf : ∀ y, Measurable (fun x ↦ f x y)) :
+    Measurable (measurableArgmax f) := by
+  refine (MeasurableEmbedding.measurable_invFun (measurableEmbedding_encode 𝓨)).comp ?_
+  refine measurable_find _ fun n ↦ ?_
+  have : {x | ∃ y, n = Encodable.encode y ∧ ∀ (z : 𝓨), f x z ≤ f x y}
+      = ⋃ y, ({x | n = Encodable.encode y} ∩ {x | ∀ z, f x z ≤ f x y}) := by ext; simp
+  rw [this]
+  refine MeasurableSet.iUnion fun y ↦ (MeasurableSet.inter (by simp) ?_)
+  exact measurableSet_isMax (by fun_prop) y
+
+lemma isMaxOn_measurableArgmax {α : Type*} [LinearOrder α]
+    [Nonempty 𝓨] [Finite 𝓨] [Encodable 𝓨] [MeasurableSingletonClass 𝓨]
+    (f : 𝓧 → 𝓨 → α)
+    [∀ x, DecidablePred fun n ↦ ∃ y, n = Encodable.encode y ∧ ∀ (z : 𝓨), f x z ≤ f x y]
+    (x : 𝓧) (z : 𝓨) :
+    f x z ≤ f x (measurableArgmax f x) := by
+  obtain ⟨y, h_eq, h_le⟩ := Nat.find_spec (exists_isMaxOn' f x)
+  refine le_trans (h_le z) (le_of_eq ?_)
+  rw [measurableArgmax, h_eq,
+    MeasurableEmbedding.leftInverse_invFun (measurableEmbedding_encode 𝓨) y]
+
+end Finite
+end MeasurableArgmax

From 26fb61028abb83d2206f21e5c3680730c7a9671e Mon Sep 17 00:00:00 2001
From: Remy Degenne <remydegenne@gmail.com>
Date: Fri, 24 Oct 2025 19:22:53 +0200
Subject: [PATCH 6/6] min_imports

---
 LeanBandits/AlgorithmBuilding.lean            | 7 ++++---
 LeanBandits/Bandit/Bandit.lean                | 5 ++---
 LeanBandits/Bandit/Regret.lean                | 3 ++-
 LeanBandits/BanditAlgorithms/ETC.lean         | 4 +---
 LeanBandits/BanditAlgorithms/UCB.lean         | 1 +
 LeanBandits/ForMathlib/CondDistrib.lean       | 4 +---
 LeanBandits/ForMathlib/IndepFun.lean          | 3 ++-
 LeanBandits/ForMathlib/Measurable.lean        | 2 +-
 LeanBandits/SequentialLearning/Algorithm.lean | 3 +--
 9 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/LeanBandits/AlgorithmBuilding.lean b/LeanBandits/AlgorithmBuilding.lean
index 1755f1b5..4cb518dd 100644
--- a/LeanBandits/AlgorithmBuilding.lean
+++ b/LeanBandits/AlgorithmBuilding.lean
@@ -3,14 +3,15 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
-import LeanBandits.Bandit.Bandit
-import LeanBandits.ForMathlib.MeasurableArgMax
+import Mathlib.Analysis.Normed.Ring.Basic
+import Mathlib.MeasureTheory.Constructions.BorelSpace.Basic
+import Mathlib.Topology.Compactness.PseudometrizableLindelof
 
 /-! # Tools to build bandit algorithms
 
 -/
 
-open MeasureTheory ProbabilityTheory Finset
+open MeasureTheory Finset
 open scoped ENNReal NNReal
 
 namespace Bandits
diff --git a/LeanBandits/Bandit/Bandit.lean b/LeanBandits/Bandit/Bandit.lean
index b15e4196..e109799a 100644
--- a/LeanBandits/Bandit/Bandit.lean
+++ b/LeanBandits/Bandit/Bandit.lean
@@ -3,10 +3,9 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne, Paulo Rauber
 -/
-import Mathlib
 import LeanBandits.SequentialLearning.Deterministic
-import LeanBandits.ForMathlib.CondDistrib
-import LeanBandits.ForMathlib.Traj
+import Mathlib.Probability.IdentDistrib
+import Mathlib.Probability.Independence.InfinitePi
 
 /-!
 # Bandit
diff --git a/LeanBandits/Bandit/Regret.lean b/LeanBandits/Bandit/Regret.lean
index 0ada931b..70ac3c2c 100644
--- a/LeanBandits/Bandit/Regret.lean
+++ b/LeanBandits/Bandit/Regret.lean
@@ -3,8 +3,9 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne, Paulo Rauber
 -/
-import Mathlib
 import LeanBandits.Bandit.Bandit
+import Mathlib.Data.ENat.Lattice
+import Mathlib.Order.CompletePartialOrder
 
 /-!
 # Regret
diff --git a/LeanBandits/BanditAlgorithms/ETC.lean b/LeanBandits/BanditAlgorithms/ETC.lean
index d212581d..eecf76f5 100644
--- a/LeanBandits/BanditAlgorithms/ETC.lean
+++ b/LeanBandits/BanditAlgorithms/ETC.lean
@@ -3,10 +3,8 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
-import Mathlib.Probability.Moments.SubGaussian
 import LeanBandits.AlgorithmAndRandomVariables
-import LeanBandits.AlgorithmBuilding
-import LeanBandits.Bandit.Regret
+import LeanBandits.ForMathlib.MeasurableArgMax
 import LeanBandits.ForMathlib.SubGaussian
 import LeanBandits.RewardByCountMeasure
 
diff --git a/LeanBandits/BanditAlgorithms/UCB.lean b/LeanBandits/BanditAlgorithms/UCB.lean
index 548bcd53..2db28bc2 100644
--- a/LeanBandits/BanditAlgorithms/UCB.lean
+++ b/LeanBandits/BanditAlgorithms/UCB.lean
@@ -5,6 +5,7 @@ Authors: Rémy Degenne
 -/
 import LeanBandits.AlgorithmBuilding
 import LeanBandits.Bandit.Regret
+import LeanBandits.ForMathlib.MeasurableArgMax
 
 /-!
 # UCB algorithm
diff --git a/LeanBandits/ForMathlib/CondDistrib.lean b/LeanBandits/ForMathlib/CondDistrib.lean
index 7689d834..8cec6c80 100644
--- a/LeanBandits/ForMathlib/CondDistrib.lean
+++ b/LeanBandits/ForMathlib/CondDistrib.lean
@@ -3,12 +3,10 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
+import LeanBandits.ForMathlib.KernelSub
 import Mathlib.MeasureTheory.Measure.ProbabilityMeasure
 import Mathlib.Probability.Independence.Basic
 import Mathlib.Probability.Independence.Conditional
-import Mathlib.Probability.Kernel.CompProdEqIff
-import Mathlib.Probability.Kernel.Composition.Lemmas
-import LeanBandits.ForMathlib.KernelSub
 
 open MeasureTheory ProbabilityTheory Finset
 open scoped ENNReal NNReal
diff --git a/LeanBandits/ForMathlib/IndepFun.lean b/LeanBandits/ForMathlib/IndepFun.lean
index 6630d037..f5e287b4 100644
--- a/LeanBandits/ForMathlib/IndepFun.lean
+++ b/LeanBandits/ForMathlib/IndepFun.lean
@@ -1,4 +1,5 @@
-import Mathlib
+import Mathlib.Probability.IdentDistrib
+import Mathlib.Probability.Independence.InfinitePi
 
 open MeasureTheory Finset
 
diff --git a/LeanBandits/ForMathlib/Measurable.lean b/LeanBandits/ForMathlib/Measurable.lean
index 3c54b21f..3f907ce9 100644
--- a/LeanBandits/ForMathlib/Measurable.lean
+++ b/LeanBandits/ForMathlib/Measurable.lean
@@ -3,7 +3,7 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
-import Mathlib
+import Mathlib.MeasureTheory.MeasurableSpace.Basic
 
 /-!
 # Measurability lemmas
diff --git a/LeanBandits/SequentialLearning/Algorithm.lean b/LeanBandits/SequentialLearning/Algorithm.lean
index e561b57b..a52ccf45 100644
--- a/LeanBandits/SequentialLearning/Algorithm.lean
+++ b/LeanBandits/SequentialLearning/Algorithm.lean
@@ -3,10 +3,9 @@ Copyright (c) 2025 Rémy Degenne. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne, Paulo Rauber
 -/
-import Mathlib
-import LeanBandits.ForMathlib.CondDistrib
 import LeanBandits.ForMathlib.Measurable
 import LeanBandits.ForMathlib.Traj
+import Mathlib.Probability.HasLaw
 
 /-!
 # Algorithms