import Erdos1135.Tao.Density.LogDensity
import Erdos1135.Terras.Core.Defs
import Mathlib.NumberTheory.Harmonic.EulerMascheroni

/-!
# Powers-Of-Two Layers For Tao Logarithmic Density

This module starts the elementary density spine for the initial `2`-adic layer in the
Syracuse-to-Collatz transfer.  It keeps the first layer independent of `Terras.twoAdicExponent`:
we work with divisibility sets first, then later bridge those sets to the factorization-based
exponent on positive naturals.
-/

namespace Erdos1135
namespace Tao

open Filter
open scoped Topology

/-- Natural numbers divisible by `q`.  The logarithmic-density API ignores membership of `0`. -/
def multiples (q : ℕ) : Set ℕ :=
  {n : ℕ | q ∣ n}

/-- The exact elementary layer where `2^a` divides but `2^(a+1)` does not. -/
def powTwoExactLayer (a : ℕ) : Set ℕ :=
  multiples (2 ^ a) \ multiples (2 ^ (a + 1))

/-- The exclusive cutoff layer union `v₂(n) < K`, represented without using `v₂`. -/
def powTwoLayerLt (K : ℕ) : Set ℕ :=
  Set.univ \ multiples (2 ^ K)

theorem multiples_one : multiples 1 = Set.univ := by
  ext n
  simp [multiples]

theorem powTwoLayerLt_zero : powTwoLayerLt 0 = ∅ := by
  simp [powTwoLayerLt, multiples]

theorem multiples_pow_succ_subset (a : ℕ) :
    multiples (2 ^ (a + 1)) ⊆ multiples (2 ^ a) := by
  intro n hn
  exact dvd_trans (Nat.pow_dvd_pow 2 (Nat.le_succ a)) hn

theorem powTwoExactLayer_eq_diff (a : ℕ) :
    powTwoExactLayer a = multiples (2 ^ a) \ multiples (2 ^ (a + 1)) :=
  rfl

theorem powTwoLayerLt_eq_diff (K : ℕ) :
    powTwoLayerLt K = Set.univ \ multiples (2 ^ K) :=
  rfl

private theorem one_div_mul_logWeight_div_of_dvd
    {q X : ℕ} (hq : 0 < q) (hdvd : q ∣ X + 1) :
    (1 / (q : ℝ)) * logWeight (X / q) = logWeight X := by
  have hdiv : (X + 1) / q = X / q + 1 := Nat.succ_div_of_dvd hdvd
  have hmul : (X / q + 1) * q = X + 1 := by
    rw [← hdiv]
    exact Nat.div_mul_cancel hdvd
  have hmulR : (((X / q + 1 : ℕ) : ℝ) * (q : ℝ)) = ((X + 1 : ℕ) : ℝ) := by
    exact_mod_cast hmul
  unfold logWeight
  rw [← hmulR]
  field_simp [Nat.cast_ne_zero.mpr hq.ne']

theorem logCount_multiples (q X : ℕ) (hq : 0 < q) :
    logCount (multiples q) X = (1 / (q : ℝ)) * logMass (X / q) := by
  induction X with
  | zero =>
      simp [logCount, logMass]
  | succ X ih =>
      by_cases hdvd : q ∣ X + 1
      · have hmem : X + 1 ∈ multiples q := by
          simpa [multiples] using hdvd
        rw [logCount_succ_of_mem hmem, ih, Nat.succ_div_of_dvd hdvd, logMass_succ]
        rw [mul_add]
        rw [one_div_mul_logWeight_div_of_dvd hq hdvd]
      · have hmem : X + 1 ∉ multiples q := by
          simpa [multiples] using hdvd
        rw [logCount_succ_of_not_mem hmem, ih, Nat.succ_div_of_not_dvd hdvd]

private theorem logMass_sub_log_tendsto_euler :
    Tendsto (fun n : ℕ => logMass n - Real.log n) atTop
      (nhds Real.eulerMascheroniConstant) := by
  simpa [logMass_eq_harmonic] using Real.tendsto_harmonic_sub_log

private theorem log_nat_div_sub_log_tendsto_log {q : ℕ} (hq : 0 < q) :
    Tendsto (fun X : ℕ => Real.log ((X / q : ℕ) : ℝ) - Real.log (X : ℝ)) atTop
      (nhds (Real.log (1 / (q : ℝ)))) := by
  have hratio := Terras.tendsto_nat_div_const_ratio hq
  have hqreal : (1 / (q : ℝ)) ≠ 0 := by
    exact one_div_ne_zero (by exact_mod_cast hq.ne')
  have hlogratio :
      Tendsto (fun X : ℕ => Real.log (((X / q : ℕ) : ℝ) / (X : ℝ))) atTop
        (nhds (Real.log (1 / (q : ℝ)))) :=
    (Real.continuousAt_log hqreal).tendsto.comp hratio
  have hdiv_atTop : Tendsto (fun X : ℕ => X / q) atTop atTop :=
    Nat.tendsto_div_const_atTop hq.ne'
  have hX_pos : ∀ᶠ X in atTop, 0 < X := by
    exact eventually_atTop.mpr ⟨1, fun X hX => lt_of_lt_of_le zero_lt_one hX⟩
  have hX_div_pos : ∀ᶠ X in atTop, 0 < X / q := by
    filter_upwards [(tendsto_atTop.mp hdiv_atTop) 1] with X hX
    exact lt_of_lt_of_le zero_lt_one hX
  refine hlogratio.congr' ?_
  filter_upwards [hX_pos, hX_div_pos] with X hX hXdiv
  have hden : ((X : ℝ) ≠ 0) := by exact_mod_cast hX.ne'
  have hnum : (((X / q : ℕ) : ℝ) ≠ 0) := by exact_mod_cast hXdiv.ne'
  rw [Real.log_div hnum hden]

private theorem logMass_div_const_sub_tendsto_log {q : ℕ} (hq : 0 < q) :
    Tendsto (fun X : ℕ => logMass (X / q) - logMass X) atTop
      (nhds (Real.log (1 / (q : ℝ)))) := by
  have hdiv_atTop : Tendsto (fun X : ℕ => X / q) atTop atTop :=
    Nat.tendsto_div_const_atTop hq.ne'
  have hA :
      Tendsto (fun X : ℕ => logMass (X / q) - Real.log ((X / q : ℕ) : ℝ)) atTop
        (nhds Real.eulerMascheroniConstant) :=
    logMass_sub_log_tendsto_euler.comp hdiv_atTop
  have hB := logMass_sub_log_tendsto_euler
  have hC := log_nat_div_sub_log_tendsto_log hq
  have hsum := (hA.sub hB).add hC
  have hsum' :
      Tendsto
        (fun X : ℕ =>
          ((logMass (X / q) - Real.log ((X / q : ℕ) : ℝ)) -
              (logMass X - Real.log (X : ℝ))) +
            (Real.log ((X / q : ℕ) : ℝ) - Real.log (X : ℝ)))
        atTop (nhds (Real.log (1 / (q : ℝ)))) := by
    simpa using hsum
  refine hsum'.congr' ?_
  filter_upwards with X
  ring

theorem logMass_div_const_ratio_tendsto_one {q : ℕ} (hq : 0 < q) :
    Tendsto (fun X : ℕ => logMass (X / q) / logMass X) atTop (nhds 1) := by
  have hdiffdiv :
      Tendsto (fun X : ℕ => (logMass (X / q) - logMass X) / logMass X) atTop
        (nhds 0) :=
    (logMass_div_const_sub_tendsto_log hq).div_atTop logMass_tendsto_atTop
  have htarget :
      Tendsto (fun X : ℕ => 1 + (logMass (X / q) - logMass X) / logMass X) atTop
        (nhds 1) := by
    simpa using (tendsto_const_nhds.add hdiffdiv)
  refine htarget.congr' ?_
  filter_upwards [eventually_atTop.mpr ⟨1, fun X hX => lt_of_lt_of_le zero_lt_one hX⟩]
    with X hX
  have hmass : logMass X ≠ 0 := (logMass_pos hX).ne'
  field_simp [hmass]
  ring

theorem hasLogDensity_multiples (q : ℕ) (hq : 0 < q) :
    HasLogDensity (multiples q) (1 / (q : ℝ)) := by
  unfold HasLogDensity logCountingRatio
  have hlim :
      Tendsto (fun X : ℕ => (1 / (q : ℝ)) * (logMass (X / q) / logMass X))
        atTop (nhds ((1 / (q : ℝ)) * 1)) :=
    tendsto_const_nhds.mul (logMass_div_const_ratio_tendsto_one hq)
  have hlim' :
      Tendsto (fun X : ℕ => (1 / (q : ℝ)) * (logMass (X / q) / logMass X))
        atTop (nhds (1 / (q : ℝ))) := by
    simpa using hlim
  refine hlim'.congr' ?_
  filter_upwards with X
  rw [logCount_multiples q X hq]
  ring

theorem powTwoExactLayer_density_value (a : ℕ) :
    (1 / (2 : ℝ) ^ a) - 1 / (2 : ℝ) ^ (a + 1) = 1 / (2 : ℝ) ^ (a + 1) := by
  have hpow : (2 : ℝ) ^ (a + 1) = (2 : ℝ) ^ a * 2 := by
    rw [pow_succ]
  rw [hpow]
  field_simp [pow_ne_zero a (by norm_num : (2 : ℝ) ≠ 0)]
  ring

theorem powTwoExactLayer_density_value_inv (a : ℕ) :
    ((2 : ℝ) ^ a)⁻¹ - ((2 : ℝ) ^ (a + 1))⁻¹ = ((2 : ℝ) ^ (a + 1))⁻¹ := by
  simpa [one_div] using powTwoExactLayer_density_value a

/--
Conditional exact-layer density reducer.  The next analytic gate is proving logarithmic densities
for fixed multiples; this theorem packages the set-difference step that consumes that gate.
-/
theorem HasLogDensity.powTwoExactLayer_of_multiples
    {a : ℕ}
    (ha : HasLogDensity (multiples (2 ^ a)) (1 / (2 : ℝ) ^ a))
    (hsucc : HasLogDensity (multiples (2 ^ (a + 1))) (1 / (2 : ℝ) ^ (a + 1))) :
    HasLogDensity (powTwoExactLayer a) (1 / (2 : ℝ) ^ (a + 1)) := by
  rw [powTwoExactLayer_eq_diff]
  simpa [one_div, powTwoExactLayer_density_value_inv a] using
    HasLogDensity.diff_of_subset ha hsucc (multiples_pow_succ_subset a)

/--
Conditional exclusive-cutoff density reducer.  The cutoff convention is `a < K`, whose density is
`1 - 2^(-K)`.
-/
theorem HasLogDensity.powTwoLayerLt_of_multiples
    {K : ℕ}
    (hK : HasLogDensity (multiples (2 ^ K)) (1 / (2 : ℝ) ^ K)) :
    HasLogDensity (powTwoLayerLt K) (1 - 1 / (2 : ℝ) ^ K) := by
  rw [powTwoLayerLt_eq_diff]
  simpa using HasLogDensity.diff_of_subset hasLogDensity_univ hK (Set.subset_univ _)

theorem hasLogDensity_multiples_powTwo (a : ℕ) :
    HasLogDensity (multiples ((2 : ℕ) ^ a)) (1 / (2 : ℝ) ^ a) := by
  have hpow : 0 < (2 : ℕ) ^ a := Nat.pow_pos (by decide : 0 < (2 : ℕ))
  simpa using hasLogDensity_multiples ((2 : ℕ) ^ a) hpow

theorem hasLogDensity_powTwoExactLayer (a : ℕ) :
    HasLogDensity (powTwoExactLayer a) (1 / (2 : ℝ) ^ (a + 1)) :=
  HasLogDensity.powTwoExactLayer_of_multiples
    (hasLogDensity_multiples_powTwo a)
    (hasLogDensity_multiples_powTwo (a + 1))

theorem hasLogDensity_powTwoLayerLt (K : ℕ) :
    HasLogDensity (powTwoLayerLt K) (1 - 1 / (2 : ℝ) ^ K) :=
  HasLogDensity.powTwoLayerLt_of_multiples (hasLogDensity_multiples_powTwo K)

theorem hasLogDensity_powTwoLayerLt_zero :
    HasLogDensity (powTwoLayerLt 0) 0 := by
  simpa using hasLogDensity_powTwoLayerLt 0

theorem hasLogDensity_powTwoLayerLt_one :
    HasLogDensity (powTwoLayerLt 1) (1 / 2) := by
  convert hasLogDensity_powTwoLayerLt 1 using 1
  norm_num

theorem hasLogDensity_powTwoExactLayer_zero :
    HasLogDensity (powTwoExactLayer 0) (1 / 2) := by
  simpa using hasLogDensity_powTwoExactLayer 0

theorem powTwo_dvd_iff_le_twoAdicExponent {n a : ℕ} (hn : n ≠ 0) :
    (2 : ℕ) ^ a ∣ n ↔ a ≤ Terras.twoAdicExponent n := by
  simpa [Terras.twoAdicExponent] using
    (Nat.Prime.pow_dvd_iff_le_factorization (p := 2) (k := a) (n := n)
      Nat.prime_two hn)

theorem mem_powTwoExactLayer_iff_twoAdicExponent_eq {a n : ℕ} (hn : 0 < n) :
    n ∈ powTwoExactLayer a ↔ Terras.twoAdicExponent n = a := by
  have hnne : n ≠ 0 := Nat.ne_of_gt hn
  simp [powTwoExactLayer, multiples, powTwo_dvd_iff_le_twoAdicExponent hnne]
  omega

theorem mem_powTwoLayerLt_iff_twoAdicExponent_lt {K n : ℕ} (hn : 0 < n) :
    n ∈ powTwoLayerLt K ↔ Terras.twoAdicExponent n < K := by
  have hnne : n ≠ 0 := Nat.ne_of_gt hn
  simp [powTwoLayerLt, multiples, powTwo_dvd_iff_le_twoAdicExponent hnne]

theorem hasLogDensity_twoAdicExponent_eq (a : ℕ) :
    HasLogDensity {n : ℕ | 0 < n ∧ Terras.twoAdicExponent n = a}
      (1 / (2 : ℝ) ^ (a + 1)) := by
  refine HasLogDensity.congr_on_pos ?_ (hasLogDensity_powTwoExactLayer a)
  intro n hn
  simp [mem_powTwoExactLayer_iff_twoAdicExponent_eq hn, hn]

theorem hasLogDensity_twoAdicExponent_lt (K : ℕ) :
    HasLogDensity {n : ℕ | 0 < n ∧ Terras.twoAdicExponent n < K}
      (1 - 1 / (2 : ℝ) ^ K) := by
  refine HasLogDensity.congr_on_pos ?_ (hasLogDensity_powTwoLayerLt K)
  intro n hn
  simp [mem_powTwoLayerLt_iff_twoAdicExponent_lt hn, hn]

/-- Exact powers-of-two layer represented by an explicit odd quotient. -/
def powTwoOddLayer (a : ℕ) : Set ℕ :=
  {N : ℕ | ∃ M : ℕ, Odd M ∧ N = (2 : ℕ) ^ a * M}

theorem mem_powTwoOddLayer_iff {a N : ℕ} :
    N ∈ powTwoOddLayer a ↔ ∃ M : ℕ, Odd M ∧ N = (2 : ℕ) ^ a * M :=
  Iff.rfl

theorem twoAdicExponent_pow_mul_odd {a M : ℕ} (hM : Odd M) :
    Terras.twoAdicExponent ((2 : ℕ) ^ a * M) = a := by
  have hpowne : (2 : ℕ) ^ a ≠ 0 := by
    exact (Nat.pow_pos (by decide : 0 < (2 : ℕ))).ne'
  have hMne : M ≠ 0 := (Odd.pos hM).ne'
  have hnotdvd : ¬(2 : ℕ) ∣ M := by
    intro hdiv
    exact (Nat.not_even_iff_odd.mpr hM) (even_iff_two_dvd.mpr hdiv)
  have hfactorM : M.factorization 2 = 0 := Nat.factorization_eq_zero_of_not_dvd hnotdvd
  unfold Terras.twoAdicExponent
  rw [Nat.factorization_mul hpowne hMne]
  rw [Nat.Prime.factorization_pow Nat.prime_two]
  simp [hfactorM]

theorem mem_powTwoOddLayer_iff_twoAdicExponent_eq {a N : ℕ} (hN : 0 < N) :
    N ∈ powTwoOddLayer a ↔ Terras.twoAdicExponent N = a := by
  constructor
  · intro h
    rcases h with ⟨M, hM, rfl⟩
    exact twoAdicExponent_pow_mul_odd hM
  · intro hExp
    rcases Nat.exists_eq_two_pow_mul_odd hN.ne' with ⟨k, M, hM, hNrep⟩
    refine ⟨M, hM, ?_⟩
    have hk : k = a := by
      have hv : Terras.twoAdicExponent N = k := by
        rw [hNrep]
        exact twoAdicExponent_pow_mul_odd hM
      omega
    rw [hNrep, hk]

theorem powTwoOddLayer_eq_powTwoExactLayer (a : ℕ) :
    powTwoOddLayer a = powTwoExactLayer a := by
  ext N
  by_cases hN : 0 < N
  · rw [mem_powTwoOddLayer_iff_twoAdicExponent_eq hN,
      mem_powTwoExactLayer_iff_twoAdicExponent_eq hN]
  · have hNzero : N = 0 := Nat.eq_zero_of_not_pos hN
    subst N
    simp [powTwoOddLayer, powTwoExactLayer, multiples]

theorem hasLogDensity_powTwoOddLayer (a : ℕ) :
    HasLogDensity (powTwoOddLayer a) (1 / (2 : ℝ) ^ (a + 1)) := by
  rw [powTwoOddLayer_eq_powTwoExactLayer]
  exact hasLogDensity_powTwoExactLayer a

theorem hasLogDensity_odd :
    HasLogDensity {N : ℕ | Odd N} (1 / 2) := by
  have hoddLayer : HasLogDensity (powTwoOddLayer 0) (1 / 2) := by
    simpa using hasLogDensity_powTwoOddLayer 0
  refine HasLogDensity.congr_on_pos ?_ hoddLayer
  intro n hn
  simp [powTwoOddLayer]

end Tao
end Erdos1135