/-
SPDX-FileCopyrightText: 2026 Advameg, Inc.
SPDX-License-Identifier: Apache-2.0
-/

import Mathlib.Data.Finset.Interval
import Mathlib.Algebra.Field.ZMod
import Mathlib.Data.ZMod.Basic
import Mathlib.FieldTheory.Finite.Basic
import Mathlib.NumberTheory.Basic
import Mathlib.Tactic

/-!
# Wolstenholme Theorem Seed

This file prepares a proof-facing statement for Wolstenholme's theorem using a
finite reciprocal sum in `ZMod (p ^ 2)`. The statement is deliberately
congruence-native so proof work can proceed through units in modular rings.
-/

open scoped BigOperators

namespace AtlasKnownTheorems.WolstenholmeTheorem

/--
The reciprocal harmonic sum `∑_{k=1}^{p-1} 1/k` in `ZMod (p^2)`. For a prime
`p > 3`, every `k` in this range is a unit modulo `p^2`; the current
definition uses field-style inverse notation in the ring.
-/
noncomputable def harmonicReciprocalSumMod (p : ℕ) : ZMod (p ^ 2) :=
  ∑ k ∈ Finset.Icc 1 (p - 1), ((k : ZMod (p ^ 2))⁻¹)

/--
Minimum Wolstenholme target: for every prime `p > 3`, the modular harmonic
reciprocal sum is zero modulo `p^2`.
-/
def WolstenholmeTheoremStatement : Prop :=
  ∀ p : ℕ, p.Prime → 3 < p → harmonicReciprocalSumMod p = 0

theorem harmonicReciprocalSumMod_zero :
    harmonicReciprocalSumMod 0 = 0 := by
  simp [harmonicReciprocalSumMod]

theorem harmonicReciprocalSumMod_one :
    harmonicReciprocalSumMod 1 = 0 := by
  simp [harmonicReciprocalSumMod]

theorem denominator_isUnit {p k : ℕ} (hp : p.Prime)
    (hk : k ∈ Finset.Icc 1 (p - 1)) :
    IsUnit (k : ZMod (p ^ 2)) := by
  rcases Finset.mem_Icc.mp hk with ⟨hk1, hkle⟩
  rw [ZMod.isUnit_iff_coprime]
  exact hp.coprime_pow_of_not_dvd (m := 2) (by
    intro hpk
    have hkpos : 0 < k := by
      omega
    have hle : p ≤ k := Nat.le_of_dvd hkpos hpk
    have hlt : k < p := by
      omega
    omega)

theorem denominator_mul_inv {p k : ℕ} (hp : p.Prime)
    (hk : k ∈ Finset.Icc 1 (p - 1)) :
    (k : ZMod (p ^ 2)) * (k : ZMod (p ^ 2))⁻¹ = 1 :=
  ZMod.mul_inv_of_unit _ (denominator_isUnit hp hk)

theorem denominator_inv_mul {p k : ℕ} (hp : p.Prime)
    (hk : k ∈ Finset.Icc 1 (p - 1)) :
    (k : ZMod (p ^ 2))⁻¹ * (k : ZMod (p ^ 2)) = 1 :=
  ZMod.inv_mul_of_unit _ (denominator_isUnit hp hk)

theorem two_isUnit_zmod_prime_sq {p : ℕ} (hp : p.Prime) (hp3 : 3 < p) :
    IsUnit (2 : ZMod (p ^ 2)) := by
  change IsUnit ((2 : ℕ) : ZMod (p ^ 2))
  rw [ZMod.isUnit_iff_coprime]
  exact hp.coprime_pow_of_not_dvd (m := 2) (a := 2) (by
    intro hp2
    have hle : p ≤ 2 := Nat.le_of_dvd (by norm_num : 0 < (2 : ℕ)) hp2
    omega)

theorem interval_sub_mem {p k : ℕ}
    (hk : k ∈ Finset.Icc 1 (p - 1)) :
    p - k ∈ Finset.Icc 1 (p - 1) := by
  rcases Finset.mem_Icc.mp hk with ⟨hk1, hkle⟩
  exact Finset.mem_Icc.mpr ⟨by omega, by omega⟩

theorem interval_sub_sub {p k : ℕ}
    (hk : k ∈ Finset.Icc 1 (p - 1)) :
    p - (p - k) = k := by
  rcases Finset.mem_Icc.mp hk with ⟨hk1, hkle⟩
  omega

theorem sum_Icc_sub_eq {p : ℕ} {M : Type} [AddCommMonoid M]
    (f : ℕ → M) :
    (∑ k ∈ Finset.Icc 1 (p - 1), f (p - k)) =
      ∑ k ∈ Finset.Icc 1 (p - 1), f k := by
  refine Finset.sum_bij (fun k _ => p - k) ?_ ?_ ?_ ?_
  · intro k hk
    exact interval_sub_mem hk
  · intro a₁ ha₁ a₂ ha₂ h
    have h' := congrArg (fun x => p - x) h
    have h₁ : p - (p - a₁) = a₁ := interval_sub_sub ha₁
    have h₂ : p - (p - a₂) = a₂ := interval_sub_sub ha₂
    simpa [h₁, h₂] using h'
  · intro b hb
    refine ⟨p - b, interval_sub_mem hb, interval_sub_sub hb⟩
  · intro a ha
    rfl

theorem reciprocal_pair_sum_eq {p k : ℕ} (hp : p.Prime)
    (hk : k ∈ Finset.Icc 1 (p - 1)) :
    (k : ZMod (p ^ 2))⁻¹ + ((p - k : ℕ) : ZMod (p ^ 2))⁻¹ =
      (p : ZMod (p ^ 2)) *
        (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹) := by
  let R := ZMod (p ^ 2)
  let a : R := k
  let b : R := (p - k : ℕ)
  have ha : IsUnit a := denominator_isUnit hp hk
  have hb : IsUnit b := denominator_isUnit hp (interval_sub_mem hk)
  have hab : IsUnit (a * b) := ha.mul hb
  apply hab.mul_left_cancel
  have hkp : k ≤ p := by
    rcases Finset.mem_Icc.mp hk with ⟨hk1, hkle⟩
    omega
  have hab_inv : (a * b) * (a * b)⁻¹ = (1 : R) :=
    ZMod.mul_inv_of_unit (a * b) hab
  have hsum : a + b = (p : R) := by
    dsimp [a, b, R]
    rw [Nat.cast_sub hkp]
    ring
  calc
    (a * b) * (a⁻¹ + b⁻¹) = a + b := by
      have ha_cancel : a * a⁻¹ = (1 : R) := ZMod.mul_inv_of_unit a ha
      have hb_cancel : b * b⁻¹ = (1 : R) := ZMod.mul_inv_of_unit b hb
      rw [mul_add]
      rw [show a * b * a⁻¹ = b by
        calc
          a * b * a⁻¹ = b * (a * a⁻¹) := by ring
          _ = b := by rw [ha_cancel, mul_one]]
      rw [show a * b * b⁻¹ = a by
        calc
          a * b * b⁻¹ = a * (b * b⁻¹) := by ring
          _ = a := by rw [hb_cancel, mul_one]]
      rw [add_comm]
    _ = (p : R) := hsum
    _ = (a * b) * ((p : R) * (a * b)⁻¹) := by
      symm
      calc
        (a * b) * ((p : R) * (a * b)⁻¹) =
            (p : R) * ((a * b) * (a * b)⁻¹) := by
          ring
        _ = (p : R) := by
          rw [hab_inv, mul_one]

theorem two_mul_harmonicReciprocalSumMod (p : ℕ) :
    (2 : ZMod (p ^ 2)) * harmonicReciprocalSumMod p =
      ∑ k ∈ Finset.Icc 1 (p - 1),
        ((k : ZMod (p ^ 2))⁻¹ + ((p - k : ℕ) : ZMod (p ^ 2))⁻¹) := by
  rw [harmonicReciprocalSumMod]
  calc
    (2 : ZMod (p ^ 2)) *
        (∑ k ∈ Finset.Icc 1 (p - 1), (k : ZMod (p ^ 2))⁻¹)
        =
      (∑ k ∈ Finset.Icc 1 (p - 1), (k : ZMod (p ^ 2))⁻¹) +
        ∑ k ∈ Finset.Icc 1 (p - 1), (k : ZMod (p ^ 2))⁻¹ := by
        ring
    _ =
      (∑ k ∈ Finset.Icc 1 (p - 1), (k : ZMod (p ^ 2))⁻¹) +
        ∑ k ∈ Finset.Icc 1 (p - 1),
          ((p - k : ℕ) : ZMod (p ^ 2))⁻¹ := by
        rw [sum_Icc_sub_eq (p := p)
          (f := fun k : ℕ => (k : ZMod (p ^ 2))⁻¹)]
    _ =
      ∑ k ∈ Finset.Icc 1 (p - 1),
        ((k : ZMod (p ^ 2))⁻¹ + ((p - k : ℕ) : ZMod (p ^ 2))⁻¹) := by
        rw [Finset.sum_add_distrib]

theorem two_mul_harmonicReciprocalSumMod_eq_p_mul_sum (p : ℕ)
    (hp : p.Prime) :
    (2 : ZMod (p ^ 2)) * harmonicReciprocalSumMod p =
      (p : ZMod (p ^ 2)) *
        ∑ k ∈ Finset.Icc 1 (p - 1),
          (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹) := by
  rw [two_mul_harmonicReciprocalSumMod]
  calc
    (∑ k ∈ Finset.Icc 1 (p - 1),
        ((k : ZMod (p ^ 2))⁻¹ + ((p - k : ℕ) : ZMod (p ^ 2))⁻¹)) =
      ∑ k ∈ Finset.Icc 1 (p - 1),
        (p : ZMod (p ^ 2)) *
          (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹) := by
        refine Finset.sum_congr rfl ?_
        intro k hk
        exact reciprocal_pair_sum_eq hp hk
    _ =
      (p : ZMod (p ^ 2)) *
        ∑ k ∈ Finset.Icc 1 (p - 1),
          (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹) := by
        rw [← Finset.mul_sum]

theorem harmonicReciprocalSumMod_eq_zero_of_p_mul_pair_sum_eq_zero
    {p : ℕ} (hp : p.Prime) (hp3 : 3 < p)
    (hpair :
      (p : ZMod (p ^ 2)) *
        ∑ k ∈ Finset.Icc 1 (p - 1),
          (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹) = 0) :
    harmonicReciprocalSumMod p = 0 := by
  have htwo :
      (2 : ZMod (p ^ 2)) * harmonicReciprocalSumMod p = 0 := by
    rw [two_mul_harmonicReciprocalSumMod_eq_p_mul_sum p hp]
    exact hpair
  apply (two_isUnit_zmod_prime_sq hp hp3).mul_left_cancel
  simpa using htwo

theorem zmod_prime_units_sum_sq_eq_zero {p : ℕ} [NeZero p]
    (hp : p.Prime) (hp3 : 3 < p) :
    (∑ x : (ZMod p)ˣ, ((x : ZMod p) ^ 2)) = 0 := by
  letI : Fact p.Prime := ⟨hp⟩
  have hnotp : ¬ p - 1 ∣ 2 := by
    intro hd
    have hle : p - 1 ≤ 2 := Nat.le_of_dvd (by norm_num : 0 < (2 : ℕ)) hd
    omega
  simpa [hnotp, ZMod.card] using FiniteField.sum_pow_units (ZMod p) 2

theorem p_mul_eq_zero_of_castHom_eq_zero {p : ℕ} (hp : p.Prime)
    (x : ZMod (p ^ 2))
    (hx :
      ZMod.castHom (show p ∣ p ^ 2 by exact dvd_pow_self p (by norm_num))
        (ZMod p) x = 0) :
    (p : ZMod (p ^ 2)) * x = 0 := by
  haveI : NeZero (p ^ 2) := ⟨pow_ne_zero 2 hp.ne_zero⟩
  have hxval : (x.val : ZMod p) = 0 := by
    rw [← hx]
    rw [← ZMod.natCast_zmod_val x]
    simp
  have hp_dvd_val : p ∣ x.val :=
    (ZMod.natCast_eq_zero_iff x.val p).mp hxval
  rcases hp_dvd_val with ⟨d, hd⟩
  rw [← ZMod.natCast_zmod_val x, hd]
  calc
    (p : ZMod (p ^ 2)) * (↑(p * d) : ZMod (p ^ 2)) =
        ((p ^ 2 * d : ℕ) : ZMod (p ^ 2)) := by
      rw [← Nat.cast_mul]
      congr 1
      ring
    _ = 0 := by
      rw [Nat.cast_mul, ZMod.natCast_self, zero_mul]

theorem castHom_inv_of_isUnit {p : ℕ} (a : ZMod (p ^ 2))
    (ha : IsUnit a) :
    ZMod.castHom (show p ∣ p ^ 2 by exact dvd_pow_self p (by norm_num))
      (ZMod p) a⁻¹ =
      (ZMod.castHom (show p ∣ p ^ 2 by exact dvd_pow_self p (by norm_num))
        (ZMod p) a)⁻¹ := by
  symm
  apply ZMod.inv_eq_of_mul_eq_one
  rw [← map_mul]
  rw [ZMod.mul_inv_of_unit a ha]
  simp

theorem castHom_pair_denominator_inv {p k : ℕ} (hp : p.Prime)
    (hk : k ∈ Finset.Icc 1 (p - 1)) :
    ZMod.castHom (show p ∣ p ^ 2 by exact dvd_pow_self p (by norm_num))
      (ZMod p)
      (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹) =
      -((k : ZMod p)⁻¹ ^ 2) := by
  letI : Fact p.Prime := ⟨hp⟩
  have hunit :
      IsUnit ((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2))) :=
    (denominator_isUnit hp hk).mul (denominator_isUnit hp (interval_sub_mem hk))
  have hkp : k ≤ p := by
    rcases Finset.mem_Icc.mp hk with ⟨hk1, hkle⟩
    omega
  rw [castHom_inv_of_isUnit _ hunit]
  simp only [map_mul, ZMod.castHom_apply]
  rw [ZMod.cast_natCast (R := ZMod p)
    (show p ∣ p ^ 2 by exact dvd_pow_self p (by norm_num)) k]
  rw [ZMod.cast_natCast (R := ZMod p)
    (show p ∣ p ^ 2 by exact dvd_pow_self p (by norm_num)) (p - k)]
  change (((k : ZMod p) * ((p - k : ℕ) : ZMod p))⁻¹) = -((k : ZMod p)⁻¹ ^ 2)
  have hsub : ((p - k : ℕ) : ZMod p) = -(k : ZMod p) := by
    rw [Nat.cast_sub hkp]
    simp
  rw [hsub]
  rw [show (k : ZMod p) * -(k : ZMod p) = -((k : ZMod p) ^ 2) by ring]
  rw [inv_neg]
  congr 1
  rw [show ((k : ZMod p) ^ 2)⁻¹ = (k : ZMod p)⁻¹ ^ 2 by
    rw [pow_two, pow_two]
    rw [mul_inv_rev]]

theorem sum_Icc_inv_sq_eq_sum_units_inv_sq {p : ℕ} [NeZero p]
    (hp : p.Prime) :
    (∑ k ∈ Finset.Icc 1 (p - 1), ((k : ZMod p)⁻¹ ^ 2)) =
      ∑ u : (ZMod p)ˣ, ((u : ZMod p)⁻¹ ^ 2) := by
  refine Finset.sum_bij
    (fun k hk =>
      ZMod.unitOfCoprime k ((hp.coprime_iff_not_dvd.mpr (by
        intro hpk
        rcases Finset.mem_Icc.mp hk with ⟨hk1, hkle⟩
        have hkpos : 0 < k := by omega
        have hle : p ≤ k := Nat.le_of_dvd hkpos hpk
        have hlt : k < p := by omega
        omega)).symm))
    ?_ ?_ ?_ ?_
  · intro k hk
    exact Finset.mem_univ _
  · intro a₁ ha₁ a₂ ha₂ h
    have hval := congrArg (fun u : (ZMod p)ˣ => (u : ZMod p).val) h
    rcases Finset.mem_Icc.mp ha₁ with ⟨ha₁1, ha₁le⟩
    rcases Finset.mem_Icc.mp ha₂ with ⟨ha₂1, ha₂le⟩
    have ha₁lt : a₁ < p := by omega
    have ha₂lt : a₂ < p := by omega
    simpa [ZMod.coe_unitOfCoprime, ZMod.val_cast_of_lt ha₁lt,
      ZMod.val_cast_of_lt ha₂lt] using hval
  · intro u hu
    have hcop : (u : ZMod p).val.Coprime p := ZMod.val_coe_unit_coprime u
    have hval_ne_zero : (u : ZMod p).val ≠ 0 := by
      intro hv
      rw [hv] at hcop
      simp at hcop
      exact hp.ne_one hcop
    have hmem : (u : ZMod p).val ∈ Finset.Icc 1 (p - 1) := by
      exact Finset.mem_Icc.mpr ⟨by omega, by
        have hlt := ZMod.val_lt (u : ZMod p)
        omega⟩
    refine ⟨(u : ZMod p).val, hmem, ?_⟩
    apply Units.ext
    simp [ZMod.coe_unitOfCoprime]
  · intro k hk
    simp [ZMod.coe_unitOfCoprime]

theorem sum_units_inv_sq_eq_sum_units_sq {p : ℕ} [NeZero p] :
    (∑ u : (ZMod p)ˣ, ((u : ZMod p)⁻¹ ^ 2)) =
      ∑ u : (ZMod p)ˣ, ((u : ZMod p) ^ 2) := by
  simpa [ZMod.inv_coe_unit] using
    (Equiv.sum_comp (Equiv.inv ((ZMod p)ˣ))
      (fun u : (ZMod p)ˣ => ((u : ZMod p) ^ 2)))

theorem castHom_pair_sum_eq_zero {p : ℕ} [NeZero p]
    (hp : p.Prime) (hp3 : 3 < p) :
    ZMod.castHom (show p ∣ p ^ 2 by exact dvd_pow_self p (by norm_num))
      (ZMod p)
      (∑ k ∈ Finset.Icc 1 (p - 1),
        (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹)) = 0 := by
  letI : Fact p.Prime := ⟨hp⟩
  calc
    ZMod.castHom (show p ∣ p ^ 2 by exact dvd_pow_self p (by norm_num))
      (ZMod p)
      (∑ k ∈ Finset.Icc 1 (p - 1),
        (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹))
        =
      ∑ k ∈ Finset.Icc 1 (p - 1),
        -((k : ZMod p)⁻¹ ^ 2) := by
        rw [map_sum]
        refine Finset.sum_congr rfl ?_
        intro k hk
        exact castHom_pair_denominator_inv hp hk
    _ = -∑ k ∈ Finset.Icc 1 (p - 1), ((k : ZMod p)⁻¹ ^ 2) := by
      rw [Finset.sum_neg_distrib]
    _ = 0 := by
      rw [sum_Icc_inv_sq_eq_sum_units_inv_sq hp]
      rw [sum_units_inv_sq_eq_sum_units_sq]
      rw [zmod_prime_units_sum_sq_eq_zero hp hp3]
      simp

theorem p_mul_pair_sum_eq_zero {p : ℕ} (hp : p.Prime) (hp3 : 3 < p) :
    (p : ZMod (p ^ 2)) *
      ∑ k ∈ Finset.Icc 1 (p - 1),
        (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹) = 0 := by
  haveI : NeZero p := ⟨hp.ne_zero⟩
  exact p_mul_eq_zero_of_castHom_eq_zero hp
    (∑ k ∈ Finset.Icc 1 (p - 1),
      (((k : ZMod (p ^ 2)) * ((p - k : ℕ) : ZMod (p ^ 2)))⁻¹))
    (castHom_pair_sum_eq_zero hp hp3)

/-- Wolstenholme's theorem in modular harmonic reciprocal-sum form. -/
theorem wolstenholmeTheorem :
    WolstenholmeTheoremStatement := by
  intro p hp hp3
  exact harmonicReciprocalSumMod_eq_zero_of_p_mul_pair_sum_eq_zero hp hp3
    (p_mul_pair_sum_eq_zero hp hp3)

/-- Checked statement wrapper for the minimum theorem. -/
theorem wolstenholmeTheorem_statement :
    WolstenholmeTheoremStatement ↔ WolstenholmeTheoremStatement := by
  rfl

end AtlasKnownTheorems.WolstenholmeTheorem