import Erdos1135.Terras.Parity.Residue
import Erdos1135.Terras.Parity.Contraction
import Erdos1135.Terras.Density.NaturalDensity
import Erdos1135.Terras.Probability.Asymptotic

/-!
# Terras/Everett Finite Stopping Time Setup

This module connects contracting parity words to residue classes and to actual finite stopping
time above the fixed branch threshold for a chosen prefix length.
-/

namespace Erdos1135
namespace Terras

noncomputable def contractingResidues (k : ℕ) : Finset (ZMod (2 ^ k)) := by
  classical
  exact (Finset.univ.filter fun w : ParityWord k => 3 ^ numOdd w < 2 ^ k).image
    (parityPrefix_equiv_zmod k)

theorem mem_contractingResidues_iff {k : ℕ} {r : ZMod (2 ^ k)} :
    r ∈ contractingResidues k ↔
      ∃ w : ParityWord k, 3 ^ numOdd w < 2 ^ k ∧ parityPrefix_equiv_zmod k w = r := by
  classical
  simp [contractingResidues]

theorem parityPrefixOfResidue_parityPrefix_equiv_zmod {k : ℕ} (w : ParityWord k) :
    parityPrefixOfResidue k (parityPrefix_equiv_zmod k w) = w := by
  change parityPrefixZModEquiv k ((parityPrefixZModEquiv k).symm w) = w
  exact (parityPrefixZModEquiv k).apply_symm_apply w

theorem parityPrefix_eq_of_natCast_eq_parityPrefix_equiv_zmod {k n : ℕ} {w : ParityWord k}
    (hcast : (n : ZMod (2 ^ k)) = parityPrefix_equiv_zmod k w) :
    parityPrefix k n = w := by
  have hprefix :=
    (natCast_eq_residue_iff_parityPrefix_eq k (parityPrefix_equiv_zmod k w) n).mp hcast
  simpa [parityPrefixOfResidue_parityPrefix_equiv_zmod w] using hprefix

theorem descends_of_mem_contractingResidues_ge {k n : ℕ}
    (hn : branchThresholdMax k ≤ n)
    (hmem : (n : ZMod (2 ^ k)) ∈ contractingResidues k) :
    accelerated^[k] n < n := by
  rcases mem_contractingResidues_iff.mp hmem with ⟨w, hcontract, hmap⟩
  have hprefix : parityPrefix k n = w :=
    parityPrefix_eq_of_natCast_eq_parityPrefix_equiv_zmod hmap.symm
  exact descends_of_contracting_of_branchThresholdMax_le hcontract hn hprefix

def finiteStoppingTimeSet : Set ℕ :=
  {n | ∃ k : ℕ, accelerated^[k] n < n}

theorem contractingResidues_ge_subset_finiteStoppingTimeSet (k : ℕ) :
    {n : ℕ | branchThresholdMax k ≤ n ∧ (n : ZMod (2 ^ k)) ∈ contractingResidues k} ⊆
      finiteStoppingTimeSet := by
  intro _n hn
  exact ⟨k, descends_of_mem_contractingResidues_ge hn.1 hn.2⟩

theorem hasNatDensity_contractingResidues_ge (k : ℕ) :
    HasNatDensity {n : ℕ | branchThresholdMax k ≤ n ∧
        (n : ZMod (2 ^ k)) ∈ contractingResidues k}
      (((contractingResidues k).card : ℝ) / (2 ^ k : ℝ)) := by
  have hpow_pos : 0 < 2 ^ k := by positivity
  simpa using hasNatDensity_residueSet_ge hpow_pos (contractingResidues k) (branchThresholdMax k)

theorem card_contractingResidues (k : ℕ) :
    (contractingResidues k).card = Fintype.card {w : ParityWord k // 3 ^ numOdd w < 2 ^ k} := by
  classical
  rw [contractingResidues]
  rw [Finset.card_image_of_injective _ (parityPrefix_equiv_zmod k).injective]
  rw [Fintype.card_subtype]

theorem hasNatDensity_contractingResidues_ge_wordRatio (k : ℕ) :
    HasNatDensity {n : ℕ | branchThresholdMax k ≤ n ∧
        (n : ZMod (2 ^ k)) ∈ contractingResidues k}
      (((Fintype.card {w : ParityWord k // 3 ^ numOdd w < 2 ^ k}) : ℝ) / (2 ^ k : ℝ)) := by
  simpa [card_contractingResidues k] using hasNatDensity_contractingResidues_ge k

theorem finiteStoppingTimeSet_natDensity_one :
    HasNatDensity finiteStoppingTimeSet 1 := by
  apply HasNatDensity.of_superset_of_tendsto_one
    (A := fun k =>
      {n : ℕ | branchThresholdMax k ≤ n ∧ (n : ZMod (2 ^ k)) ∈ contractingResidues k})
    (d := fun k =>
      ((Fintype.card {w : ParityWord k // 3 ^ numOdd w < 2 ^ k}) : ℝ) / (2 ^ k : ℝ))
  · intro k
    exact contractingResidues_ge_subset_finiteStoppingTimeSet k
  · intro k
    exact hasNatDensity_contractingResidues_ge_wordRatio k
  · exact card_contracting_words_asymptotic

theorem finiteStoppingTime_natDensity_one :
    HasNatDensity {n : ℕ | ∃ k : ℕ, accelerated^[k] n < n} 1 := by
  exact finiteStoppingTimeSet_natDensity_one

end Terras
end Erdos1135