import Mathlib
import MaxLeafFormal.MLS
import MaxLeafFormal.CDS

/-!
# Budgeted trials (Theorem 3.14) and the connected-domination wrapper

Formalises the threshold lemma on `[0,∞]` (`lem:fw-threshold-prob`), budgeted trials
(`thm:fw-budgeted`), the first accepted size (`lem:cds-first-accepted`) and the
budgeted trials for connected domination (`cor:cds-budgeted`).
-/

namespace MaxLeafFormal.Budget

open MeasureTheory ProbabilityTheory Finset
open scoped ENNReal

/-- `2^{-x}` on `[0,∞]`, with `2^{-∞} = 0`. -/
noncomputable def twoPowNeg (x : ℝ≥0∞) : ℝ :=
  if x = ⊤ then 0 else (2:ℝ) ^ (-x.toReal)

/-- The event `{ξ ≤ θ}` for a real threshold `θ`. -/
def leEvent {Ξ : Type*} (ξ : Ξ → ℝ≥0∞) (θ : ℝ) : Set Ξ :=
  {x | ξ x ≠ ⊤ ∧ (ξ x).toReal ≤ θ}

/-! ### `twoPowNeg` -/

theorem measurable_twoPowNeg : Measurable twoPowNeg := by
  unfold twoPowNeg
  refine Measurable.ite (measurableSet_singleton ⊤) measurable_const ?_
  exact measurable_const.pow (ENNReal.measurable_toReal.neg)

theorem twoPowNeg_mem_Icc (x : ℝ≥0∞) : twoPowNeg x ∈ Set.Icc (0:ℝ) 1 := by
  unfold twoPowNeg
  split_ifs with h
  · norm_num
  · refine ⟨Real.rpow_nonneg (by norm_num) _, ?_⟩
    exact Real.rpow_le_one_of_one_le_of_nonpos (by norm_num)
      (neg_nonpos.2 ENNReal.toReal_nonneg)

theorem cap_condition_iff (E : ℝ) (hE : 0 < E) (Θ : ℝ) :
    (2:ℝ) ^ (-Θ) ≤ E / 2 ↔ 1 + Real.logb 2 (1 / E) ≤ Θ := by
  have hE2 : (0:ℝ) < E / 2 := by positivity
  have hlog : Real.logb 2 (E / 2) = -Real.logb 2 (1 / E) - 1 := by
    rw [one_div, Real.logb_inv, Real.logb_div hE.ne' (by norm_num),
      Real.logb_self_eq_one (by norm_num)]
    ring
  have h1 : (2:ℝ) ^ (-Θ) ≤ E / 2 ↔ (2:ℝ) ^ (-Θ) ≤ (2:ℝ) ^ (Real.logb 2 (E / 2)) := by
    rw [Real.rpow_logb (by norm_num) (by norm_num) hE2]
  rw [h1, Real.rpow_le_rpow_left_iff (by norm_num : (1:ℝ) < 2), hlog]
  constructor <;> intro h <;> linarith

theorem integral_twoPowNeg_pos {Ξ : Type*} [MeasurableSpace Ξ] (Q : Measure Ξ)
    [IsProbabilityMeasure Q] (ξ : Ξ → ℝ≥0∞) (hξ : Measurable ξ)
    (hfin : 0 < Q {x | ξ x < ⊤}) :
    0 < ∫ x, twoPowNeg (ξ x) ∂Q := by
  have hnn : 0 ≤ fun x => twoPowNeg (ξ x) := fun x => (twoPowNeg_mem_Icc (ξ x)).1
  have hint : Integrable (fun x => twoPowNeg (ξ x)) Q := by
    refine Integrable.of_bound ((measurable_twoPowNeg.comp hξ).aestronglyMeasurable) 1
      (Filter.Eventually.of_forall (fun x => ?_))
    rw [Real.norm_of_nonneg (twoPowNeg_mem_Icc (ξ x)).1]
    exact (twoPowNeg_mem_Icc (ξ x)).2
  rw [integral_pos_iff_support_of_nonneg hnn hint]
  have hs : Function.support (fun x => twoPowNeg (ξ x)) = {x | ξ x < ⊤} := by
    ext x
    have hpos : (2:ℝ) ^ (-(ξ x).toReal) ≠ 0 := (Real.rpow_pos_of_pos (by norm_num) _).ne'
    by_cases h : ξ x = ⊤
    · simp [Function.mem_support, twoPowNeg, h]
    · simp [Function.mem_support, twoPowNeg, h, hpos, lt_top_iff_ne_top]
  rw [hs]
  exact hfin

/-! ### `lem:fw-threshold-prob` -/

/-- Supporting form of (a) without `hfin` (used by `budgeted_success`). -/
theorem twoPowNeg_le_pointwise {Ξ : Type*} (ξ : Ξ → ℝ≥0∞) (Θ : ℕ) (x : Ξ) :
    twoPowNeg (ξ x) ≤ (2:ℝ)^(-(Θ:ℝ)) +
      ∑ θ ∈ range (Θ+1), 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x := by
  have hterm : ∀ θ : ℕ, 0 ≤ 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x :=
    fun θ => mul_nonneg (by norm_num)
      (Set.indicator_nonneg (fun _ _ => Real.rpow_nonneg (by norm_num) _) x)
  have hsum : 0 ≤ ∑ θ ∈ range (Θ+1), 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x :=
    Finset.sum_nonneg (fun θ _ => hterm θ)
  have hc0 : (0:ℝ) ≤ (2:ℝ)^(-(Θ:ℝ)) := Real.rpow_nonneg (by norm_num) _
  by_cases htop : ξ x = ⊤
  · have h0 : twoPowNeg (ξ x) = 0 := by
      unfold twoPowNeg
      rw [ite_eq_left htop]
    rw [h0]
    linarith
  · have hval : twoPowNeg (ξ x) = (2:ℝ)^(-(ξ x).toReal) := by
      unfold twoPowNeg
      rw [ite_eq_right htop]
    rw [hval]
    by_cases hr : (ξ x).toReal ≤ (Θ:ℝ)
    · have hc : ⌈(ξ x).toReal⌉₊ ≤ Θ := Nat.ceil_le.2 hr
      have hr0 : (0:ℝ) ≤ (ξ x).toReal := ENNReal.toReal_nonneg
      have hxmem : x ∈ leEvent ξ ((⌈(ξ x).toReal⌉₊ : ℕ) : ℝ) :=
        ⟨htop, Nat.le_ceil _⟩
      have hlt : ((⌈(ξ x).toReal⌉₊ : ℕ) : ℝ) < (ξ x).toReal + 1 :=
        Nat.ceil_lt_add_one hr0
      have hpow : (2:ℝ)^(-(ξ x).toReal) ≤ 2 * (2:ℝ)^(-((⌈(ξ x).toReal⌉₊ : ℕ) : ℝ)) := by
        have hid : (2:ℝ)^(-(ξ x).toReal) = 2 * (2:ℝ)^(-(ξ x).toReal - 1) := by
          rw [Real.rpow_sub (by norm_num), Real.rpow_one]
          ring
        rw [hid]
        refine mul_le_mul_of_nonneg_left ?_ (by norm_num)
        refine Real.rpow_le_rpow_of_exponent_le (by norm_num) ?_
        linarith
      have hsingle : 2 * (leEvent ξ ((⌈(ξ x).toReal⌉₊ : ℕ) : ℝ)).indicator
            (fun _ => (2:ℝ)^(-((⌈(ξ x).toReal⌉₊ : ℕ) : ℝ))) x
          ≤ ∑ θ ∈ range (Θ+1), 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x :=
        Finset.single_le_sum
          (f := fun θ : ℕ => 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x)
          (fun θ _ => hterm θ) (Finset.mem_range.2 (Nat.lt_succ_of_le hc))
      simp only [Set.indicator_of_mem hxmem] at hsingle
      linarith
    · push Not at hr
      have : (2:ℝ)^(-(ξ x).toReal) ≤ (2:ℝ)^(-(Θ:ℝ)) :=
        Real.rpow_le_rpow_of_exponent_le (by norm_num) (by linarith)
      linarith

theorem threshold_prob_aux {Ξ : Type*} [MeasurableSpace Ξ] (Q : Measure Ξ)
    [IsProbabilityMeasure Q] (ξ : Ξ → ℝ≥0∞) (hξ : Measurable ξ)
    (Θ : ℕ) (hΘ : (2:ℝ) ^ (-(Θ : ℝ)) ≤ (∫ x, twoPowNeg (ξ x) ∂Q) / 2) :
    ∃ θ : ℕ, θ ≤ Θ ∧
      (∫ x, twoPowNeg (ξ x) ∂Q) / (4 * ((Θ : ℝ) + 1))
        ≤ Q.real (leEvent ξ θ) * (2:ℝ) ^ (-(θ : ℝ)) := by
  have hmeasS : ∀ t : ℝ, MeasurableSet (leEvent ξ t) := fun t => by
    have hs : MeasurableSet {y : ℝ≥0∞ | y ≠ ⊤ ∧ y.toReal ≤ t} := by
      refine MeasurableSet.inter ((measurableSet_singleton ⊤).compl) ?_
      exact measurableSet_le ENNReal.measurable_toReal measurable_const
    exact hξ hs
  have hgint : Integrable (fun x => twoPowNeg (ξ x)) Q := by
    refine Integrable.of_bound ((measurable_twoPowNeg.comp hξ).aestronglyMeasurable) 1
      (Filter.Eventually.of_forall (fun x => ?_))
    rw [Real.norm_of_nonneg (twoPowNeg_mem_Icc (ξ x)).1]
    exact (twoPowNeg_mem_Icc (ξ x)).2
  have hterm : ∀ θ : ℕ, Integrable
      (fun x => 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x) Q :=
    fun θ => ((integrable_const _).indicator (hmeasS _)).const_mul 2
  have hRint : Integrable (fun x => (2:ℝ)^(-(Θ:ℝ)) +
      ∑ θ ∈ range (Θ+1), 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x) Q :=
    (integrable_const _).add (integrable_finsetSum _ (fun θ _ => hterm θ))
  have hE : ∫ x, twoPowNeg (ξ x) ∂Q ≤ ∫ x, ((2:ℝ)^(-(Θ:ℝ)) +
      ∑ θ ∈ range (Θ+1), 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x) ∂Q :=
    integral_mono hgint hRint (fun x => twoPowNeg_le_pointwise ξ Θ x)
  have hRval : ∫ x, ((2:ℝ)^(-(Θ:ℝ)) +
      ∑ θ ∈ range (Θ+1), 2 * (leEvent ξ (θ:ℝ)).indicator (fun _ => (2:ℝ)^(-(θ:ℝ))) x) ∂Q
      = (2:ℝ)^(-(Θ:ℝ)) + 2 * ∑ θ ∈ range (Θ+1), Q.real (leEvent ξ (θ:ℝ)) * (2:ℝ)^(-(θ:ℝ)) := by
    rw [integral_add (integrable_const _) (integrable_finsetSum _ (fun θ _ => hterm θ)),
      integral_const, probReal_univ, one_smul,
      integral_finsetSum _ (fun θ _ => hterm θ), Finset.mul_sum]
    congr 1
    apply Finset.sum_congr rfl
    intro θ _
    rw [integral_const_mul, integral_indicator_const _ (hmeasS _), smul_eq_mul]
  rw [hRval] at hE
  by_contra hcon
  push Not at hcon
  have hΘ1 : (0:ℝ) < (Θ:ℝ) + 1 := by positivity
  have hS : ∑ θ ∈ range (Θ+1), Q.real (leEvent ξ (θ:ℝ)) * (2:ℝ)^(-(θ:ℝ))
      < (∫ x, twoPowNeg (ξ x) ∂Q) / 4 := by
    calc ∑ θ ∈ range (Θ+1), Q.real (leEvent ξ (θ:ℝ)) * (2:ℝ)^(-(θ:ℝ))
        < ∑ θ ∈ range (Θ+1), (∫ x, twoPowNeg (ξ x) ∂Q) / (4 * ((Θ:ℝ)+1)) :=
          Finset.sum_lt_sum_of_nonempty ⟨0, by simp⟩
            (fun θ hθ => hcon θ (Nat.lt_succ_iff.1 (Finset.mem_range.1 hθ)))
      _ = (∫ x, twoPowNeg (ξ x) ∂Q) / 4 := by
          rw [Finset.sum_const, Finset.card_range, nsmul_eq_mul]
          push_cast
          field_simp
  linarith

-- lem:fw-threshold-prob (a)
theorem threshold_prob {Ξ : Type*} [MeasurableSpace Ξ] (Q : Measure Ξ)
    [IsProbabilityMeasure Q] (ξ : Ξ → ℝ≥0∞) (hξ : Measurable ξ)
    (_hfin : 0 < Q {x | ξ x < ⊤})
    (Θ : ℕ) (hΘ : (2:ℝ) ^ (-(Θ : ℝ)) ≤ (∫ x, twoPowNeg (ξ x) ∂Q) / 2) :
    ∃ θ : ℕ, θ ≤ Θ ∧
      (∫ x, twoPowNeg (ξ x) ∂Q) / (4 * ((Θ : ℝ) + 1))
        ≤ Q.real (leEvent ξ θ) * (2:ℝ) ^ (-(θ : ℝ)) :=
  threshold_prob_aux Q ξ hξ Θ hΘ

-- lem:fw-threshold-prob (b)
theorem threshold_prob_upper {Ξ : Type*} [MeasurableSpace Ξ] (Q : Measure Ξ)
    [IsProbabilityMeasure Q] (ξ : Ξ → ℝ≥0∞) (hξ : Measurable ξ)
    (_hfin : 0 < Q {x | ξ x < ⊤}) (θ : ℤ) :
    Q.real (leEvent ξ θ) * (2:ℝ) ^ (-(θ : ℝ)) ≤ ∫ x, twoPowNeg (ξ x) ∂Q := by
  have hmeas : MeasurableSet (leEvent ξ (θ : ℝ)) := by
    have hs : MeasurableSet {y : ℝ≥0∞ | y ≠ ⊤ ∧ y.toReal ≤ (θ : ℝ)} := by
      refine MeasurableSet.inter ((measurableSet_singleton ⊤).compl) ?_
      exact measurableSet_le ENNReal.measurable_toReal measurable_const
    exact hξ hs
  have hint : Integrable (fun x => twoPowNeg (ξ x)) Q := by
    refine Integrable.of_bound ((measurable_twoPowNeg.comp hξ).aestronglyMeasurable) 1
      (Filter.Eventually.of_forall (fun x => ?_))
    rw [Real.norm_of_nonneg (twoPowNeg_mem_Icc (ξ x)).1]
    exact (twoPowNeg_mem_Icc (ξ x)).2
  have h1 : ∫ x, (leEvent ξ (θ : ℝ)).indicator (fun _ => (2:ℝ) ^ (-(θ : ℝ))) x ∂Q
      = Q.real (leEvent ξ (θ : ℝ)) * (2:ℝ) ^ (-(θ : ℝ)) := by
    rw [integral_indicator_const _ hmeas, smul_eq_mul]
  rw [← h1]
  apply integral_mono ((integrable_const _).indicator hmeas) hint
  intro x
  by_cases hx : x ∈ leEvent ξ (θ : ℝ)
  · rw [Set.indicator_of_mem hx]
    have hx' : ξ x ≠ ⊤ ∧ (ξ x).toReal ≤ (θ : ℝ) := hx
    show _ ≤ (if ξ x = ⊤ then 0 else (2:ℝ) ^ (-(ξ x).toReal))
    rw [ite_eq_right hx'.1]
    exact Real.rpow_le_rpow_of_exponent_le (by norm_num) (neg_le_neg hx'.2)
  · rw [Set.indicator_of_notMem hx]
    exact (twoPowNeg_mem_Icc (ξ x)).1

/-! ### alg:fw-budgeted -/

abbrev BSlot (ι : Type*) (Θ : ℕ) (m : ι → ℕ → ℕ) : Type _ :=
  Σ p : ι × Fin (Θ + 1), Fin (m p.1 p.2)

abbrev BOutcome {ι : Type*} (Ω : ι → ℕ → Type*) (Θ : ℕ) (m : ι → ℕ → ℕ) : Type _ :=
  (j : BSlot ι Θ m) → Ω j.1.1 j.1.2

noncomputable def budgetMeasure {ι : Type*} [Fintype ι] {Ω : ι → ℕ → Type*}
    [∀ i θ, MeasurableSpace (Ω i θ)] (P : ∀ i θ, Measure (Ω i θ))
    (Θ : ℕ) (m : ι → ℕ → ℕ) : Measure (BOutcome Ω Θ m) :=
  Measure.pi fun j => P j.1.1 j.1.2

def BAccepts {ι : Type*} {Ω : ι → ℕ → Type*} (A : ∀ i θ, Set (Ω i θ))
    {Θ : ℕ} {m : ι → ℕ → ℕ} (o : BOutcome Ω Θ m) : Prop :=
  ∃ j, o j ∈ A j.1.1 j.1.2

/-- `M_{ι,θ}` of eq:fw-params. -/
noncomputable def Mcount (n : ℕ) (ω Γ κ : ℝ) (Θ θ : ℕ) : ℝ :=
  4 * κ * Γ * ((Θ : ℝ) + 1) * (2:ℝ) ^ (ω * n - θ)

/-- eq:fw-params, with `ι` in place of `t`. -/
def AdmissibleParams {ι : Type*} (n : ℕ) (ω Γ κ : ℝ) (Θ : ℕ) (m : ι → ℕ → ℕ) : Prop :=
  ω * n + Real.logb 2 Γ + 1 ≤ Θ ∧ (Θ : ℝ) ≤ ω * n + Real.logb 2 Γ + 3 ∧
  ∀ i θ, θ ≤ Θ → Mcount n ω Γ κ Θ θ ≤ m i θ ∧ (m i θ : ℝ) ≤ 2 * Mcount n ω Γ κ Θ θ + 1

-- thm:fw-budgeted, soundness (every outcome)
theorem budgeted_sound {U : Type*} (sol : Finset U → Prop) (k : ℕ)
    {ι : Type*} {Ω : ι → ℕ → Type*} (A : ∀ i θ, Set (Ω i θ))
    (hT2 : ∀ i θ, ∀ x ∈ A i θ, ∃ S, sol S ∧ S.card ≤ k)
    {Θ : ℕ} {m : ι → ℕ → ℕ} (o : BOutcome Ω Θ m) (h : BAccepts A o) :
    ∃ S, sol S ∧ S.card ≤ k := by
  obtain ⟨j, hj⟩ := h
  exact hT2 _ _ _ hj

/-- The acceptance event is measurable. -/
theorem measurableSet_bAccepts {ι : Type*} [Fintype ι] {Ω : ι → ℕ → Type*}
    [∀ i θ, MeasurableSpace (Ω i θ)] (A : ∀ i θ, Set (Ω i θ))
    (hA : ∀ i θ, MeasurableSet (A i θ)) (Θ : ℕ) (m : ι → ℕ → ℕ) :
    MeasurableSet {o : BOutcome Ω Θ m | BAccepts A o} := by
  classical
  have h : {o : BOutcome Ω Θ m | BAccepts A o}
      = ⋃ j : BSlot ι Θ m, (fun o : BOutcome Ω Θ m => o j) ⁻¹' A j.1.1 j.1.2 := by
    ext o
    simp only [Set.mem_ofPred_eq, Set.mem_iUnion, Set.mem_preimage, BAccepts]
  rw [h]
  exact MeasurableSet.iUnion fun j => measurable_pi_apply j (hA _ _)

-- Supporting: all `m i θ` copies of one trial fail
theorem budget_fail_le {ι : Type*} [Fintype ι]
    {Ω : ι → ℕ → Type*} [∀ i θ, MeasurableSpace (Ω i θ)]
    (P : ∀ i θ, Measure (Ω i θ)) [∀ i θ, IsProbabilityMeasure (P i θ)]
    (A : ∀ i θ, Set (Ω i θ)) (hA : ∀ i θ, MeasurableSet (A i θ))
    (Θ : ℕ) (m : ι → ℕ → ℕ) (i : ι) (θ : ℕ) (hθ : θ ≤ Θ) :
    (budgetMeasure P Θ m).real {o | ¬ BAccepts A o}
      ≤ (1 - (P i θ).real (A i θ)) ^ (m i θ) := by
  classical
  have hsub : {o : BOutcome Ω Θ m | ¬ BAccepts A o} ⊆
      Set.pi Set.univ (fun j : BSlot ι Θ m =>
        if j.1 = (i, (⟨θ, Nat.lt_succ_of_le hθ⟩ : Fin (Θ + 1))) then (A j.1.1 j.1.2)ᶜ
        else Set.univ) := by
    intro o ho j _
    dsimp only
    split_ifs with hj
    · exact fun h => ho ⟨j, h⟩
    · exact Set.mem_univ _
  have hc : ((P i θ) (A i θ)ᶜ).toReal = 1 - ((P i θ) (A i θ)).toReal := by
    rw [prob_compl_eq_one_sub (hA _ _), ENNReal.toReal_sub_of_le prob_le_one ENNReal.one_ne_top,
      ENNReal.toReal_one]
  have hB : (budgetMeasure P Θ m).real (Set.pi Set.univ (fun j : BSlot ι Θ m =>
        if j.1 = (i, (⟨θ, Nat.lt_succ_of_le hθ⟩ : Fin (Θ + 1))) then (A j.1.1 j.1.2)ᶜ
        else Set.univ))
      = (1 - (P i θ).real (A i θ)) ^ (m i θ) := by
    unfold budgetMeasure
    rw [measureReal_def, Measure.pi_pi, ENNReal.toReal_prod, Fintype.prod_sigma,
      Fintype.prod_eq_single (i, (⟨θ, Nat.lt_succ_of_le hθ⟩ : Fin (Θ + 1)))]
    · simp only [ite_true]
      rw [Finset.prod_congr rfl (fun _ _ => hc), Finset.prod_const, Finset.card_univ,
        Fintype.card_fin]
      rfl
    · intro a ha
      refine Finset.prod_eq_one fun b _ => ?_
      simp [ha]
  have : IsFiniteMeasure (budgetMeasure P Θ m) := by
    unfold budgetMeasure
    infer_instance
  calc (budgetMeasure P Θ m).real {o | ¬ BAccepts A o}
      ≤ (budgetMeasure P Θ m).real (Set.pi Set.univ (fun j : BSlot ι Θ m =>
        if j.1 = (i, (⟨θ, Nat.lt_succ_of_le hθ⟩ : Fin (Θ + 1))) then (A j.1.1 j.1.2)ᶜ
        else Set.univ)) := measureReal_mono hsub (measure_ne_top _ _)
    _ = _ := hB

-- thm:fw-budgeted, success probability
theorem budgeted_success {ι : Type*} [Fintype ι]
    {Ω : ι → ℕ → Type*} [∀ i θ, MeasurableSpace (Ω i θ)]
    (P : ∀ i θ, Measure (Ω i θ)) [∀ i θ, IsProbabilityMeasure (P i θ)]
    (A : ∀ i θ, Set (Ω i θ)) (hA : ∀ i θ, MeasurableSet (A i θ))
    {Ξ : ι → Type*} [∀ i, MeasurableSpace (Ξ i)]
    (Q : ∀ i, Measure (Ξ i)) [∀ i, IsProbabilityMeasure (Q i)]
    (ξ : ∀ i, Ξ i → ℝ≥0∞) (hξ : ∀ i, Measurable (ξ i))
    (hT3 : ∀ i (θ : ℕ), (Q i).real (leEvent (ξ i) θ) ≤ (P i θ).real (A i θ))
    (n : ℕ) (ω Γ κ : ℝ) (_hω : 0 ≤ ω) (hΓ : 1 ≤ Γ) (hκ : 0 < κ)
    (Θ : ℕ) (hΘ : ω * n + Real.logb 2 Γ + 1 ≤ Θ)
    (m : ι → ℕ → ℕ) (hm : ∀ i θ, θ ≤ Θ → Mcount n ω Γ κ Θ θ ≤ m i θ)
    (hsucc : ∃ i, (2:ℝ) ^ (-(ω * n)) / Γ ≤ ∫ x, twoPowNeg (ξ i x) ∂(Q i)) :
    1 - Real.exp (-κ) ≤ (budgetMeasure P Θ m).real {o | BAccepts A o} := by
  obtain ⟨i, hi⟩ := hsucc
  set E := ∫ x, twoPowNeg (ξ i x) ∂(Q i) with hE
  have hΓ0 : 0 < Γ := by linarith
  -- the cap is high enough
  have hcap : (2:ℝ) ^ (-(Θ : ℝ)) ≤ E / 2 := by
    calc (2:ℝ) ^ (-(Θ : ℝ)) ≤ (2:ℝ) ^ (-(ω * n + Real.logb 2 Γ + 1)) := by
          apply Real.rpow_le_rpow_of_exponent_le (by norm_num)
          linarith
      _ = (2:ℝ) ^ (-(ω * n)) / Γ / 2 := by
          rw [show -(ω * n + Real.logb 2 Γ + 1) = -(ω * n) + -Real.logb 2 Γ + -1 by ring,
            Real.rpow_add (by norm_num), Real.rpow_add (by norm_num),
            Real.rpow_neg (by norm_num : (0:ℝ) ≤ 2) (Real.logb 2 Γ),
            Real.rpow_logb (by norm_num) (by norm_num) hΓ0, Real.rpow_neg_one]
          ring
      _ ≤ E / 2 := by linarith
  obtain ⟨θ, hθΘ, hθ⟩ := threshold_prob_aux (Q i) (ξ i) (hξ i) Θ hcap
  set q := (Q i).real (leEvent (ξ i) θ) with hq
  set p := (P i θ).real (A i θ) with hp
  have hqp : q ≤ p := hT3 i θ
  have hq0 : 0 ≤ q := measureReal_nonneg
  have hp1 : p ≤ 1 := measureReal_le_one
  have hM := hm i θ hθΘ
  have hΘ1 : (0:ℝ) < Θ + 1 := by positivity
  have hinv : (2:ℝ) ^ (ω * n) * 2 ^ (-(ω * n)) = 1 := by
    rw [← Real.rpow_add (by norm_num)]; simp
  have hsplit : (2:ℝ) ^ (ω * n - θ) = 2 ^ (ω * n) * 2 ^ (-(θ : ℝ)) := by
    rw [sub_eq_add_neg, Real.rpow_add (by norm_num)]
  -- `κ ≤ p * m`
  have hpm : κ ≤ p * m i θ := by
    have h1 : κ ≤ E * κ * Γ * 2 ^ (ω * n) := by
      have := mul_le_mul_of_nonneg_right hi
        (by positivity : (0:ℝ) ≤ κ * Γ * 2 ^ (ω * n))
      calc κ = (2:ℝ) ^ (-(ω * n)) / Γ * (κ * Γ * 2 ^ (ω * n)) := by
            rw [show (2:ℝ) ^ (-(ω * n)) / Γ * (κ * Γ * 2 ^ (ω * n))
                = κ * (2 ^ (ω * n) * 2 ^ (-(ω * n))) by field_simp, hinv, mul_one]
        _ ≤ E * (κ * Γ * 2 ^ (ω * n)) := this
        _ = E * κ * Γ * 2 ^ (ω * n) := by ring
    have h2 : E * κ * Γ * 2 ^ (ω * n) ≤ q * Mcount n ω Γ κ Θ θ := by
      have := mul_le_mul_of_nonneg_right hθ
        (by positivity : (0:ℝ) ≤ 4 * κ * Γ * ((Θ : ℝ) + 1) * 2 ^ (ω * n))
      calc E * κ * Γ * 2 ^ (ω * n)
          = E / (4 * ((Θ : ℝ) + 1)) * (4 * κ * Γ * ((Θ : ℝ) + 1) * 2 ^ (ω * n)) := by
            field_simp
        _ ≤ q * 2 ^ (-(θ : ℝ)) * (4 * κ * Γ * ((Θ : ℝ) + 1) * 2 ^ (ω * n)) := this
        _ = q * Mcount n ω Γ κ Θ θ := by
            unfold Mcount; rw [hsplit]; ring
    have h3 : q * Mcount n ω Γ κ Θ θ ≤ p * m i θ :=
      mul_le_mul hqp hM (by unfold Mcount; positivity) (hq0.trans hqp)
    linarith
  -- failure probability
  have hfail := budget_fail_le P A hA Θ m i θ hθΘ
  have hexp : (1 - p) ^ (m i θ) ≤ Real.exp (-κ) := by
    calc (1 - p) ^ (m i θ) ≤ (Real.exp (-p)) ^ (m i θ) := by
          apply pow_le_pow_left₀ (by linarith)
          linarith [Real.add_one_le_exp (-p)]
      _ = Real.exp (-(p * m i θ)) := by rw [← Real.exp_nat_mul]; ring_nf
      _ ≤ Real.exp (-κ) := Real.exp_le_exp.2 (by linarith)
  have hprob : IsProbabilityMeasure (budgetMeasure P Θ m) := by
    unfold budgetMeasure; infer_instance
  have hcompl := measureReal_add_measureReal_compl (μ := budgetMeasure P Θ m)
    (measurableSet_bAccepts A hA Θ m)
  have hc : ({o : BOutcome Ω Θ m | BAccepts A o})ᶜ = {o | ¬ BAccepts A o} := rfl
  rw [hc, probReal_univ] at hcompl
  linarith

-- thm:fw-budgeted, success probability, literal form
theorem budgeted_success_of_admissible {ι : Type*} [Fintype ι]
    {Ω : ι → ℕ → Type*} [∀ i θ, MeasurableSpace (Ω i θ)]
    (P : ∀ i θ, Measure (Ω i θ)) [∀ i θ, IsProbabilityMeasure (P i θ)]
    (A : ∀ i θ, Set (Ω i θ)) (hA : ∀ i θ, MeasurableSet (A i θ))
    {Ξ : ι → Type*} [∀ i, MeasurableSpace (Ξ i)]
    (Q : ∀ i, Measure (Ξ i)) [∀ i, IsProbabilityMeasure (Q i)]
    (ξ : ∀ i, Ξ i → ℝ≥0∞) (hξ : ∀ i, Measurable (ξ i))
    (hT3 : ∀ i (θ : ℕ), (Q i).real (leEvent (ξ i) θ) ≤ (P i θ).real (A i θ))
    (n : ℕ) (ω Γ κ : ℝ) (hω : 0 ≤ ω) (hΓ : 1 ≤ Γ) (hκ : 0 < κ)
    (Θ : ℕ) (m : ι → ℕ → ℕ) (hpar : AdmissibleParams n ω Γ κ Θ m)
    (hsucc : ∃ i, (2:ℝ) ^ (-(ω * n)) / Γ ≤ ∫ x, twoPowNeg (ξ i x) ∂(Q i)) :
    1 - Real.exp (-κ) ≤ (budgetMeasure P Θ m).real {o | BAccepts A o} :=
  budgeted_success P A hA Q ξ hξ hT3 n ω Γ κ hω hΓ hκ Θ hpar.1 m
    (fun i θ hθ => (hpar.2.2 i θ hθ).1) hsucc

theorem params_admissible (ι : Type*) (n : ℕ) (ω Γ κ : ℝ)
    (hω : 0 ≤ ω) (hΓ : 1 ≤ Γ) (hκ : 0 < κ) :
    AdmissibleParams (ι := ι) n ω Γ κ (MLS.cap n ω Γ) (fun _ θ => MLS.trials n ω Γ κ θ) := by
  have hx : 0 ≤ ω * n + Real.logb 2 Γ :=
    add_nonneg (by positivity) (Real.logb_nonneg (by norm_num) hΓ)
  have hΓ0 : 0 < Γ := by linarith
  have hc1 := Nat.le_ceil (ω * n + Real.logb 2 Γ)
  have hc2 := Nat.ceil_lt_add_one hx
  have hcap : ((MLS.cap n ω Γ : ℕ) : ℝ) = (⌈ω * n + Real.logb 2 Γ⌉₊ : ℝ) + 1 := by
    unfold MLS.cap
    push_cast
    ring
  have hM : ∀ θ : ℕ, 0 ≤ Mcount n ω Γ κ (MLS.cap n ω Γ) θ := fun θ => by
    unfold Mcount
    exact mul_nonneg (mul_nonneg (mul_nonneg (mul_nonneg (by norm_num) hκ.le) hΓ0.le)
      (by positivity)) (by positivity)
  have hT : ∀ θ : ℕ, MLS.trials n ω Γ κ θ = ⌈Mcount n ω Γ κ (MLS.cap n ω Γ) θ⌉₊ := by
    intro θ
    unfold MLS.trials Mcount
    congr 1
  unfold AdmissibleParams
  refine ⟨?_, ?_, fun i θ _ => ⟨?_, ?_⟩⟩
  · rw [hcap]
    linarith
  · rw [hcap]
    linarith
  · dsimp only
    rw [hT θ]
    exact Nat.le_ceil _
  · dsimp only
    rw [hT θ]
    have := Nat.ceil_lt_add_one (hM θ)
    linarith [hM θ]

/-! ### lem:cds-first-accepted -/

open Classical in
noncomputable def firstAccepted (n : ℕ) (acc : ℕ → Prop) : ℕ :=
  if h : ((Finset.Icc 1 n).filter acc).Nonempty
  then ((Finset.Icc 1 n).filter acc).min' h else n

theorem first_accepted {U : Type*} [Fintype U] (sol : Finset U → Prop)
    (hne : ∃ S, sol S) (hempty : ¬ sol ∅)
    {Ω : Type*} [MeasurableSpace Ω] (P : Measure Ω) [IsProbabilityMeasure P]
    (acc : ℕ → Ω → Prop)
    (hsound : ∀ s o, 1 ≤ s → s ≤ Fintype.card U → acc s o → ∃ S, sol S ∧ S.card ≤ s) :
    let opt := sInf {s | ∃ S, sol S ∧ S.card = s}
    (∀ o, opt ≤ firstAccepted (Fintype.card U) (fun s => acc s o)) ∧
    P.real {o | acc opt o}
      ≤ P.real {o | firstAccepted (Fintype.card U) (fun s => acc s o) = opt} := by
  intro opt
  classical
  have hmem : ∃ s, ∃ S, sol S ∧ S.card = s := by
    obtain ⟨S, hS⟩ := hne
    exact ⟨S.card, S, hS, rfl⟩
  have hopt_mem : opt ∈ {s | ∃ S, sol S ∧ S.card = s} := Nat.sInf_mem hmem
  obtain ⟨S0, hS0, hS0card⟩ := hopt_mem
  have hopt_eq : opt = S0.card := hS0card.symm
  have hS0ne : S0 ≠ ∅ := by
    rintro rfl
    exact hempty hS0
  have hopt_pos : 1 ≤ opt := by
    rw [hopt_eq]
    exact Finset.card_pos.mpr (Finset.nonempty_iff_ne_empty.mpr hS0ne)
  have hopt_le : opt ≤ Fintype.card U := by
    rw [hopt_eq]
    exact Finset.card_le_univ S0
  have key : ∀ s o, 1 ≤ s → s ≤ Fintype.card U → acc s o → opt ≤ s := by
    intro s o h1 h2 h3
    obtain ⟨S', hS', hcard⟩ := hsound s o h1 h2 h3
    exact le_trans (Nat.sInf_le ⟨S', hS', rfl⟩) hcard
  have hfa : ∀ o, opt ≤ firstAccepted (Fintype.card U) (fun s => acc s o) := by
    intro o
    unfold firstAccepted
    split_ifs with h
    · have hm := Finset.min'_mem _ h
      rw [Finset.mem_filter, Finset.mem_Icc] at hm
      exact key _ o hm.1.1 hm.1.2 hm.2
    · exact hopt_le
  have hpt : ∀ o, acc opt o → firstAccepted (Fintype.card U) (fun s => acc s o) = opt := by
    intro o ho
    have hopt_mem' : opt ∈ (Finset.Icc 1 (Fintype.card U)).filter (fun s => acc s o) :=
      Finset.mem_filter.mpr ⟨Finset.mem_Icc.mpr ⟨hopt_pos, hopt_le⟩, ho⟩
    have hne' : ((Finset.Icc 1 (Fintype.card U)).filter (fun s => acc s o)).Nonempty :=
      ⟨opt, hopt_mem'⟩
    have hge := hfa o
    unfold firstAccepted at hge ⊢
    rw [dite_eq_left hne'] at hge ⊢
    exact le_antisymm (Finset.min'_le _ _ hopt_mem') hge
  refine ⟨hfa, ?_⟩
  exact measureReal_mono (fun o ho => hpt o ho) (measure_ne_top P _)

/-! ### cor:cds-budgeted -/

abbrev RunOutcome (Ω : ℕ → ℕ → ℕ → Type*) (Θ : ℕ → ℕ) (m : ℕ → ℕ → ℕ → ℕ)
    (s : ℕ) : Type _ :=
  BOutcome (ι := Fin (s + 1)) (fun t θ => Ω s t θ) (Θ s) (fun t θ => m s t θ)

/-- One independent run for each `s = i + 1`, `i : Fin n`. -/
abbrev CdsOutcome (n : ℕ) (Ω : ℕ → ℕ → ℕ → Type*) (Θ : ℕ → ℕ)
    (m : ℕ → ℕ → ℕ → ℕ) : Type _ :=
  (i : Fin n) → RunOutcome Ω Θ m ((i : ℕ) + 1)

noncomputable def cdsMeasure (n : ℕ) {Ω : ℕ → ℕ → ℕ → Type*}
    [∀ s t θ, MeasurableSpace (Ω s t θ)] (P : ∀ s t θ, Measure (Ω s t θ))
    (Θ : ℕ → ℕ) (m : ℕ → ℕ → ℕ → ℕ) : Measure (CdsOutcome n Ω Θ m) :=
  Measure.pi fun i =>
    budgetMeasure (fun (t : Fin ((i : ℕ) + 1 + 1)) θ => P ((i : ℕ) + 1) t θ)
      (Θ ((i : ℕ) + 1)) (fun t θ => m ((i : ℕ) + 1) t θ)

def cdsAcc {n : ℕ} {Ω : ℕ → ℕ → ℕ → Type*} {Θ : ℕ → ℕ} {m : ℕ → ℕ → ℕ → ℕ}
    (A : ∀ s t θ, Set (Ω s t θ)) (o : CdsOutcome n Ω Θ m) (s : ℕ) : Prop :=
  ∃ i : Fin n, (i : ℕ) + 1 = s ∧
    BAccepts (fun (t : Fin ((i : ℕ) + 1 + 1)) θ => A ((i : ℕ) + 1) t θ) (o i)

noncomputable def cdsOutput (n : ℕ) {Ω : ℕ → ℕ → ℕ → Type*} {Θ : ℕ → ℕ}
    {m : ℕ → ℕ → ℕ → ℕ} (A : ∀ s t θ, Set (Ω s t θ))
    (o : CdsOutcome n Ω Θ m) : ℕ :=
  firstAccepted n (cdsAcc A o)

open MaxLeafFormal.CDS

/-- Helper: the whole vertex set is a CDS of a connected graph. -/
theorem isCDS_univ {V : Type*} [Fintype V] [DecidableEq V] (G : SimpleGraph V)
    (hG : G.Connected) : IsCDS G univ := by
  refine ⟨?_, fun v => Or.inl (Finset.mem_univ v)⟩
  rw [Finset.coe_univ]
  exact ((SimpleGraph.induceUnivIso G).connected_iff).mpr hG

theorem cds_budgeted {V : Type*} [Fintype V] [DecidableEq V]
    (G : SimpleGraph V) [DecidableRel G.Adj] (hG : G.Connected)
    {Ω : ℕ → ℕ → ℕ → Type*} [∀ s t θ, MeasurableSpace (Ω s t θ)]
    (P : ∀ s t θ, Measure (Ω s t θ))
    (hP : ∀ s t θ, 1 ≤ s → s ≤ Fintype.card V → t ≤ s → IsProbabilityMeasure (P s t θ))
    (A : ∀ s t θ, Set (Ω s t θ))
    (hA : ∀ s t θ, 1 ≤ s → s ≤ Fintype.card V → t ≤ s → MeasurableSet (A s t θ))
    (hT2 : ∀ s t θ, 1 ≤ s → s ≤ Fintype.card V → t ≤ s → ∀ x ∈ A s t θ, γc G ≤ s)
    (ω Γ κ : ℝ) (hω : 0 ≤ ω) (hΓ : 1 ≤ Γ) (hκ : 0 < κ)
    (Θ : ℕ → ℕ)
    (hΘ : ∀ s, 1 ≤ s → s ≤ Fintype.card V → ω * Fintype.card V + Real.logb 2 Γ + 1 ≤ Θ s)
    (m : ℕ → ℕ → ℕ → ℕ)
    (hm : ∀ s t θ, 1 ≤ s → s ≤ Fintype.card V → t ≤ s → θ ≤ Θ s →
      Mcount (Fintype.card V) ω Γ κ (Θ s) θ ≤ m s t θ)
    {Ξ : ℕ → Type*} [∀ t, MeasurableSpace (Ξ t)] (Q : ∀ t, Measure (Ξ t))
    (hQ : ∀ t ≤ γc G, IsProbabilityMeasure (Q t))
    (ξ : ∀ t, Ξ t → ℝ≥0∞) (hξ : ∀ t ≤ γc G, Measurable (ξ t))
    (hT3 : ∀ t (θ : ℕ), t ≤ γc G →
      (Q t).real (leEvent (ξ t) θ) ≤ (P (γc G) t θ).real (A (γc G) t θ))
    (hmax : ∃ t ≤ γc G,
      (2:ℝ) ^ (-(ω * Fintype.card V)) / Γ ≤ ∫ x, twoPowNeg (ξ t x) ∂(Q t)) :
    (∀ o : CdsOutcome (Fintype.card V) Ω Θ m, γc G ≤ cdsOutput (Fintype.card V) A o) ∧
    1 - Real.exp (-κ)
      ≤ (cdsMeasure (Fintype.card V) P Θ m).real
          {o | cdsOutput (Fintype.card V) A o = γc G} := by
  set n := Fintype.card V with hn
  -- a minimum CDS
  have hne : {s | ∃ S, IsCDS G S ∧ S.card = s}.Nonempty := ⟨_, _, isCDS_univ G hG, rfl⟩
  obtain ⟨S0, hS0, hS0c⟩ := Nat.sInf_mem hne
  change S0.card = γc G at hS0c
  have hempty : ¬ IsCDS G (∅ : Finset V) := by
    rintro ⟨hc, -⟩
    obtain ⟨v⟩ := hc.nonempty
    simpa using v.2
  have hγ1 : 1 ≤ γc G := by
    rw [← hS0c, Nat.one_le_iff_ne_zero]
    intro h0
    rw [Finset.card_eq_zero] at h0
    exact hempty (h0 ▸ hS0)
  have hγn : γc G ≤ n := hS0c ▸ Finset.card_le_univ S0
  -- the product measure is a probability measure
  have hPi : ∀ (i : Fin n) (t : Fin ((i : ℕ) + 1 + 1)) (θ : ℕ),
      IsProbabilityMeasure (P ((i : ℕ) + 1) t θ) := fun i t θ =>
    hP _ _ _ (by omega) (by omega) (by omega)
  have hB : ∀ i : Fin n, IsProbabilityMeasure
      (budgetMeasure (fun (t : Fin ((i : ℕ) + 1 + 1)) θ => P ((i : ℕ) + 1) t θ)
        (Θ ((i : ℕ) + 1)) (fun t θ => m ((i : ℕ) + 1) t θ)) := fun i => by
    unfold budgetMeasure; infer_instance
  have hprob : IsProbabilityMeasure (cdsMeasure n P Θ m) := by
    unfold cdsMeasure budgetMeasure; infer_instance
  -- first accepted size
  have hsound : ∀ s (o : CdsOutcome n Ω Θ m), 1 ≤ s → s ≤ Fintype.card V →
      cdsAcc A o s → ∃ S, IsCDS G S ∧ S.card ≤ s := by
    rintro s o - - ⟨i, hi, j, hj⟩
    have := hT2 _ _ _ (by omega) (by omega) (by omega) _ hj
    exact ⟨S0, hS0, by omega⟩
  obtain ⟨hlow, hacc⟩ := first_accepted (IsCDS G) ⟨_, isCDS_univ G hG⟩ hempty
    (cdsMeasure n P Θ m) (fun s o => cdsAcc A o s) hsound
  refine ⟨hlow, le_trans ?_ hacc⟩
  -- the run at `s = γc G`
  set i0 : Fin n := ⟨γc G - 1, by omega⟩ with hi0
  have hs : (i0 : ℕ) + 1 = γc G := by simp only [hi0]; omega
  have hsub : {o : CdsOutcome n Ω Θ m |
      BAccepts (fun (t : Fin ((i0 : ℕ) + 1 + 1)) θ => A ((i0 : ℕ) + 1) t θ) (o i0)}
      ⊆ {o | cdsAcc A o (γc G)} := fun o ho => ⟨i0, hs, ho⟩
  refine le_trans ?_ (measureReal_mono hsub)
  have hmeasA : MeasurableSet {o' : RunOutcome Ω Θ m ((i0 : ℕ) + 1) |
      BAccepts (fun (t : Fin ((i0 : ℕ) + 1 + 1)) θ => A ((i0 : ℕ) + 1) t θ) o'} :=
    measurableSet_bAccepts _ (fun t θ => hA _ _ _ (by omega) (by omega) (by omega)) _ _
  have hpres := measurePreserving_eval
    (fun i : Fin n => budgetMeasure (fun (t : Fin ((i : ℕ) + 1 + 1)) θ => P ((i : ℕ) + 1) t θ)
      (Θ ((i : ℕ) + 1)) (fun t θ => m ((i : ℕ) + 1) t θ)) i0
  have heq : (cdsMeasure n P Θ m).real {o : CdsOutcome n Ω Θ m |
      BAccepts (fun (t : Fin ((i0 : ℕ) + 1 + 1)) θ => A ((i0 : ℕ) + 1) t θ) (o i0)}
      = (budgetMeasure (fun (t : Fin ((i0 : ℕ) + 1 + 1)) θ => P ((i0 : ℕ) + 1) t θ)
          (Θ ((i0 : ℕ) + 1)) (fun t θ => m ((i0 : ℕ) + 1) t θ)).real
          {o' | BAccepts (fun (t : Fin ((i0 : ℕ) + 1 + 1)) θ => A ((i0 : ℕ) + 1) t θ) o'} := by
    unfold Measure.real
    rw [← hpres.measure_preimage hmeasA.nullMeasurableSet]
    rfl
  rw [heq]
  -- budgeted trials at `s = i0 + 1`
  rw [← hs] at hT3 hQ hξ hmax
  have : ∀ t : Fin ((i0 : ℕ) + 1 + 1), IsProbabilityMeasure (Q t) := fun t =>
    hQ t (by omega)
  obtain ⟨t0, ht0, hmax0⟩ := hmax
  exact budgeted_success (fun (t : Fin ((i0 : ℕ) + 1 + 1)) θ => P ((i0 : ℕ) + 1) t θ)
    (fun t θ => A ((i0 : ℕ) + 1) t θ) (fun t θ => hA _ _ _ (by omega) (by omega) (by omega))
    (Ξ := fun t => Ξ t) (fun t => Q t) (fun t => ξ t) (fun t => hξ t (by omega))
    (fun t θ => hT3 t θ (by omega)) n ω Γ κ hω hΓ hκ (Θ ((i0 : ℕ) + 1))
    (hΘ _ (by omega) (by omega)) (fun t θ => m ((i0 : ℕ) + 1) t θ)
    (fun t θ hθ => hm _ _ _ (by omega) (by omega) (by omega) hθ)
    ⟨⟨t0, by omega⟩, hmax0⟩

theorem cds_budgeted_maxLeaf {V : Type*} [Fintype V] [DecidableEq V]
    (G : SimpleGraph V) [DecidableRel G.Adj] (hG : G.Connected)
    {Ω : ℕ → ℕ → ℕ → Type*} [∀ s t θ, MeasurableSpace (Ω s t θ)]
    (P : ∀ s t θ, Measure (Ω s t θ))
    (hP : ∀ s t θ, 1 ≤ s → s ≤ Fintype.card V → t ≤ s → IsProbabilityMeasure (P s t θ))
    (A : ∀ s t θ, Set (Ω s t θ))
    (hA : ∀ s t θ, 1 ≤ s → s ≤ Fintype.card V → t ≤ s → MeasurableSet (A s t θ))
    (hT2 : ∀ s t θ, 1 ≤ s → s ≤ Fintype.card V → t ≤ s → ∀ x ∈ A s t θ, γc G ≤ s)
    (ω Γ κ : ℝ) (hω : 0 ≤ ω) (hΓ : 1 ≤ Γ) (hκ : 0 < κ)
    (Θ : ℕ → ℕ)
    (hΘ : ∀ s, 1 ≤ s → s ≤ Fintype.card V → ω * Fintype.card V + Real.logb 2 Γ + 1 ≤ Θ s)
    (m : ℕ → ℕ → ℕ → ℕ)
    (hm : ∀ s t θ, 1 ≤ s → s ≤ Fintype.card V → t ≤ s → θ ≤ Θ s →
      Mcount (Fintype.card V) ω Γ κ (Θ s) θ ≤ m s t θ)
    {Ξ : ℕ → Type*} [∀ t, MeasurableSpace (Ξ t)] (Q : ∀ t, Measure (Ξ t))
    (hQ : ∀ t ≤ γc G, IsProbabilityMeasure (Q t))
    (ξ : ∀ t, Ξ t → ℝ≥0∞) (hξ : ∀ t ≤ γc G, Measurable (ξ t))
    (hT3 : ∀ t (θ : ℕ), t ≤ γc G →
      (Q t).real (leEvent (ξ t) θ) ≤ (P (γc G) t θ).real (A (γc G) t θ))
    (hmax : ∃ t ≤ γc G,
      (2:ℝ) ^ (-(ω * Fintype.card V)) / Γ ≤ ∫ x, twoPowNeg (ξ t x) ∂(Q t))
    (hn : 3 ≤ Fintype.card V) :
    1 - Real.exp (-κ)
      ≤ (cdsMeasure (Fintype.card V) P Θ m).real
          {o | IsGreatest {ℓ | ∃ F : SimpleGraph V, F ≤ G ∧ F.IsTree ∧ numLeaves F = ℓ}
                 (Fintype.card V - cdsOutput (Fintype.card V) A o)} := by
  have hPi : ∀ (i : Fin (Fintype.card V)) (t : Fin ((i : ℕ) + 1 + 1)) (θ : ℕ),
      IsProbabilityMeasure (P ((i : ℕ) + 1) t θ) := fun i t θ =>
    hP _ _ _ (by omega) (by omega) (by omega)
  have hB : ∀ i : Fin (Fintype.card V), IsProbabilityMeasure
      (budgetMeasure (fun (t : Fin ((i : ℕ) + 1 + 1)) θ => P ((i : ℕ) + 1) t θ)
        (Θ ((i : ℕ) + 1)) (fun t θ => m ((i : ℕ) + 1) t θ)) := fun i => by
    unfold budgetMeasure; infer_instance
  have hprob : IsProbabilityMeasure (cdsMeasure (Fintype.card V) P Θ m) := by
    unfold cdsMeasure; infer_instance
  obtain ⟨-, h⟩ := cds_budgeted G hG P hP A hA hT2 ω Γ κ hω hΓ hκ Θ hΘ m hm Q hQ ξ hξ hT3 hmax
  refine h.trans (measureReal_mono ?_)
  intro o ho
  simp only [Set.mem_ofPred_eq] at ho ⊢
  rw [ho]
  exact maxLeaf_eq G hG hn

end MaxLeafFormal.Budget