import Erdos1135.Terras.Probability.Bitstrings

/-!
# Terras/Everett Binomial-Tail Reduction

This module starts the elementary tail side of the Terras/Everett density route.  It proves that
noncontracting parity words must lie in a simple upper tail, using only the rational bound
`3 / 5 < log_3 2` encoded as exact power arithmetic.
-/

namespace Erdos1135
namespace Terras

set_option maxHeartbeats 400000

/--
The exact arithmetic threshold used to avoid logarithms: if the odd-count ratio is at most `3/5`,
then the accelerated branch multiplier is strictly smaller than one.
-/
theorem pow_three_lt_pow_two_of_five_mul_le_three_mul {r k : ℕ}
    (hk : 0 < k) (h : 5 * r ≤ 3 * k) : 3 ^ r < 2 ^ k := by
  by_cases hr : r = 0
  · subst r
    exact Nat.one_lt_two_pow (Nat.ne_of_gt hk)
  · have hbase : (3 : ℕ) ^ 5 < 2 ^ 8 := by norm_num
    have hraise : ((3 : ℕ) ^ 5) ^ r < (2 ^ 8) ^ r := by
      exact Nat.pow_lt_pow_left hbase hr
    have hleft_raise : ((3 : ℕ) ^ 5) ^ r = 3 ^ (5 * r) := by
      rw [← Nat.pow_mul]
    have hright_raise : ((2 : ℕ) ^ 8) ^ r = 2 ^ (8 * r) := by
      rw [← Nat.pow_mul]
    have hraise' : (3 : ℕ) ^ (5 * r) < 2 ^ (8 * r) := by
      rwa [hleft_raise, hright_raise] at hraise
    have hright_le : 2 ^ (8 * r) ≤ 2 ^ (5 * k) := by
      have h8 : 8 * r ≤ 5 * k := by omega
      exact Nat.pow_le_pow_right (by norm_num : 0 < (2 : ℕ)) h8
    have htmp : (3 : ℕ) ^ (5 * r) < 2 ^ (5 * k) :=
      lt_of_lt_of_le hraise' hright_le
    have hleft_eq : ((3 : ℕ) ^ r) ^ 5 = 3 ^ (5 * r) := by
      rw [← Nat.pow_mul]
      congr 1
      omega
    have hright_eq : ((2 : ℕ) ^ k) ^ 5 = 2 ^ (5 * k) := by
      rw [← Nat.pow_mul]
      congr 1
      omega
    have hcomp : ((3 : ℕ) ^ r) ^ 5 < (2 ^ k) ^ 5 := by
      rw [hleft_eq, hright_eq]
      exact htmp
    exact (Nat.pow_lt_pow_iff_left (by omega : (5 : ℕ) ≠ 0)).mp hcomp

theorem contractingOddCount_of_five_mul_le_three_mul {k r : ℕ}
    (hk : 0 < k) (hr : r ≤ k) (h : 5 * r ≤ 3 * k) :
    r ∈ contractingOddCounts k := by
  rw [mem_contractingOddCounts_iff]
  exact ⟨hr, pow_three_lt_pow_two_of_five_mul_le_three_mul hk h⟩

theorem contracting_of_five_mul_numOdd_le_three_mul {k : ℕ} (w : ParityWord k)
    (hk : 0 < k) (h : 5 * numOdd w ≤ 3 * k) :
    3 ^ numOdd w < 2 ^ k :=
  pow_three_lt_pow_two_of_five_mul_le_three_mul hk h

/-- Odd-count values whose branch multiplier does not contract. -/
def noncontractingOddCounts (k : ℕ) : Finset ℕ :=
  (Finset.range (k + 1)).filter fun r => ¬ 3 ^ r < 2 ^ k

/-- A simple upper tail that contains every noncontracting odd-count value. -/
def highOddCounts (k : ℕ) : Finset ℕ :=
  (Finset.range (k + 1)).filter fun r => 3 * k < 5 * r

lemma binomial_weighted_sum_eq_five_pow (k : ℕ) :
    (∑ r ∈ Finset.range (k + 1), Nat.choose k r * 3 ^ r * 2 ^ (k - r)) = 5 ^ k := by
  calc
    (∑ r ∈ Finset.range (k + 1), Nat.choose k r * 3 ^ r * 2 ^ (k - r))
        = ∑ r ∈ Finset.range (k + 1), 3 ^ r * 2 ^ (k - r) * k.choose r := by
          apply Finset.sum_congr rfl
          intro r hr
          ring_nf
    _ = (3 + 2 : ℕ) ^ k := by
          exact (add_pow (3 : ℕ) (2 : ℕ) k).symm
    _ = 5 ^ k := by norm_num

lemma choose_mul_three_pow_mul_two_pow_le_five_pow (k r : ℕ) :
    Nat.choose k r * 3 ^ r * 2 ^ (k - r) ≤ 5 ^ k := by
  by_cases hr : r ≤ k
  · have hmem : r ∈ Finset.range (k + 1) := by
      rw [Finset.mem_range]
      omega
    have hterm :
        Nat.choose k r * 3 ^ r * 2 ^ (k - r) ≤
          ∑ m ∈ Finset.range (k + 1), Nat.choose k m * 3 ^ m * 2 ^ (k - m) := by
      exact Finset.single_le_sum
        (s := Finset.range (k + 1))
        (f := fun m => Nat.choose k m * 3 ^ m * 2 ^ (k - m))
        (by intro m hm; exact Nat.zero_le _) hmem
    exact hterm.trans_eq (binomial_weighted_sum_eq_five_pow k)
  · have hzero : Nat.choose k r = 0 := Nat.choose_eq_zero_of_lt (by omega)
    simp [hzero]

theorem highOddCounts_weighted_sum_le_five_pow (k : ℕ) :
    (∑ r ∈ highOddCounts k, Nat.choose k r * 3 ^ r * 2 ^ (k - r)) ≤ 5 ^ k := by
  have hsub : highOddCounts k ⊆ Finset.range (k + 1) := by
    intro r hr
    have hrange : r ≤ k ∧ 3 * k < 5 * r := by
      simpa [highOddCounts] using hr
    rw [Finset.mem_range]
    omega
  calc
    (∑ r ∈ highOddCounts k, Nat.choose k r * 3 ^ r * 2 ^ (k - r))
        ≤ ∑ r ∈ Finset.range (k + 1), Nat.choose k r * 3 ^ r * 2 ^ (k - r) := by
          exact Finset.sum_le_sum_of_subset_of_nonneg hsub
            (by intro x hx hxnot; exact Nat.zero_le _)
    _ = 5 ^ k := binomial_weighted_sum_eq_five_pow k

lemma highOddCount_weight_lower_bound {k r : ℕ} (hr : r ∈ highOddCounts k) :
    3 ^ (3 * (k / 5)) * 2 ^ (k - 3 * (k / 5)) ≤ 3 ^ r * 2 ^ (k - r) := by
  have hr' : r ≤ k ∧ 3 * k < 5 * r := by
    simpa [highOddCounts] using hr
  let a := 3 * (k / 5)
  have h5q : 5 * (k / 5) ≤ k := Nat.mul_div_le k 5
  have ha_le_r : a ≤ r := by
    dsimp [a]
    omega
  have hr_le_k : r ≤ k := hr'.1
  let d := r - a
  have hr_eq : r = a + d := by
    dsimp [d]
    omega
  have hk_sub : k - a = d + (k - r) := by
    dsimp [d]
    omega
  calc
    3 ^ a * 2 ^ (k - a)
        = 3 ^ a * 2 ^ (d + (k - r)) := by rw [hk_sub]
    _ = 3 ^ a * (2 ^ d * 2 ^ (k - r)) := by rw [Nat.pow_add]
    _ ≤ 3 ^ a * (3 ^ d * 2 ^ (k - r)) := by
          exact Nat.mul_le_mul_left _
            (Nat.mul_le_mul_right _ (Nat.pow_le_pow_left (by norm_num : (2 : ℕ) ≤ 3) d))
    _ = 3 ^ (a + d) * 2 ^ (k - r) := by
          rw [Nat.pow_add]
          ring
    _ = 3 ^ r * 2 ^ (k - r) := by rw [hr_eq]

theorem highOddCounts_sum_choose_mul_weight_lower_le_five_pow (k : ℕ) :
    (3 ^ (3 * (k / 5)) * 2 ^ (k - 3 * (k / 5))) *
        (∑ r ∈ highOddCounts k, Nat.choose k r) ≤ 5 ^ k := by
  let W := 3 ^ (3 * (k / 5)) * 2 ^ (k - 3 * (k / 5))
  have hweighted :
      W * (∑ r ∈ highOddCounts k, Nat.choose k r) ≤
        ∑ r ∈ highOddCounts k, Nat.choose k r * 3 ^ r * 2 ^ (k - r) := by
    rw [Finset.mul_sum]
    apply Finset.sum_le_sum
    intro r hr
    have hW : W ≤ 3 ^ r * 2 ^ (k - r) := by
      simpa [W] using highOddCount_weight_lower_bound hr
    calc
      W * Nat.choose k r ≤ (3 ^ r * 2 ^ (k - r)) * Nat.choose k r :=
        Nat.mul_le_mul_right _ hW
      _ = Nat.choose k r * 3 ^ r * 2 ^ (k - r) := by ring
  exact hweighted.trans (highOddCounts_weighted_sum_le_five_pow k)

lemma mem_noncontractingOddCounts_iff {k r : ℕ} :
    r ∈ noncontractingOddCounts k ↔ r ≤ k ∧ ¬ 3 ^ r < 2 ^ k := by
  simp [noncontractingOddCounts]

lemma mem_highOddCounts_iff {k r : ℕ} :
    r ∈ highOddCounts k ↔ r ≤ k ∧ 3 * k < 5 * r := by
  simp [highOddCounts]

lemma numOdd_mem_noncontractingOddCounts_iff {k : ℕ} (w : ParityWord k) :
    numOdd w ∈ noncontractingOddCounts k ↔ ¬ 3 ^ numOdd w < 2 ^ k := by
  rw [mem_noncontractingOddCounts_iff]
  exact ⟨fun h => h.2, fun h => ⟨numOdd_le w, h⟩⟩

lemma three_mul_lt_five_mul_of_not_contracting {k r : ℕ} (hk : 0 < k)
    (hnot : ¬ 3 ^ r < 2 ^ k) :
    3 * k < 5 * r := by
  by_contra hle
  have hlin : 5 * r ≤ 3 * k := by omega
  exact hnot (pow_three_lt_pow_two_of_five_mul_le_three_mul hk hlin)

lemma noncontractingOddCounts_subset_highOddCounts {k : ℕ} (hk : 0 < k) :
    noncontractingOddCounts k ⊆ highOddCounts k := by
  intro r hr
  rw [mem_noncontractingOddCounts_iff] at hr
  rw [mem_highOddCounts_iff]
  exact ⟨hr.1, three_mul_lt_five_mul_of_not_contracting hk hr.2⟩

noncomputable def noncontractingWordsEquivSigmaOddCount (k : ℕ) :
    {w : ParityWord k // ¬ 3 ^ numOdd w < 2 ^ k} ≃
      Sigma (fun r : {r : ℕ // r ∈ noncontractingOddCounts k} =>
        {w : ParityWord k // numOdd w = r.1}) where
  toFun w :=
    ⟨⟨numOdd w.1, (numOdd_mem_noncontractingOddCounts_iff w.1).mpr w.2⟩, ⟨w.1, rfl⟩⟩
  invFun rw :=
    ⟨rw.2.1, by
      rw [rw.2.2]
      exact (mem_noncontractingOddCounts_iff.mp rw.1.2).2⟩
  left_inv w := by
    ext i
    rfl
  right_inv rw := by
    cases rw with
    | mk r w =>
      cases r with
      | mk r hr =>
        cases w with
        | mk w hw =>
          simp at hw ⊢
          constructor
          · exact hw
          · subst r
            simp

theorem card_noncontracting_words_eq_sum_choose (k : ℕ) :
    Fintype.card {w : ParityWord k // ¬ 3 ^ numOdd w < 2 ^ k} =
      ∑ r ∈ noncontractingOddCounts k, Nat.choose k r := by
  classical
  rw [Fintype.card_congr (noncontractingWordsEquivSigmaOddCount k)]
  rw [Fintype.card_sigma]
  simp [card_words_with_numOdd]
  rw [Finset.sum_attach]

theorem card_noncontracting_words_le_high_odd_tail {k : ℕ} (hk : 0 < k) :
    Fintype.card {w : ParityWord k // ¬ 3 ^ numOdd w < 2 ^ k} ≤
      ∑ r ∈ highOddCounts k, Nat.choose k r := by
  rw [card_noncontracting_words_eq_sum_choose]
  exact Finset.sum_le_sum_of_subset_of_nonneg
    (noncontractingOddCounts_subset_highOddCounts hk)
    (by intro x hx hxnot; exact Nat.zero_le _)

end Terras
end Erdos1135
