import Erdos1135.Terras.LogTimeStopping

/-!
# Effective logarithmic-time Terras bound

This module makes the logarithmic stopping-time argument quantitative.  At a chosen parity depth
`k`, it gives an explicit upper bound for the number of starting values below `N` that have not
descended within `C * log n` accelerated steps.  The three terms in the final real-valued bound
are the explicit branch/logarithm threshold, the geometric binomial-tail bound, and the one-period
rounding error.
-/

namespace Erdos1135
namespace Terras

open Filter
open scoped Topology

lemma branchScaleList_le_three_pow_length : ∀ ps : List Bool,
    branchScaleList ps ≤ (3 : ℤ) ^ ps.length
  | [] => by simp [branchScaleList]
  | isOdd :: ps => by
      have ih := branchScaleList_le_three_pow_length ps
      have hscale : branchScale isOdd ≤ 3 := by
        cases isOdd <;> simp [branchScale]
      have hnonneg := branchScaleList_nonneg ps
      simp only [branchScaleList, List.length_cons, pow_succ]
      nlinarith

lemma branchConstList_le_three_pow_length : ∀ ps : List Bool,
    branchConstList ps ≤ (3 : ℤ) ^ ps.length
  | [] => by simp [branchConstList]
  | isOdd :: ps => by
      have hscale := branchScaleList_le_three_pow_length ps
      have hconst := branchConstList_le_three_pow_length ps
      have hscale_nonneg := branchScaleList_nonneg ps
      have hconst_nonneg := branchConstList_nonneg ps
      have hoffset : branchOffset isOdd ≤ 1 := by
        cases isOdd <;> simp [branchOffset]
      have hoffset_nonneg := branchOffset_nonneg isOdd
      simp only [branchConstList, List.length_cons, pow_succ]
      nlinarith

theorem branchConst_le_three_pow {k : ℕ} (w : ParityWord k) :
    branchConst w ≤ (3 : ℤ) ^ k := by
  simpa [branchConst, parityWordToList_length] using
    branchConstList_le_three_pow_length (parityWordToList w)

theorem branchThreshold_le_three_pow_add_one {k : ℕ} (w : ParityWord k) :
    branchThreshold w ≤ 3 ^ k + 1 := by
  have hconst_nonneg := branchConst_nonneg w
  have hconst := branchConst_le_three_pow w
  have habs_cast : (Int.natAbs (branchConst w) : ℤ) = branchConst w :=
    Int.natAbs_of_nonneg hconst_nonneg
  have habs : Int.natAbs (branchConst w) ≤ 3 ^ k := by
    exact_mod_cast (habs_cast.symm ▸ hconst)
  simpa [branchThreshold] using Nat.add_le_add_right habs 1

theorem branchThresholdMax_le_three_pow_add_one (k : ℕ) :
    branchThresholdMax k ≤ 3 ^ k + 1 := by
  classical
  unfold branchThresholdMax
  exact Finset.sup_le fun w _hw => branchThreshold_le_three_pow_add_one w

/-- A closed-form threshold that simultaneously dominates every length-`k` branch constant and
ensures that `k` accelerated steps fit inside the `C * log n` budget. -/
noncomputable def effectiveLogTimeThreshold (C : ℝ) (k : ℕ) : ℕ :=
  max (3 ^ k + 1) ⌈Real.exp ((k : ℝ) / C)⌉₊

theorem branchThresholdMax_le_effectiveLogTimeThreshold (C : ℝ) (k : ℕ) :
    branchThresholdMax k ≤ effectiveLogTimeThreshold C k :=
  (branchThresholdMax_le_three_pow_add_one k).trans (Nat.le_max_left _ _)

theorem steps_le_log_budget_of_effectiveLogTimeThreshold_le
    {C : ℝ} (hC : 0 < C) {k n : ℕ}
    (hn : effectiveLogTimeThreshold C k ≤ n) :
    (k : ℝ) ≤ C * Real.log (n : ℝ) := by
  have hceil : ⌈Real.exp ((k : ℝ) / C)⌉₊ ≤ n :=
    (Nat.le_max_right _ _).trans hn
  have hexp : Real.exp ((k : ℝ) / C) ≤ (n : ℝ) := Nat.ceil_le.mp hceil
  have hn_pos : (0 : ℝ) < (n : ℝ) := lt_of_lt_of_le (Real.exp_pos _) hexp
  have hlog : (k : ℝ) / C ≤ Real.log (n : ℝ) :=
    (Real.le_log_iff_exp_le hn_pos).2 hexp
  calc
    (k : ℝ) = C * ((k : ℝ) / C) := by field_simp [hC.ne']
    _ ≤ C * Real.log (n : ℝ) := mul_le_mul_of_nonneg_left hlog hC.le

/-- Residues whose length-`k` parity word is not contracting. -/
noncomputable def noncontractingResidues (k : ℕ) : Finset (ZMod (2 ^ k)) :=
  Finset.univ \ contractingResidues k

theorem mem_noncontractingResidues_iff {k : ℕ} {r : ZMod (2 ^ k)} :
    r ∈ noncontractingResidues k ↔ r ∉ contractingResidues k := by
  simp [noncontractingResidues]

theorem card_noncontractingResidues (k : ℕ) :
    (noncontractingResidues k).card =
      Fintype.card {w : ParityWord k // ¬ 3 ^ numOdd w < 2 ^ k} := by
  classical
  have hsplit := card_contracting_words_add_card_noncontracting_words k
  have hdiff :
      (noncontractingResidues k).card = 2 ^ k - (contractingResidues k).card := by
    simp [noncontractingResidues, Finset.card_sdiff, ZMod.card]
  rw [hdiff, card_contractingResidues]
  omega

theorem natCount_union_le (s t : Set ℕ) (N : ℕ) :
    natCount (s ∪ t) N ≤ natCount s N + natCount t N := by
  induction N with
  | zero => simp [natCount]
  | succ N ih =>
      unfold natCount at ih ⊢
      repeat rw [Nat.count_succ]
      by_cases hs : N ∈ s <;> by_cases ht : N ∈ t <;> simp [hs, ht] at * <;> omega

theorem natCount_finsetSetUnion_of_pairwise_disjoint
    {ι : Type*} [DecidableEq ι]
    (I : Finset ι) (A : ι → Set ℕ)
    (hdisj : ∀ i ∈ I, ∀ j ∈ I, i ≠ j → Disjoint (A i) (A j))
    (N : ℕ) :
    natCount (finsetSetUnion I A) N = ∑ i ∈ I, natCount (A i) N := by
  induction I using Finset.induction_on with
  | empty => simp [finsetSetUnion_empty, natCount]
  | @insert a I ha ih =>
      have hdisj_tail : ∀ i ∈ I, ∀ j ∈ I, i ≠ j → Disjoint (A i) (A j) := by
        intro i hi j hj hij
        exact hdisj i (by simp [hi]) j (by simp [hj]) hij
      have hhead_tail : Disjoint (A a) (finsetSetUnion I A) := by
        rw [Set.disjoint_left]
        intro n hhead htail
        rcases htail with ⟨j, hj, hAj⟩
        have haj : a ≠ j := by
          intro heq
          exact ha (heq ▸ hj)
        exact Set.disjoint_left.mp
          (hdisj a (by simp) j (by simp [hj]) haj) hhead hAj
      rw [finsetSetUnion_insert ha]
      rw [natCount_union_of_disjoint hhead_tail]
      rw [ih hdisj_tail]
      simp [Finset.sum_insert ha]

/-- A union of `R.card` residue classes modulo `m` has at most
`R.card * (N / m + 1)` representatives below `N`. -/
theorem natCount_residueSet_le {m : ℕ} (hm : 0 < m)
    (R : Finset (ZMod m)) (N : ℕ) :
    natCount (residueSet m R) N ≤ R.card * (N / m + 1) := by
  rw [residueSet_eq_finsetSetUnion_val_modEq hm]
  rw [natCount_finsetSetUnion_of_pairwise_disjoint R
    (fun z : ZMod m => {n : ℕ | n ≡ z.val [MOD m]})
    (by
      intro a _ha b _hb hab
      exact modEqResidueClass_disjoint_of_ne hm hab)
    N]
  calc
    ∑ z ∈ R, natCount {n : ℕ | n ≡ z.val [MOD m]} N
        ≤ ∑ _z ∈ R, (N / m + 1) := by
          apply Finset.sum_le_sum
          intro z _hz
          rw [natCount_modEq hm]
          split <;> omega
    _ = R.card * (N / m + 1) := by simp

/-- Starting values that have not descended within their `C * log n` accelerated-step budget. -/
def finiteStoppingTimeLogBudgetFailureSet (C : ℝ) : Set ℕ :=
  (finiteStoppingTimeLogBudgetSet C)ᶜ

theorem finiteStoppingTimeLogBudgetFailureSet_subset
    {C : ℝ} (hC : 0 < C) (k : ℕ) :
    finiteStoppingTimeLogBudgetFailureSet C ⊆
      {n : ℕ | n < effectiveLogTimeThreshold C k} ∪
        residueSet (2 ^ k) (noncontractingResidues k) := by
  intro n hn
  by_cases hlow : n < effectiveLogTimeThreshold C k
  · exact Or.inl hlow
  · apply Or.inr
    have hn_threshold : effectiveLogTimeThreshold C k ≤ n := le_of_not_gt hlow
    have hn_branch : branchThresholdMax k ≤ n :=
      (branchThresholdMax_le_effectiveLogTimeThreshold C k).trans hn_threshold
    have hn_budget : (k : ℝ) ≤ C * Real.log (n : ℝ) :=
      steps_le_log_budget_of_effectiveLogTimeThreshold_le hC hn_threshold
    change (n : ZMod (2 ^ k)) ∈ noncontractingResidues k
    rw [mem_noncontractingResidues_iff]
    intro hcontracting
    change n ∉ finiteStoppingTimeLogBudgetSet C at hn
    exact hn ⟨k, hn_budget,
      descends_of_mem_contractingResidues_ge hn_branch hcontracting⟩

/-- Fully explicit finite-count form of the quantitative Terras theorem.  The right side consists
of the closed branch/log threshold plus one rounded count for each noncontracting residue class. -/
theorem natCount_finiteStoppingTimeLogBudgetFailureSet_le
    {C : ℝ} (hC : 0 < C) (k N : ℕ) :
    natCount (finiteStoppingTimeLogBudgetFailureSet C) N ≤
      effectiveLogTimeThreshold C k +
        (noncontractingResidues k).card * (N / 2 ^ k + 1) := by
  calc
    natCount (finiteStoppingTimeLogBudgetFailureSet C) N
        ≤ natCount
            ({n : ℕ | n < effectiveLogTimeThreshold C k} ∪
              residueSet (2 ^ k) (noncontractingResidues k)) N :=
          natCount_mono (finiteStoppingTimeLogBudgetFailureSet_subset hC k) N
    _ ≤ natCount {n : ℕ | n < effectiveLogTimeThreshold C k} N +
          natCount (residueSet (2 ^ k) (noncontractingResidues k)) N :=
        natCount_union_le _ _ N
    _ ≤ effectiveLogTimeThreshold C k +
          (noncontractingResidues k).card * (N / 2 ^ k + 1) :=
        Nat.add_le_add
          (natCount_Iio_le (effectiveLogTimeThreshold C k) N)
          (natCount_residueSet_le (by positivity) (noncontractingResidues k) N)

/-- Effective geometric-rate form.  For every chosen positive depth `k` and cutoff `N`, the
proportion of exceptions below `N` is bounded by

* the explicit initial threshold divided by `N`;
* `625 * (3125 / 3456)^(k / 5)` from the checked binomial-tail estimate; and
* `2^k / N`, the finite residue-class rounding error.
-/
theorem natCountingRatio_finiteStoppingTimeLogBudgetFailureSet_le
    {C : ℝ} (hC : 0 < C) {k N : ℕ} (hk : 0 < k) (hN : 0 < N) :
    natCountingRatio (finiteStoppingTimeLogBudgetFailureSet C) N ≤
      (effectiveLogTimeThreshold C k : ℝ) / (N : ℝ) +
        625 * ((3125 : ℝ) / 3456) ^ (k / 5) +
        (2 ^ k : ℝ) / (N : ℝ) := by
  let R : ℕ := (noncontractingResidues k).card
  let M : ℕ := 2 ^ k
  have hM : 0 < M := by simp [M]
  have hN_real : (0 : ℝ) < (N : ℝ) := by exact_mod_cast hN
  have hcount := natCount_finiteStoppingTimeLogBudgetFailureSet_le hC k N
  have hcount_real :
      (natCount (finiteStoppingTimeLogBudgetFailureSet C) N : ℝ) ≤
        ((effectiveLogTimeThreshold C k + R * (N / M + 1) : ℕ) : ℝ) := by
    exact_mod_cast hcount
  have hquotient :
      ((R : ℝ) * (N / M : ℕ)) / (N : ℝ) ≤ (R : ℝ) / (M : ℝ) := by
    have hdiv : ((N / M : ℕ) : ℝ) ≤ (N : ℝ) / (M : ℝ) := Nat.cast_div_le
    have hmul :
        (R : ℝ) * (N / M : ℕ) ≤ (R : ℝ) * ((N : ℝ) / (M : ℝ)) :=
      mul_le_mul_of_nonneg_left hdiv (by positivity)
    calc
      ((R : ℝ) * (N / M : ℕ)) / (N : ℝ)
          ≤ ((R : ℝ) * ((N : ℝ) / (M : ℝ))) / (N : ℝ) :=
        div_le_div_of_nonneg_right hmul hN_real.le
      _ = (R : ℝ) / (M : ℝ) := by
        field_simp [hN_real.ne', (show (M : ℝ) ≠ 0 by positivity)]
  have hR_le_M_nat : R ≤ M := by
    dsimp [R, M]
    simpa [ZMod.card] using Finset.card_le_univ (noncontractingResidues k)
  have hRround : (R : ℝ) / (N : ℝ) ≤ (M : ℝ) / (N : ℝ) := by
    exact div_le_div_of_nonneg_right (by exact_mod_cast hR_le_M_nat) hN_real.le
  have hratio :
      (R : ℝ) / (M : ℝ) = noncontractingWordRatio k := by
    simp [R, M, noncontractingWordRatio, card_noncontractingResidues]
  have htail :
      noncontractingWordRatio k ≤
        625 * ((3125 : ℝ) / 3456) ^ (k / 5) :=
    (noncontractingWordRatio_le_highOddTailRatio hk).trans
      (highOddTailRatio_le_geometric_bound k)
  unfold natCountingRatio
  calc
    (natCount (finiteStoppingTimeLogBudgetFailureSet C) N : ℝ) / (N : ℝ)
        ≤ ((effectiveLogTimeThreshold C k + R * (N / M + 1) : ℕ) : ℝ) /
            (N : ℝ) := div_le_div_of_nonneg_right hcount_real hN_real.le
    _ = (effectiveLogTimeThreshold C k : ℝ) / (N : ℝ) +
          ((R : ℝ) * (N / M : ℕ)) / (N : ℝ) +
          (R : ℝ) / (N : ℝ) := by
        push_cast
        ring
    _ ≤ (effectiveLogTimeThreshold C k : ℝ) / (N : ℝ) +
          (R : ℝ) / (M : ℝ) + (M : ℝ) / (N : ℝ) := by
        gcongr
    _ = (effectiveLogTimeThreshold C k : ℝ) / (N : ℝ) +
          noncontractingWordRatio k + (2 ^ k : ℝ) / (N : ℝ) := by
        rw [hratio]
        simp [M]
    _ ≤ (effectiveLogTimeThreshold C k : ℝ) / (N : ℝ) +
          625 * ((3125 : ℝ) / 3456) ^ (k / 5) +
          (2 ^ k : ℝ) / (N : ℝ) := by
        gcongr

/-- Purely geometric error bound after an explicit cutoff.  This is the convenient effective
form: choosing a larger parity depth reduces all three error terms geometrically, and the required
cutoff is the closed expression `effectiveLogTimeThreshold C k * 2^k`. -/
theorem natCountingRatio_finiteStoppingTimeLogBudgetFailureSet_le_geometric
    {C : ℝ} (hC : 0 < C) {k N : ℕ} (hk : 0 < k)
    (hN : effectiveLogTimeThreshold C k * 2 ^ k ≤ N) :
    natCountingRatio (finiteStoppingTimeLogBudgetFailureSet C) N ≤
      1 / (2 ^ k : ℝ) +
        625 * ((3125 : ℝ) / 3456) ^ (k / 5) +
        1 / (3 ^ k : ℝ) := by
  have hthreshold_pos : 0 < effectiveLogTimeThreshold C k := by
    unfold effectiveLogTimeThreshold
    positivity
  have hpow_two_pos : 0 < 2 ^ k := by positivity
  have hN_pos : 0 < N :=
    lt_of_lt_of_le (Nat.mul_pos hthreshold_pos hpow_two_pos) hN
  have hthreshold_ge_three : 3 ^ k ≤ effectiveLogTimeThreshold C k := by
    have hleft : 3 ^ k + 1 ≤ effectiveLogTimeThreshold C k := by
      exact Nat.le_max_left _ _
    omega
  have hthree_mul_two : 3 ^ k * 2 ^ k ≤ N :=
    (Nat.mul_le_mul_right (2 ^ k) hthreshold_ge_three).trans hN
  have hN_real : (0 : ℝ) < (N : ℝ) := by exact_mod_cast hN_pos
  have htwo_real : (0 : ℝ) < (2 ^ k : ℝ) := by positivity
  have hthree_real : (0 : ℝ) < (3 ^ k : ℝ) := by positivity
  have hthreshold_term :
      (effectiveLogTimeThreshold C k : ℝ) / (N : ℝ) ≤ 1 / (2 ^ k : ℝ) := by
    rw [div_le_div_iff₀ hN_real htwo_real]
    have hcast :
        (effectiveLogTimeThreshold C k : ℝ) * (2 ^ k : ℝ) ≤ (N : ℝ) := by
      exact_mod_cast hN
    simpa using hcast
  have hrounding_term : (2 ^ k : ℝ) / (N : ℝ) ≤ 1 / (3 ^ k : ℝ) := by
    rw [div_le_div_iff₀ hN_real hthree_real]
    have hcast : (3 ^ k : ℝ) * (2 ^ k : ℝ) ≤ (N : ℝ) := by
      exact_mod_cast hthree_mul_two
    simpa [mul_comm] using hcast
  calc
    natCountingRatio (finiteStoppingTimeLogBudgetFailureSet C) N
        ≤ (effectiveLogTimeThreshold C k : ℝ) / (N : ℝ) +
            625 * ((3125 : ℝ) / 3456) ^ (k / 5) +
            (2 ^ k : ℝ) / (N : ℝ) :=
          natCountingRatio_finiteStoppingTimeLogBudgetFailureSet_le hC hk hN_pos
    _ ≤ 1 / (2 ^ k : ℝ) +
          625 * ((3125 : ℝ) / 3456) ^ (k / 5) +
          1 / (3 ^ k : ℝ) := by
        gcongr

/-- Five-step block specialization.  At depth `5m`, the exceptional proportion is bounded by
`32^{-m} + 625 * (3125/3456)^m + 243^{-m}` beyond the displayed computable cutoff. -/
theorem natCountingRatio_finiteStoppingTimeLogBudgetFailureSet_le_five_mul
    {C : ℝ} (hC : 0 < C) {m N : ℕ} (hm : 0 < m)
    (hN : effectiveLogTimeThreshold C (5 * m) * 2 ^ (5 * m) ≤ N) :
    natCountingRatio (finiteStoppingTimeLogBudgetFailureSet C) N ≤
      1 / (32 : ℝ) ^ m +
        625 * ((3125 : ℝ) / 3456) ^ m +
        1 / (243 : ℝ) ^ m := by
  have hk : 0 < 5 * m := Nat.mul_pos (by norm_num) hm
  have h := natCountingRatio_finiteStoppingTimeLogBudgetFailureSet_le_geometric
    hC hk hN
  norm_num [pow_mul] at h ⊢
  exact h

/-- Literal `log n` specialization of the effective theorem. -/
theorem natCountingRatio_finiteStoppingTime_log_failure_le_five_mul
    {m N : ℕ} (hm : 0 < m)
    (hN : effectiveLogTimeThreshold 1 (5 * m) * 2 ^ (5 * m) ≤ N) :
    natCountingRatio
        {n : ℕ | ¬ ∃ k : ℕ,
          (k : ℝ) ≤ Real.log (n : ℝ) ∧ accelerated^[k] n < n}
        N ≤
      1 / (32 : ℝ) ^ m +
        625 * ((3125 : ℝ) / 3456) ^ m +
        1 / (243 : ℝ) ^ m := by
  have hset :
      finiteStoppingTimeLogBudgetFailureSet 1 =
        {n : ℕ | ¬ ∃ k : ℕ,
          (k : ℝ) ≤ Real.log (n : ℝ) ∧ accelerated^[k] n < n} := by
    ext n
    simp [finiteStoppingTimeLogBudgetFailureSet,
      finiteStoppingTimeLogBudgetSet]
  rw [← hset]
  exact natCountingRatio_finiteStoppingTimeLogBudgetFailureSet_le_five_mul
    (C := (1 : ℝ)) zero_lt_one hm hN

end Terras
end Erdos1135
