import Erdos1135.Tao.Density.FiniteV2Lift
import Erdos1135.Terras.Density.NaturalDensity

/-!
# Inclusive Natural Counts Under Dilation

This file begins the P4 hard-interface pilot.  It transports the project's
half-open natural count through the inherited quotient-form dilation while
keeping the source's inclusive endpoint explicit.
-/

namespace Erdos1135
namespace ND

open Filter
open scoped Topology

private theorem natCount_succ_of_mem
    {A : Set ℕ} {N : ℕ} (hN : N ∈ A) :
    Terras.natCount A (N + 1) = Terras.natCount A N + 1 := by
  unfold Terras.natCount
  rw [Nat.count_succ]
  simp [hN]

private theorem natCount_succ_of_not_mem
    {A : Set ℕ} {N : ℕ} (hN : N ∉ A) :
    Terras.natCount A (N + 1) = Terras.natCount A N := by
  unfold Terras.natCount
  rw [Nat.count_succ]
  simp [hN]

/-- Dilation by a positive natural maps the inclusive prefix `[0,X]`
exactly to the inclusive source prefix `[0,X/q]`.

Both successors are load-bearing because `Terras.natCount A N` counts the
half-open prefix `[0,N)`.  The statement includes `X = 0` and preserves
membership of zero in the arbitrary source set `S`.
-/
theorem natCount_natDilate_succ
    (S : Set ℕ) (q X : ℕ) (hq : 0 < q) :
    Terras.natCount (Tao.natDilate q S) (X + 1) =
      Terras.natCount S (X / q + 1) := by
  induction X with
  | zero =>
      simp only [Nat.zero_div]
      by_cases hzero : 0 ∈ S
      · have hmem : 0 ∈ Tao.natDilate q S :=
          (Tao.mem_natDilate_iff_exists_mul hq).2 ⟨0, hzero, by simp⟩
        rw [natCount_succ_of_mem hmem, natCount_succ_of_mem hzero]
        simp [Terras.natCount]
      · have hnot : 0 ∉ Tao.natDilate q S := by
          intro hmem
          exact hzero (by simpa using hmem.2)
        rw [natCount_succ_of_not_mem hnot,
          natCount_succ_of_not_mem hzero]
        simp [Terras.natCount]
  | succ X ih =>
      change
        Terras.natCount (Tao.natDilate q S) ((X + 1) + 1) =
          Terras.natCount S ((X + 1) / q + 1)
      by_cases hdvd : q ∣ X + 1
      · have hdiv : (X + 1) / q = X / q + 1 :=
          Nat.succ_div_of_dvd hdvd
        by_cases hS : X / q + 1 ∈ S
        · have hmem : X + 1 ∈ Tao.natDilate q S := by
            exact ⟨hdvd, by simpa only [hdiv] using hS⟩
          rw [natCount_succ_of_mem hmem, ih, hdiv,
            natCount_succ_of_mem hS]
        · have hmem : X + 1 ∉ Tao.natDilate q S := by
            intro hmem
            exact hS (by simpa only [hdiv] using hmem.2)
          rw [natCount_succ_of_not_mem hmem, ih, hdiv,
            natCount_succ_of_not_mem hS]
      · have hdiv : (X + 1) / q = X / q :=
          Nat.succ_div_of_not_dvd hdvd
        have hnot : X + 1 ∉ Tao.natDilate q S := by
          intro hmem
          exact hdvd hmem.1
        rw [natCount_succ_of_not_mem hnot, ih, hdiv]

/-- Natural dilation by a positive factor scales natural density by the
reciprocal factor.

The proof deliberately works at successor endpoints.  The checked finite
identity above sends the inclusive target prefix `[0,X]` to the inclusive
source prefix `[0,X/q]`; the resulting endpoint factor tends to `1/q`.
-/
theorem hasNatDensity_natDilate
    {S : Set ℕ} {d : ℝ}
    (hS : Terras.HasNatDensity S d) (q : ℕ) (hq : 0 < q) :
    Terras.HasNatDensity (Tao.natDilate q S) (d / (q : ℝ)) := by
  unfold Terras.HasNatDensity at hS ⊢
  have hsourceIndex : Tendsto (fun X : ℕ => X / q + 1) atTop atTop :=
    (Filter.tendsto_add_atTop_nat 1).comp
      (Nat.tendsto_div_const_atTop hq.ne')
  have hsource :
      Tendsto (fun X : ℕ => Terras.natCountingRatio S (X / q + 1))
        atTop (nhds d) :=
    hS.comp hsourceIndex
  have hscaleRaw :
      Tendsto
        (fun X : ℕ =>
          (((X / q : ℕ) : ℝ) / (X : ℝ)) *
              ((X : ℝ) / ((X : ℝ) + 1)) +
            (1 : ℝ) / ((X : ℝ) + 1))
        atTop (nhds (1 / (q : ℝ))) := by
    simpa using
      ((Terras.tendsto_nat_div_const_ratio hq).mul
        (tendsto_natCast_div_add_atTop (𝕜 := ℝ) (1 : ℝ))).add
        (tendsto_one_div_add_atTop_nhds_zero_nat (𝕜 := ℝ))
  have hscale :
      Tendsto
        (fun X : ℕ =>
          ((X / q + 1 : ℕ) : ℝ) / ((X + 1 : ℕ) : ℝ))
        atTop (nhds (1 / (q : ℝ))) := by
    refine hscaleRaw.congr' ?_
    filter_upwards [eventually_gt_atTop (0 : ℕ)] with X hX
    have hXne : (X : ℝ) ≠ 0 := by
      exact_mod_cast hX.ne'
    have hX1ne : (X : ℝ) + 1 ≠ 0 := by
      positivity
    simp only [Nat.cast_add, Nat.cast_one]
    field_simp [hXne, hX1ne] <;> ring
  have hproduct := hsource.mul hscale
  have hshifted :
      Tendsto
        (fun X : ℕ =>
          Terras.natCountingRatio (Tao.natDilate q S) (X + 1))
        atTop (nhds (d / (q : ℝ))) := by
    have hproduct' :
        Tendsto
          (fun X : ℕ =>
            Terras.natCountingRatio S (X / q + 1) *
              (((X / q + 1 : ℕ) : ℝ) / ((X + 1 : ℕ) : ℝ)))
          atTop (nhds (d / (q : ℝ))) := by
      simpa [div_eq_mul_inv] using hproduct
    refine hproduct'.congr' ?_
    filter_upwards with X
    unfold Terras.natCountingRatio
    rw [natCount_natDilate_succ S q X hq]
    have hsourceNe : ((X / q + 1 : ℕ) : ℝ) ≠ 0 := by
      positivity
    have htargetNe : ((X + 1 : ℕ) : ℝ) ≠ 0 := by
      positivity
    field_simp [hsourceNe, htargetNe]
  exact (Filter.tendsto_add_atTop_iff_nat 1).mp hshifted

end ND
end Erdos1135