diff --git a/LeanBandits.lean b/LeanBandits.lean
index acbeb2ee..457152e6 100644
--- a/LeanBandits.lean
+++ b/LeanBandits.lean
@@ -8,6 +8,7 @@ import LeanBandits.ForMathlib.IdentDistrib
 import LeanBandits.ForMathlib.KernelCompositionLemmas
 import LeanBandits.ForMathlib.KernelCompositionParallelComp
 import LeanBandits.ForMathlib.KernelSub
+import LeanBandits.ForMathlib.Measurable
 import LeanBandits.ForMathlib.SubGaussian
 import LeanBandits.ForMathlib.Traj
 import LeanBandits.Regret
diff --git a/LeanBandits/Algorithm.lean b/LeanBandits/Algorithm.lean
index 46fb402e..cc7a9ce1 100644
--- a/LeanBandits/Algorithm.lean
+++ b/LeanBandits/Algorithm.lean
@@ -6,6 +6,7 @@ Authors: Rémy Degenne, Paulo Rauber
 import Mathlib
 import LeanBandits.ForMathlib.CondDistrib
 import LeanBandits.ForMathlib.KernelCompositionLemmas
+import LeanBandits.ForMathlib.Measurable
 import LeanBandits.ForMathlib.Traj
 
 /-!
@@ -223,10 +224,8 @@ lemma measurable_hist_filtration (n : ℕ) : Measurable[Learning.filtration α R
 
 lemma measurable_action_filtration (n : ℕ) : Measurable[Learning.filtration α R n] (action n) := by
   simp only [Learning.filtration, Filtration.piLE_eq_comap_frestrictLe, ← hist_eq_frestrictLe]
-  rw [action_eq_eval_comp_hist, measurable_iff_comap_le, ← MeasurableSpace.comap_comp]
-  refine MeasurableSpace.comap_mono ?_
-  rw [← measurable_iff_comap_le]
-  fun_prop
+  rw [action_eq_eval_comp_hist]
+  exact measurable_comp_comap _ (by fun_prop)
 
 lemma adapted_action [TopologicalSpace α] [TopologicalSpace.PseudoMetrizableSpace α]
     [SecondCountableTopology α] [OpensMeasurableSpace α] :
@@ -235,10 +234,8 @@ lemma adapted_action [TopologicalSpace α] [TopologicalSpace.PseudoMetrizableSpa
 
 lemma measurable_reward_filtration (n : ℕ) : Measurable[Learning.filtration α R n] (reward n) := by
   simp only [Learning.filtration, Filtration.piLE_eq_comap_frestrictLe, ← hist_eq_frestrictLe]
-  rw [reward_eq_eval_comp_hist, measurable_iff_comap_le, ← MeasurableSpace.comap_comp]
-  refine MeasurableSpace.comap_mono ?_
-  rw [← measurable_iff_comap_le]
-  fun_prop
+  rw [reward_eq_eval_comp_hist]
+  exact measurable_comp_comap _ (by fun_prop)
 
 lemma adapted_reward [TopologicalSpace R] [TopologicalSpace.PseudoMetrizableSpace R]
     [SecondCountableTopology R] [OpensMeasurableSpace R] :
diff --git a/LeanBandits/Bandit.lean b/LeanBandits/Bandit.lean
index 28435b1b..f9bb90f5 100644
--- a/LeanBandits/Bandit.lean
+++ b/LeanBandits/Bandit.lean
@@ -146,6 +146,13 @@ lemma indepFun_eval_streamMeasure' (ν : Kernel α R) [IsMarkovKernel ν] {a b :
     IndepFun (fun ω n ↦ ω n a) (fun ω n ↦ ω n b) (Bandit.streamMeasure ν) := by
   sorry
 
+lemma indepFun_eval_snd_measure (alg : Algorithm α R) (ν : Kernel α R) [IsMarkovKernel ν]
+    {a b : α} (h : a ≠ b) :
+    IndepFun (fun ω n ↦ ω.2 n a) (fun ω n ↦ ω.2 n b) (Bandit.measure alg ν) := by
+  refine indepFun_snd_prod ?_ ?_ (indepFun_eval_streamMeasure' ν h) (Bandit.trajMeasure alg ν)
+  · exact Measurable.aemeasurable (by fun_prop)
+  · exact Measurable.aemeasurable (by fun_prop)
+
 end StreamMeasure
 
 /-- `arm n` is the arm pulled at time `n`. This is a random variable on the measurable space
diff --git a/LeanBandits/ETC.lean b/LeanBandits/ETC.lean
index 0ff3076b..3b9561d0 100644
--- a/LeanBandits/ETC.lean
+++ b/LeanBandits/ETC.lean
@@ -6,7 +6,6 @@ Authors: Rémy Degenne
 import Mathlib.Probability.Moments.SubGaussian
 import LeanBandits.AlgorithmAndRandomVariables
 import LeanBandits.AlgorithmBuilding
-import LeanBandits.ForMathlib.IdentDistrib
 import LeanBandits.ForMathlib.SubGaussian
 import LeanBandits.Regret
 import LeanBandits.RewardByCountMeasure
@@ -178,12 +177,26 @@ lemma identDistrib_aux (m : ℕ) (a b : Fin K) :
     IdentDistrib
       (fun ω ↦ (∑ s ∈ Icc 1 m, rewardByCount a s ω.1 ω.2, ∑ s ∈ Icc 1 m, rewardByCount b s ω.1 ω.2))
       (fun ω ↦ (∑ s ∈ range m, ω.2 s a, ∑ s ∈ range m, ω.2 s b)) 𝔓 𝔓 := by
+  have : Nonempty (Fin K) := Fin.pos_iff_nonempty.mp hK
   have h1 (a : Fin K) :
-      IdentDistrib (fun ω s ↦ rewardByCount a s ω.1 ω.2) (fun ω s ↦ ω.2 s a) 𝔓 𝔓 := by
-    sorry
+      IdentDistrib (fun ω s ↦ rewardByCount a (s + 1) ω.1 ω.2) (fun ω s ↦ ω.2 s a) 𝔓 𝔓 :=
+    identDistrib_rewardByCount_stream a
   have h2 (a : Fin K) : IdentDistrib (fun ω ↦ ∑ s ∈ Icc 1 m, rewardByCount a s ω.1 ω.2)
       (fun ω ↦ ∑ s ∈ range m, ω.2 s a) 𝔓 𝔓 := by
-    sorry
+    have h_eq (ω : (ℕ → Fin K × ℝ) × (ℕ → Fin K → ℝ)) : ∑ s ∈ Icc 1 m, rewardByCount a s ω.1 ω.2
+        = ∑ s ∈ range m, rewardByCount a (s + 1) ω.1 ω.2 := by
+      let e : Icc 1 m ≃ range m :=
+      { toFun x := ⟨x - 1, by have h := x.2; simp only [mem_Icc] at h; grind⟩
+        invFun x := ⟨x + 1, by have h := x.2; simp only [mem_Icc]; grind⟩
+        left_inv x := by have h := x.2; simp only [mem_Icc] at h; grind
+        right_inv x := by have h := x.2; grind }
+      rw [← sum_coe_sort (Icc 1 m), ← sum_coe_sort (range m), sum_equiv e]
+      · simp
+      · simp only [univ_eq_attach, mem_attach, forall_const, Subtype.forall, mem_Icc,
+          forall_and_index]
+        grind
+    simp_rw [h_eq]
+    exact IdentDistrib.comp (h1 a) (u := fun p ↦ ∑ s ∈ range m, p s) (by fun_prop)
   by_cases hab : a = b
   · simp only [hab]
     exact (h2 b).comp (u := fun p ↦ (p, p)) (by fun_prop)
@@ -192,11 +205,11 @@ lemma identDistrib_aux (m : ℕ) (a b : Fin K) :
         𝔓 by
       exact this.comp (φ := fun p ↦ ∑ i ∈ Icc 1 m, p i) (ψ := fun p ↦ ∑ j ∈ Icc 1 m, p j)
         (by fun_prop) (by fun_prop)
-    sorry
+    exact indepFun_rewardByCount_of_ne hab
   · suffices IndepFun (fun ω s ↦ ω.2 s a) (fun ω s ↦ ω.2 s b) 𝔓 by
       exact this.comp (φ := fun p ↦ ∑ i ∈ range m, p i) (ψ := fun p ↦ ∑ j ∈ range m, p j)
         (by fun_prop) (by fun_prop)
-    sorry
+    exact indepFun_eval_snd_measure _ ν hab
 
 lemma prob_arm_mul_eq_le (hν : ∀ a, HasSubgaussianMGF (fun x ↦ x - (ν a)[id]) 1 (ν a)) (a : Fin K)
     (hm : m ≠ 0) :
diff --git a/LeanBandits/ForMathlib/CondDistrib.lean b/LeanBandits/ForMathlib/CondDistrib.lean
index 26f7eef7..91e525a6 100644
--- a/LeanBandits/ForMathlib/CondDistrib.lean
+++ b/LeanBandits/ForMathlib/CondDistrib.lean
@@ -18,7 +18,7 @@ variable {α β γ δ Ω Ω' : Type*}
   {m mα : MeasurableSpace α} {μ : Measure α} {mβ : MeasurableSpace β} {mγ : MeasurableSpace γ}
   {mδ : MeasurableSpace δ}
   [MeasurableSpace Ω] [StandardBorelSpace Ω] [Nonempty Ω]
-  [MeasurableSpace Ω'] [StandardBorelSpace Ω'] [Nonempty Ω']
+  [mΩ' : MeasurableSpace Ω'] [StandardBorelSpace Ω'] [Nonempty Ω']
   {X : α → β} {Y : α → Ω} {Z : α → Ω'} {T : α → γ}
 
 lemma MeasurableSpace.comap_prodMk (X : α → β) (Y : α → γ) :
@@ -844,6 +844,87 @@ lemma condIndepFun_of_exists_condDistrib_prod_ae_eq_prodMkLeft
   rw [Kernel.prodMkLeft_ae_eq_iff, Measure.snd_map_prodMk (by fun_prop)]
   exact hη_eq.symm
 
+/-- Law of `Y` conditioned on `X`. -/
+notation "𝓛[" Y " | " X "; " μ "]" => condDistrib Y X μ
+
+-- generalize to map instead of fst
+omit [Nonempty Ω'] in
+lemma condIndepFun_fst_prod [StandardBorelSpace α] [StandardBorelSpace β] [Nonempty β]
+    [StandardBorelSpace γ]
+    (hX : Measurable X) (hY : Measurable Y) (hZ : Measurable Z)
+    (ν : Measure γ) [IsProbabilityMeasure ν]
+    (h_indep : CondIndepFun (mΩ'.comap Z) hZ.comap_le Y X μ) :
+    CondIndepFun (mΩ'.comap (fun ω ↦ Z ω.1)) (hZ.comp measurable_fst).comap_le
+      (fun ω ↦ Y ω.1) (fun ω ↦ X ω.1) (μ.prod ν) := by
+  rw [condIndepFun_comap_iff_map_prod_eq_prod_condDistrib_prod_condDistrib (by fun_prop)
+    (by fun_prop) (by fun_prop)] at h_indep ⊢
+  have h1 : 𝓛[fun ω ↦ Y ω.1 | fun ω ↦ Z ω.1; μ.prod ν] =ᵐ[μ.map Z] 𝓛[Y | Z; μ] :=
+    condDistrib_fst_prod (Y := Y) (X := Z) (ν := ν) (μ := μ) (by fun_prop) (by fun_prop)
+  have h2 : 𝓛[fun ω ↦ X ω.1 | fun ω ↦ Z ω.1; μ.prod ν] =ᵐ[μ.map Z] 𝓛[X | Z; μ] :=
+    condDistrib_fst_prod (Y := X) (X := Z) (ν := ν) (μ := μ) (by fun_prop) (by fun_prop)
+  have h_fst1 : (μ.prod ν).map (fun ω ↦ Z ω.1) = μ.map Z := by
+    conv_rhs => rw [← Measure.fst_prod (μ := μ) (ν := ν), Measure.fst,
+      Measure.map_map (by fun_prop) (by fun_prop)]
+    rfl
+  have h_fst2 : (μ.prod ν).map (fun ω ↦ (Z ω.1, Y ω.1, X ω.1))
+      = μ.map (fun ω ↦ (Z ω, Y ω, X ω)) := by
+    conv_rhs => rw [← Measure.fst_prod (μ := μ) (ν := ν), Measure.fst,
+      Measure.map_map (by fun_prop) (by fun_prop)]
+    rfl
+  rw [h_fst1, h_fst2, h_indep]
+  refine Measure.bind_congr_right ?_
+  filter_upwards [h1, h2] with x hx1 hx2
+  simp_rw [Kernel.prod_apply]
+  rw [hx1, ← hx2]
+
+omit [StandardBorelSpace Ω] [Nonempty Ω] in
+lemma indepFun_fst_prod (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) (h_indep : IndepFun X Y μ)
+    (ν : Measure γ) [IsProbabilityMeasure ν] :
+    IndepFun (fun ω ↦ X ω.1) (fun ω ↦ Y ω.1) (μ.prod ν) := by
+  rw [indepFun_iff_map_prod_eq_prod_map_map (by fun_prop) (by fun_prop)] at h_indep ⊢
+  have :  AEMeasurable (fun ω ↦ (X ω, Y ω)) (Measure.map Prod.fst (μ.prod ν)) := by
+    simp only [Measure.map_fst_prod, measure_univ, one_smul]
+    fun_prop
+  have :  AEMeasurable X (Measure.map Prod.fst (μ.prod ν)) := by
+    simp only [Measure.map_fst_prod, measure_univ, one_smul]
+    fun_prop
+  have :  AEMeasurable Y (Measure.map Prod.fst (μ.prod ν)) := by
+    simp only [Measure.map_fst_prod, measure_univ, one_smul]
+    fun_prop
+  have h : (μ.prod ν).map (fun ω ↦ (X ω.1, Y ω.1)) = μ.map (fun ω ↦ (X ω, Y ω)) := by
+    conv_rhs => rw [← Measure.fst_prod (μ := μ) (ν := ν), Measure.fst,
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
+    rfl
+  rw [h, h_indep]
+  conv_lhs => rw [← Measure.fst_prod (μ := μ) (ν := ν), Measure.fst,
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop),
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
+  rfl
+
+omit [StandardBorelSpace Ω] [Nonempty Ω] in
+lemma indepFun_snd_prod (hX : AEMeasurable X μ) (hY : AEMeasurable Y μ) (h_indep : IndepFun X Y μ)
+    (ν : Measure γ) [IsProbabilityMeasure ν] :
+    IndepFun (fun ω ↦ X ω.2) (fun ω ↦ Y ω.2) (ν.prod μ) := by
+  rw [indepFun_iff_map_prod_eq_prod_map_map (by fun_prop) (by fun_prop)] at h_indep ⊢
+  have :  AEMeasurable (fun ω ↦ (X ω, Y ω)) (Measure.map Prod.snd (ν.prod μ)) := by
+    simp only [Measure.map_snd_prod, measure_univ, one_smul]
+    fun_prop
+  have :  AEMeasurable X (Measure.map Prod.snd (ν.prod μ)) := by
+    simp only [Measure.map_snd_prod, measure_univ, one_smul]
+    fun_prop
+  have :  AEMeasurable Y (Measure.map Prod.snd (ν.prod μ)) := by
+    simp only [Measure.map_snd_prod, measure_univ, one_smul]
+    fun_prop
+  have h : (ν.prod μ).map (fun ω ↦ (X ω.2, Y ω.2)) = μ.map (fun ω ↦ (X ω, Y ω)) := by
+    conv_rhs => rw [← Measure.snd_prod (μ := ν) (ν := μ), Measure.snd,
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
+    rfl
+  rw [h, h_indep]
+  conv_lhs => rw [← Measure.snd_prod (μ := ν) (ν := μ), Measure.snd,
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop),
+      AEMeasurable.map_map_of_aemeasurable (by fun_prop) (by fun_prop)]
+  rfl
+
 end CondDistrib
 
 section Cond
diff --git a/LeanBandits/ForMathlib/IdentDistrib.lean b/LeanBandits/ForMathlib/IdentDistrib.lean
index e1b650b6..0e542109 100644
--- a/LeanBandits/ForMathlib/IdentDistrib.lean
+++ b/LeanBandits/ForMathlib/IdentDistrib.lean
@@ -4,13 +4,14 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
 import Mathlib.Probability.IdentDistrib
+import Mathlib.Probability.Independence.InfinitePi
 
 open MeasureTheory ProbabilityTheory Finset
 open scoped ENNReal NNReal
 
 namespace ProbabilityTheory
 
-variable {Ω Ω' : Type*} {mΩ : MeasurableSpace Ω} {mΩ' : MeasurableSpace Ω'}
+variable {Ω Ω' ι : Type*} {mΩ : MeasurableSpace Ω} {mΩ' : MeasurableSpace Ω'}
     {μ : Measure Ω} {ν : Measure Ω'} {X Y : Ω → ℝ} {Z W : Ω' → ℝ}
 
 lemma IdentDistrib.prod [IsFiniteMeasure μ] [IsFiniteMeasure ν]
@@ -24,4 +25,20 @@ lemma IdentDistrib.prod [IsFiniteMeasure μ] [IsFiniteMeasure ν]
       (indepFun_iff_map_prod_eq_prod_map_map hXZ.aemeasurable_snd hYW.aemeasurable_snd).mp hZW,
       hXZ.map_eq, hYW.map_eq]
 
+lemma IdentDistrib.pi [Countable ι] [IsProbabilityMeasure μ] [IsProbabilityMeasure ν]
+    {X : (i : ι) → Ω → ℝ} {Y : (i : ι) → Ω' → ℝ}
+    (h : ∀ i, IdentDistrib (X i) (Y i) μ ν) (hX_ind : iIndepFun X μ) (hY_ind : iIndepFun Y ν) :
+    IdentDistrib (fun ω ↦ (X · ω)) (fun ω ↦ (Y · ω)) μ ν where
+  aemeasurable_fst := by
+    rw [aemeasurable_pi_iff]
+    exact fun i ↦ (h i).aemeasurable_fst
+  aemeasurable_snd := by
+    rw [aemeasurable_pi_iff]
+    exact fun i ↦ (h i).aemeasurable_snd
+  map_eq := by
+    rw [(iIndepFun_iff_map_fun_eq_infinitePi_map₀' (fun i ↦ (h i).aemeasurable_fst)).mp hX_ind,
+      (iIndepFun_iff_map_fun_eq_infinitePi_map₀' (fun i ↦ (h i).aemeasurable_snd)).mp hY_ind]
+    congr with i
+    rw [(h i).map_eq]
+
 end ProbabilityTheory
diff --git a/LeanBandits/ForMathlib/Measurable.lean b/LeanBandits/ForMathlib/Measurable.lean
new file mode 100644
index 00000000..ae5ce178
--- /dev/null
+++ b/LeanBandits/ForMathlib/Measurable.lean
@@ -0,0 +1,22 @@
+/-
+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
+
+/-!
+# Measurability lemmas
+-/
+
+namespace MeasureTheory
+
+lemma measurable_comp_comap {α β γ : Type*} {mβ : MeasurableSpace β} {mγ : MeasurableSpace γ}
+    (f : α → β) {g : β → γ} (hg : Measurable g) :
+    Measurable[mβ.comap f] (g ∘ f) := by
+  rw [measurable_iff_comap_le, ← MeasurableSpace.comap_comp]
+  refine MeasurableSpace.comap_mono ?_
+  rw [← measurable_iff_comap_le]
+  exact hg
+
+end MeasureTheory
diff --git a/LeanBandits/RewardByCountMeasure.lean b/LeanBandits/RewardByCountMeasure.lean
index eb9e1ba2..7c482f1c 100644
--- a/LeanBandits/RewardByCountMeasure.lean
+++ b/LeanBandits/RewardByCountMeasure.lean
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Rémy Degenne
 -/
 import LeanBandits.Bandit
+import LeanBandits.ForMathlib.IdentDistrib
 import LeanBandits.Regret
 
 /-! # Laws of `stepsUntil` and `rewardByCount`
@@ -101,8 +102,6 @@ notation "𝓛[" Y " | " s "; " μ "]" => Measure.map Y (μ[|s])
 notation "𝓛[" Y " | " X " in " s "; " μ "]" => Measure.map Y (μ[|X ⁻¹' s])
 /-- Law of `Y` conditioned on the event that `X` equals `x`. -/
 notation "𝓛[" Y " | " X " ← " x "; " μ "]" => Measure.map Y (μ[|X ⁻¹' {x}])
-/-- Law of `Y` conditioned on `X`. -/
-notation "𝓛[" Y " | " X "; " μ "]" => condDistrib Y X μ
 
 omit [DecidableEq α] [MeasurableSingletonClass α] in
 lemma condDistrib_reward'' [StandardBorelSpace α] [Nonempty α] (n : ℕ) :
@@ -139,18 +138,52 @@ lemma reward_cond_arm [StandardBorelSpace α] [Nonempty α] [Countable α] (a :
   rw [h_ra] at h_eq
   exact h_eq.symm
 
-lemma condIndepFun_reward_stepsUntil_arm [StandardBorelSpace α] [Countable α] [Nonempty α]
+-- after the Mathlib stopping time refactor, we will be able to prove that stepsUntil is a
+-- stopping time
+lemma measurable_comap_indicator_stepsUntil_eq (a : α) (m n : ℕ) :
+    Measurable[MeasurableSpace.comap (fun ω : ℕ → α × ℝ ↦ (hist (n-1) ω, arm n ω)) inferInstance]
+      ({ω | stepsUntil a m ω = ↑n}.indicator fun _ ↦ 1) := by
+  let k : ((Iic (n - 1) → α × ℝ) × α) → (ℕ → α × ℝ) := fun x i ↦
+    if hi : i ∈ Iic (n - 1) then (x.1 ⟨i, hi⟩) else if i = n then (x.2, 0) else (a, 0)
+  have hk : Measurable k := by
+    unfold k
+    rw [measurable_pi_iff]
+    intro i
+    split_ifs <;> fun_prop
+  let φ : ((Iic (n - 1) → α × ℝ) × α) → ℕ := fun x ↦ if stepsUntil a m (k x) = ↑n then 1 else 0
+  have hφ : Measurable φ :=
+    Measurable.ite ((measurableSet_singleton _).preimage (by fun_prop)) (by fun_prop) (by fun_prop)
+  suffices {ω | stepsUntil a m ω = ↑n}.indicator (fun x ↦ 1)
+      = φ ∘ fun ω ↦ (hist (n - 1) ω, arm n ω) from this ▸ measurable_comp_comap _ hφ
+  ext ω
+  classical
+  simp only [Set.indicator_apply, Set.mem_setOf_eq, Function.comp_apply, φ]
+  congr 1
+  rw [stepsUntil_eq_congr]
+  intro i hin
+  simp only [arm, mem_Iic, hist, dite_eq_ite, k]
+  grind
+
+lemma measurable_indicator_stepsUntil_eq (a : α) (m n : ℕ) :
+    Measurable ({ω | stepsUntil a m ω = ↑n}.indicator fun _ ↦ 1) := by
+  refine (measurable_comap_indicator_stepsUntil_eq a m n).mono ?_ le_rfl
+  refine Measurable.comap_le ?_
+  fun_prop
+
+lemma measurableSet_stepsUntil_eq (a : α) (m n : ℕ) :
+    MeasurableSet[MeasurableSpace.comap (fun ω : ℕ → α × ℝ ↦ (hist (n-1) ω, arm n ω)) inferInstance]
+      {ω : ℕ → α × ℝ | stepsUntil a m ω = ↑n} := by
+  let mProd := MeasurableSpace.comap (fun ω : ℕ → α × ℝ ↦ (hist (n-1) ω, arm n ω)) inferInstance
+  suffices Measurable[mProd] ({ω | stepsUntil a m ω = ↑n}.indicator fun x ↦ 1) by
+    rwa [measurable_indicator_const_iff] at this
+  exact measurable_comap_indicator_stepsUntil_eq a m n
+
+lemma condIndepFun_reward_stepsUntil_arm' [StandardBorelSpace α] [Countable α] [Nonempty α]
     (a : α) (m n : ℕ) (hm : m ≠ 0) :
-    CondIndepFun (mα.comap (fun ω ↦ arm n ω.1)) ((measurable_arm n).comp measurable_fst).comap_le
-      (fun ω ↦ reward n ω.1) ({ω | stepsUntil a m ω.1 = ↑n}.indicator (fun _ ↦ 1))
-      (Bandit.measure alg ν) := by
-  -- first restrict to the `trajMeasure` side
-  suffices h_indep :
-      CondIndepFun (mα.comap (arm n)) (measurable_arm n).comap_le
-        (reward n) ({ω | stepsUntil a m ω = ↑n}.indicator (fun _ ↦ 1))
-        (Bandit.trajMeasure alg ν) by
-    sorry
-  -- Now prove the independence : the indicator of `stepsUntil ... = n` is a function of
+    CondIndepFun (mα.comap (arm n)) (measurable_arm n).comap_le
+      (reward n) ({ω | stepsUntil a m ω = ↑n}.indicator (fun _ ↦ 1))
+      (Bandit.trajMeasure alg ν) := by
+  -- the indicator of `stepsUntil ... = n` is a function of
   -- `hist (n-1)` and `arm n`.
   -- It thus suffices to prove the independence of `reward n` and `hist (n-1)` conditionally
   -- on `arm n`.
@@ -175,45 +208,25 @@ lemma condIndepFun_reward_stepsUntil_arm [StandardBorelSpace α] [Countable α]
   have h_indep' : CondIndepFun (mα.comap (arm n)) (measurable_arm n).comap_le (reward n)
       (fun ω ↦ (hist (n - 1) ω, arm n ω)) (Bandit.trajMeasure alg ν) :=
     h_indep.prod_right (by fun_prop) (by fun_prop) (by fun_prop)
-  suffices ∃ φ : ((Iic (n - 1) → α × ℝ) × α) → ℕ, Measurable φ ∧
-      ({ω : ℕ → α × ℝ | stepsUntil a m ω = ↑n}.indicator (fun _ ↦ 1))
-        = φ ∘ (fun ω : ℕ → α × ℝ ↦ (hist (n - 1) ω, arm n ω)) by
-    obtain ⟨φ, hφ_meas, h_eq⟩ := this
-    rw [h_eq]
-    exact h_indep'.comp measurable_id hφ_meas
-  -- it would follow from measurability wrt the sigma-algebra generated by
-  -- `hist (n-1)` and `arm n`, but we can also give an explicit function
-  let k : ((Iic (n - 1) → α × ℝ) × α) → (ℕ → α × ℝ) := fun x i ↦
-    if hi : i ∈ Iic (n - 1) then (x.1 ⟨i, hi⟩) else if i = n then (x.2, 0) else (a, 0)
-  let φ : ((Iic (n - 1) → α × ℝ) × α) → ℕ := fun x ↦ if stepsUntil a m (k x) = ↑n then 1 else 0
-  classical
-  have hφ_meas : Measurable φ := by
-    refine Measurable.ite ?_ (by fun_prop) (by fun_prop)
-    refine (measurableSet_singleton _).preimage ?_
-    refine (measurable_stepsUntil a m).comp ?_
-    unfold k
-    rw [measurable_pi_iff]
-    intro i
-    split_ifs <;> fun_prop
-  refine ⟨φ, hφ_meas, ?_⟩
-  ext ω
-  classical
-  simp only [Set.indicator_apply, Set.mem_setOf_eq, Function.comp_apply, φ]
-  congr 1
-  rw [stepsUntil_eq_congr]
-  intro i hin
-  simp only [arm, mem_Iic, hist, dite_eq_ite, k]
-  split_ifs with h1 h2
-  · rfl
-  · simp [h2]
-  · exfalso
-    grind
+  obtain ⟨φ, hφ_meas, h_eq⟩ : ∃ φ : ((Iic (n - 1) → α × ℝ) × α) → ℕ, Measurable φ ∧
+      {ω | stepsUntil a m ω = ↑n}.indicator (fun _ ↦ 1) = φ ∘ (fun ω ↦ (hist (n - 1) ω, arm n ω)) :=
+    (measurable_comap_indicator_stepsUntil_eq a m n).exists_eq_measurable_comp
+  rw [h_eq]
+  exact h_indep'.comp measurable_id hφ_meas
+
+lemma condIndepFun_reward_stepsUntil_arm [StandardBorelSpace α] [Countable α] [Nonempty α]
+    (a : α) (m n : ℕ) (hm : m ≠ 0) :
+    CondIndepFun (mα.comap (fun ω ↦ arm n ω.1)) ((measurable_arm n).comp measurable_fst).comap_le
+      (fun ω ↦ reward n ω.1) ({ω | stepsUntil a m ω.1 = ↑n}.indicator (fun _ ↦ 1))
+      (Bandit.measure alg ν) :=
+  condIndepFun_fst_prod (ν := Bandit.streamMeasure ν)
+    (measurable_indicator_stepsUntil_eq a m n) (by fun_prop) (by fun_prop)
+    (condIndepFun_reward_stepsUntil_arm' a m n hm)
 
 lemma reward_cond_stepsUntil [StandardBorelSpace α] [Countable α] [Nonempty α] (a : α) (m n : ℕ)
     (hm : m ≠ 0)
     (hμn : (Bandit.measure alg ν) ((fun ω ↦ stepsUntil a m ω.1) ⁻¹' {↑n}) ≠ 0) :
-    𝓛[fun ω ↦ reward n ω.1 | fun ω ↦ stepsUntil a m ω.1 ← (n : ℕ∞);
-      Bandit.measure alg ν] = ν a := by
+    𝓛[fun ω ↦ reward n ω.1 | fun ω ↦ stepsUntil a m ω.1 ← ↑n; Bandit.measure alg ν] = ν a := by
   let μ := Bandit.measure alg ν
   have hμna :
       μ ((fun ω ↦ stepsUntil a m ω.1) ⁻¹' {↑n} ∩ (fun ω ↦ arm n ω.1) ⁻¹' {a}) ≠ 0 := by
@@ -326,8 +339,36 @@ lemma identDistrib_rewardByCount_eval [Countable α] [StandardBorelSpace α] [No
       (Bandit.measure alg ν) (Bandit.streamMeasure ν) :=
   (identDistrib_rewardByCount_id a n hn).trans (identDistrib_eval_eval_id_streamMeasure ν m a).symm
 
+lemma iIndepFun_rewardByCount' (alg : Algorithm α ℝ) (ν : Kernel α ℝ) [IsMarkovKernel ν] (a : α) :
+    iIndepFun (fun n ω ↦ rewardByCount a n ω.1 ω.2) (Bandit.measure alg ν) := by
+  sorry
+
 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) (ν := ν)
+  · sorry
+  · exact iIndepFun_eval_streamMeasure'' ν a
+
+lemma identDistrib_rewardByCount_stream [Countable α] [StandardBorelSpace α] [Nonempty α]
+    (a : α) :
+    IdentDistrib (fun ω n ↦ rewardByCount a (n + 1) ω.1 ω.2) (fun ω n ↦ ω.2 n a)
+      (Bandit.measure alg ν) (Bandit.measure alg ν) := by
+  refine (identDistrib_rewardByCount_stream' a).trans ?_
+  refine IdentDistrib.pi (fun n ↦ ?_) ?_ ?_
+  · sorry
+  · sorry
+  · sorry
+
+lemma indepFun_rewardByCount_of_ne {a b : α} (hab : a ≠ b) :
+    IndepFun (fun ω s ↦ rewardByCount a s ω.1 ω.2) (fun ω s ↦ rewardByCount b s ω.1 ω.2)
+      (Bandit.measure alg ν) := by
+  sorry
+
 end Bandits