import Mathlib.Data.Int.ModEq
import Mathlib.Data.Nat.GCD.Basic
import Mathlib.Data.Finset.Basic
import Mathlib.Data.List.Basic
import Mathlib.Tactic

open Finset

namespace CoreFormalism.CRTSidonN

/-!
# CRT Sidon Preservation: n-moduli generalization

This module extends the 2-moduli `sidon_preserved_mod` theorem from
`CRTSidon.lean` to the general n-moduli case.

**Theorem (n-moduli CRT Sidon Preservation):**
If A is a Sidon set, moduli L₀, L₁, ..., Lₖ₋₁ are pairwise coprime,
all pairwise sums a+b < M = ∏Lᵢ, S ≥ all labels, and:
  (a+b) % L₀ = (c+d) % L₀                    [identity component]
  ((S-a)+(S-b)) % Lᵢ = ((S-c)+(S-d)) % Lᵢ    [reflection components, i ≥ 1]
then {a,b} = {c,d}.

**Proof strategy:** By the generalized CRT uniqueness theorem:
if x ≡ y (mod Lᵢ) for all i, pairwise coprime moduli, x,y < ∏Lᵢ,
then x = y. This is proven by induction on the moduli list, using
the 2-moduli case (`mod_eq_of_coprime`) as the inductive step.

The reflection components (i ≥ 1) all reduce to (a+b) ≡ (c+d) (mod Lᵢ)
via the same algebraic manipulation as the 2-moduli case:
  (S-a) + (S-b) = 2S - (a+b), so
  (2S - (a+b)) % Lᵢ = (2S - (c+d)) % Lᵢ implies (a+b) ≡ (c+d) (mod Lᵢ). -/


/-- A Sidon set: all unordered pairwise sums are distinct. -/
def IsSidon (A : Finset ℕ) : Prop :=
  ∀ ⦃a b c d : ℕ⦄, a ∈ A → b ∈ A → c ∈ A → d ∈ A → a + b = c + d →
    (a = c ∧ b = d) ∨ (a = d ∧ b = c)

/-- Pairwise coprime list. -/
def PairwiseCoprime (ls : List ℕ) : Prop :=
  ∀ i j (hi : i < ls.length) (hj : j < ls.length), i < j →
    (ls.get ⟨i, hi⟩).gcd (ls.get ⟨j, hj⟩) = 1

/-- All elements of a list are positive. -/
def AllPos (ls : List ℕ) : Prop :=
  ∀ i (hi : i < ls.length), 0 < ls.get ⟨i, hi⟩

/-- Convert ℕ modular equality to ℤ divisibility: a%n = b%n → n ∣ (a-b) in ℤ. -/
lemma mod_eq_dvd (a b n : ℕ) (h : a % n = b % n) : (n : ℤ) ∣ ((a : ℤ) - (b : ℤ)) := by
  have hz : (a : ℤ) % n = (b : ℤ) % n := by exact_mod_cast h
  have h_eq : (a : ℤ) ≡ (b : ℤ) [ZMOD n] := hz
  have h_sub : (a : ℤ) - (b : ℤ) ≡ 0 [ZMOD n] := by
    calc
      (a : ℤ) - (b : ℤ) ≡ (b : ℤ) - (b : ℤ) [ZMOD n] := Int.ModEq.sub h_eq (Int.ModEq.refl _)
      _ = 0 := by ring
  exact (Int.modEq_zero_iff_dvd.mp h_sub)

/-- Convert ℤ divisibility to ℕ modular equality. -/
lemma dvd_mod_eq (a b n : ℕ) (h : (n : ℤ) ∣ ((a : ℤ) - (b : ℤ))) : a % n = b % n := by
  have h_sub : ((a : ℤ) - (b : ℤ)) ≡ 0 [ZMOD n] := by
    rw [Int.modEq_zero_iff_dvd]
    exact h
  have h_mod : (a : ℤ) ≡ (b : ℤ) [ZMOD n] := by
    calc
      (a : ℤ) = ((a : ℤ) - (b : ℤ)) + (b : ℤ) := by ring
      _ ≡ 0 + (b : ℤ) [ZMOD n] := Int.ModEq.add h_sub (Int.ModEq.refl _)
      _ = (b : ℤ) := by ring
  exact_mod_cast h_mod

/-! ## Generalized CRT Uniqueness -/

/-- If gcd(m, n) = 1 and n ∣ d, then m*n ∣ m*d. -/
lemma dvd_mul_of_dvd_right {m n d : ℕ} (hmn : Nat.Coprime m n) (hn : n ∣ d) :
    m * n ∣ m * d := by
  obtain ⟨q, hq⟩ := hn
  exact ⟨q, by rw [hq, mul_assoc]⟩

/-! ### Key lemma: product of coprime moduli

If L is a pairwise coprime list, and each Lᵢ divides d, then
∏Lᵢ divides d. This is the heart of the generalized CRT.

We prove this by strong induction on the list length.

-/
/-- If L₀ is coprime to every element of tail, then L₀ is coprime to the product. -/
lemma coprime_to_product (L₀ : ℕ) : ∀ (tail : List ℕ), (∀ x ∈ tail, Nat.Coprime L₀ x) → Nat.Coprime L₀ tail.prod := by
  intro tail hall
  induction tail with
  | nil => simp
  | cons h t ih =>
    have h_coprime_h : Nat.Coprime L₀ h := hall h (by simp)
    have h_coprime_t : ∀ x ∈ t, Nat.Coprime L₀ x := by
      intro x hx
      exact hall x (by simp [hx])
    simpa [List.prod_cons] using (Nat.Coprime.mul_right h_coprime_h (ih h_coprime_t))

/-- If (L₀ :: tail) is pairwise coprime, then L₀ is coprime to every element of tail. -/
lemma pairwise_coprime_cons_all_coprime (L₀ : ℕ) (tail : List ℕ) (hCoprime : PairwiseCoprime (L₀ :: tail)) :
    ∀ x ∈ tail, Nat.Coprime L₀ x := by
  intro x hx
  rcases List.mem_iff_get.mp hx with ⟨n, hn⟩
  -- n : Fin tail.length, hn : tail.get n = x
  have hi_len : n.1 < tail.length := n.2
  have h_src : n.1 + 1 < (L₀ :: tail).length := by
    simpa [List.length_cons] using hi_len
  have h := hCoprime 0 (n.1 + 1) (by simp) h_src (by omega)
  -- h simplifies to L₀.gcd (tail.get n) = 1
  have h_gcd : L₀.gcd (tail.get n) = 1 := by
    simpa using h
  -- Need L₀.Coprime x, i.e., L₀.gcd x = 1
  rw [hn] at h_gcd
  exact h_gcd

/-- Product of a list of ℕ. -/
def listProd (ls : List ℕ) : ℕ := ls.prod

/-- If all elements in a list divide d, and the list is pairwise coprime,
    then the product divides d.

    Proof by induction on the list:
    - Base case ([]): product = 1, 1 ∣ d trivially
    - Inductive step (L₀ :: tail): L₀ ∣ d and (∏tail) ∣ d
      By IH: (∏tail) ∣ d
      Since L₀ is coprime to each element of tail, L₀ is coprime to ∏tail
      (coprime to product = coprime to each factor)
      Since L₀ ∣ d and (∏tail) ∣ d and gcd(L₀, ∏tail) = 1:
    -/
theorem pairwise_coprime_product_dvd (L : List ℕ) (hCoprime : PairwiseCoprime L)
    (hDiv : ∀ i (hi : i < L.length), (L.get ⟨i, hi⟩) ∣ d) :
    L.prod ∣ d := by
  induction L using List.rec with
  | nil => simp [List.prod, Nat.one_dvd]
  | cons L₀ tail IH =>
    -- L = L₀ :: tail, product = L₀ * tail.prod
    -- L₀ ∣ d (from hDiv at index 0)
    have hL0_dvd : L₀ ∣ d := hDiv 0 (by
      simpa using Nat.zero_lt_succ (tail.length))
    -- tail elements divide d (from hDiv at shifted indices)
    have hTail_dvd : ∀ i (hi : i < tail.length), (tail.get ⟨i, hi⟩) ∣ d := by
      intro i hi
      exact hDiv (i + 1) (by simp [hi])
    -- tail is pairwise coprime (inherited from L)
    have hTailCoprime : PairwiseCoprime tail := by
      intro i j hi hj hij
      exact hCoprime (i + 1) (j + 1) (by simp [hi]) (by simp [hj]) (by omega)
    -- By IH: tail.prod ∣ d
    have hTailProd_dvd : tail.prod ∣ d := IH hTailCoprime hTail_dvd
    -- L₀ is coprime to tail.prod
    -- (because L₀ is coprime to each element of tail)
    have hL0_coprime_tail_prod : Nat.Coprime L₀ tail.prod := by
      apply coprime_to_product L₀ tail
      apply pairwise_coprime_cons_all_coprime L₀ tail hCoprime
    -- Since L₀ ∣ d and tail.prod ∣ d and gcd(L₀, tail.prod) = 1:
    -- L₀ * tail.prod ∣ d
    -- Use: Nat.Coprime.dvd_mul or Nat.mul_dvd_of_coprime
    have hprod_dvd : L₀ * tail.prod ∣ d :=
      Nat.Coprime.mul_dvd_of_dvd_of_dvd hL0_coprime_tail_prod hL0_dvd hTailProd_dvd
    -- product of (L₀ :: tail) = L₀ * tail.prod
    simpa [List.prod_cons] using hprod_dvd

/-! ### Generalized CRT uniqueness (n moduli) -/

/-- If a ≡ b (mod Lᵢ) for all i, L is pairwise coprime, and a, b < ∏Lᵢ,
    then a = b.

    This generalizes `mod_eq_of_coprime` from 2 moduli to n moduli. -/
theorem mod_eq_of_coprime_list {a b : ℕ} (L : List ℕ)
    (hCoprime : PairwiseCoprime L) (hPos : AllPos L)
    (hcong : ∀ i (hi : i < L.length), a % (L.get ⟨i, hi⟩) = b % (L.get ⟨i, hi⟩))
    (ha : a < L.prod) (hb : b < L.prod) : a = b := by
  -- Each Lᵢ divides (a - b) in ℤ
  have hL_dvd : ∀ i (hi : i < L.length), (L.get ⟨i, hi⟩ : ℤ) ∣ ((a : ℤ) - (b : ℤ)) := by
    intro i hi
    exact mod_eq_dvd a b (L.get ⟨i, hi⟩) (hcong i hi)
  -- Convert to ℕ divisibility (need to handle the ℤ → ℕ direction)
  -- Actually, we need: each Lᵢ divides |a - b| in ℕ
  -- Since a, b ∈ ℕ, either a ≥ b or b > a
  -- Case 1: a ≥ b → d = a - b ≥ 0, each Lᵢ ∣ d in ℕ
  -- Case 2: b > a → d = b - a ≥ 0, each Lᵢ ∣ d in ℕ (by symmetry)
  -- In both cases, ∏Lᵢ ∣ d, and d < ∏Lᵢ, so d = 0, i.e., a = b
  by_cases hab : a ≥ b
  · -- a ≥ b: d = a - b
    set d := a - b
    -- Each Lᵢ ∣ d in ℕ
    have hL_dvd_nat : ∀ i (hi : i < L.length), (L.get ⟨i, hi⟩) ∣ d := by
      intro i hi
      have h := hL_dvd i hi
      have hd_eq : ((a : ℤ) - (b : ℤ)) = (d : ℤ) := by
        dsimp [d]
        have hsub : a - b = d := rfl
        omega
      rw [hd_eq] at h
      rw [Int.natCast_dvd_natCast] at h
      exact h
    -- ∏Lᵢ ∣ d
    have hprod_dvd : L.prod ∣ d := pairwise_coprime_product_dvd L hCoprime hL_dvd_nat
    -- |a - b| < ∏Lᵢ (since a, b < ∏Lᵢ)
    have hd_lt : d < L.prod := by
      dsimp [d]
      -- a - b ≤ a < L.prod  (follows from hab : a ≥ b and ha : a < L.prod)
      omega
    -- ∏Lᵢ ∣ d and d < ∏Lᵢ → d = 0 → a = b
    obtain ⟨q, hq⟩ := hprod_dvd
    by_cases hq0 : q = 0
    · -- q = 0 → d = 0 → a = b
      dsimp [d] at hq
      rw [hq0, mul_zero] at hq
      omega
    · -- q ≥ 1 → d = ∏Lᵢ * q ≥ ∏Lᵢ > d (contradiction)
      have hprod_le : L.prod ≤ d := by
        rw [hq]
        exact Nat.mul_le_mul_left L.prod (by omega : 1 ≤ q)
      omega
  · -- b > a: d = b - a, symmetric argument
    set d := b - a
    have hd_eq : ((a : ℤ) - (b : ℤ)) = -(d : ℤ) := by
      dsimp [d]; omega
    -- Each Lᵢ ∣ (b - a) in ℕ (by symmetry of modular arithmetic)
    have hL_dvd_nat : ∀ i (hi : i < L.length), (L.get ⟨i, hi⟩) ∣ d := by
      intro i hi
      have h := hL_dvd i hi
      rw [hd_eq] at h
      -- (Lᵢ : ℤ) ∣ -(d : ℤ) → (Lᵢ : ℤ) ∣ (d : ℤ)
      have h' : (L.get ⟨i, hi⟩ : ℤ) ∣ (d : ℤ) := Int.dvd_neg.mp h
      rwa [Int.natCast_dvd_natCast] at h'
    have hprod_dvd : L.prod ∣ d := pairwise_coprime_product_dvd L hCoprime hL_dvd_nat
    have hd_lt : d < L.prod := by
      dsimp [d]
      omega
    obtain ⟨q, hq⟩ := hprod_dvd
    by_cases hq0 : q = 0
    · dsimp [d] at hq
      rw [hq0, mul_zero] at hq
      omega
    · have hprod_le : L.prod ≤ d := by
        rw [hq]
        have : 1 ≤ q := by omega
        calc
          L.prod = L.prod * 1 := by simp
          _ ≤ L.prod * q := Nat.mul_le_mul_left L.prod this
      omega

/-! ### Reflection → sum congruence -/

/-- The reflection component (S-a)+(S-b) ≡ (S-c)+(S-d) (mod Lᵢ)
    implies (a+b) ≡ (c+d) (mod Lᵢ).

    This is the same algebraic manipulation as the 2-moduli case:
    (S-a)+(S-b) = 2S-(a+b), so
    (2S-(a+b)) % Lᵢ = (2S-(c+d)) % Lᵢ
    → (a+b) ≡ (c+d) (mod Lᵢ)  -/
lemma reflection_implies_sum_cong {a b c d S Lᵢ : ℕ}
    (hS_a : a ≤ S) (hS_b : b ≤ S) (hS_c : c ≤ S) (hS_d : d ≤ S)
    (h_ref : ((S - a) + (S - b)) % Lᵢ = ((S - c) + (S - d)) % Lᵢ) :
    (a + b) % Lᵢ = (c + d) % Lᵢ := by
  -- (S-a) + (S-b) = 2S - (a+b)
  have h_ab : (S - a) + (S - b) = 2 * S - (a + b) := by omega
  have h_cd : (S - c) + (S - d) = 2 * S - (c + d) := by omega
  rw [h_ab, h_cd] at h_ref
  -- h_ref: (2S - (a+b)) % Lᵢ = (2S - (c+d)) % Lᵢ
  -- In ℤ: 2S-(a+b) ≡ 2S-(c+d) (ZMOD Lᵢ) → (a+b) ≡ (c+d) (ZMOD Lᵢ)
  have hz : ((2 * S - (a + b) : ℕ) : ℤ) % Lᵢ = ((2 * S - (c + d) : ℕ) : ℤ) % Lᵢ :=
    by exact_mod_cast h_ref
  have h_eq : ((2 * S - (a + b) : ℕ) : ℤ) ≡ ((2 * S - (c + d) : ℕ) : ℤ) [ZMOD Lᵢ] := hz
  have h_sub : ((2 * S - (c + d) : ℕ) : ℤ) - ((2 * S - (a + b) : ℕ) : ℤ) ≡ 0 [ZMOD Lᵢ] := by
    have h_sub_eq : ((2 * S - (c + d) : ℕ) : ℤ) - ((2 * S - (a + b) : ℕ) : ℤ) ≡
      ((2 * S - (c + d) : ℕ) : ℤ) - ((2 * S - (c + d) : ℕ) : ℤ) [ZMOD Lᵢ] :=
      Int.ModEq.sub (Int.ModEq.refl _) h_eq
    have hzero : ((2 * S - (c + d) : ℕ) : ℤ) - ((2 * S - (c + d) : ℕ) : ℤ) = 0 := by ring
    simpa [hzero] using h_sub_eq
  have h_sub_simp : ((2 * S - (c + d) : ℕ) : ℤ) - ((2 * S - (a + b) : ℕ) : ℤ) =
    ((a + b : ℤ) - (c + d : ℤ)) := by omega
  rw [h_sub_simp] at h_sub
  exact dvd_mod_eq (a + b) (c + d) Lᵢ (Int.modEq_zero_iff_dvd.mp h_sub)

/-! ### Main theorem: n-moduli CRT Sidon preservation -/

/-- **Theorem (n-moduli CRT Sidon Preservation)**.
    If A is a Sidon set, moduli L = [L₀, L₁, ..., Lₖ₋₁] are pairwise coprime
    and positive, all pairwise sums < M = ∏Lᵢ, S ≥ all labels, and:
      - Component 0 (identity): (a+b) % L₀ = (c+d) % L₀
      - Component i ≥ 1 (reflection): ((S-a)+(S-b)) % Lᵢ = ((S-c)+(S-d)) % Lᵢ
    then {a,b} = {c,d}.

    This generalizes `sidon_preserved_mod` from 2 moduli to n moduli. -/
theorem sidon_preserved_mod_n (A : Finset ℕ) (hSidon : IsSidon A) (S : ℕ)
    (L : List ℕ) (hCoprime : PairwiseCoprime L) (hPos : AllPos L) (hNonempty : L ≠ [])
    (hS : ∀ a ∈ A, a ≤ S)
    (hBound : ∀ a ∈ A, ∀ b ∈ A, a + b < L.prod) :
    ∀ ⦃a b c d : ℕ⦄, a ∈ A → b ∈ A → c ∈ A → d ∈ A →
      -- identity component (index 0)
      (a % (L.head hNonempty) + b % (L.head hNonempty)) % (L.head hNonempty) =
      (c % (L.head hNonempty) + d % (L.head hNonempty)) % (L.head hNonempty) →
      -- reflection components (indices ≥ 1)
      (∀ i (hi : 1 ≤ i) (hi' : i < L.length),
        ((S - a) % (L.get ⟨i, hi'⟩) + (S - b) % (L.get ⟨i, hi'⟩)) % (L.get ⟨i, hi'⟩) =
        ((S - c) % (L.get ⟨i, hi'⟩) + (S - d) % (L.get ⟨i, hi'⟩)) % (L.get ⟨i, hi'⟩)) →
        (a = c ∧ b = d) ∨ (a = d ∧ b = c) := by
  intro a b c d ha hb hc hd h_id h_ref
  -- Step 1: identity component → (a+b) % L₀ = (c+d) % L₀
  have hlen0 : 0 < L.length := List.length_pos_of_ne_nil hNonempty
  have hL0_cong : (a + b) % (L.head hNonempty) = (c + d) % (L.head hNonempty) := by
    calc
      (a + b) % (L.head hNonempty) = (a % (L.head hNonempty) + b % (L.head hNonempty)) % (L.head hNonempty) := by rw [Nat.add_mod]
      _ = (c % (L.head hNonempty) + d % (L.head hNonempty)) % (L.head hNonempty) := h_id
      _ = (c + d) % (L.head hNonempty) := by rw [← Nat.add_mod]
  -- Step 2: reflection components → (a+b) % Lᵢ = (c+d) % Lᵢ for all i ≥ 1
  have hRef_cong : ∀ i (hi : 1 ≤ i) (hi' : i < L.length),
      (a + b) % (L.get ⟨i, hi'⟩) = (c + d) % (L.get ⟨i, hi'⟩) := by
    intro i hi hi'
    have hS_a : a ≤ S := hS a ha
    have hS_b : b ≤ S := hS b hb
    have hS_c : c ≤ S := hS c hc
    have hS_d : d ≤ S := hS d hd
    have h_ref_simple : ((S - a) + (S - b)) % (L.get ⟨i, hi'⟩) = ((S - c) + (S - d)) % (L.get ⟨i, hi'⟩) := by
      calc
        ((S - a) + (S - b)) % (L.get ⟨i, hi'⟩) = ((S - a) % (L.get ⟨i, hi'⟩) + (S - b) % (L.get ⟨i, hi'⟩)) % (L.get ⟨i, hi'⟩) := by rw [Nat.add_mod]
        _ = ((S - c) % (L.get ⟨i, hi'⟩) + (S - d) % (L.get ⟨i, hi'⟩)) % (L.get ⟨i, hi'⟩) := h_ref i hi hi'
        _ = ((S - c) + (S - d)) % (L.get ⟨i, hi'⟩) := by rw [← Nat.add_mod]
    exact reflection_implies_sum_cong hS_a hS_b hS_c hS_d h_ref_simple
  -- Step 3: All components congruent: (a+b) ≡ (c+d) (mod Lᵢ) for ALL i
  have hAll_cong : ∀ i (hi : i < L.length),
      (a + b) % (L.get ⟨i, hi⟩) = (c + d) % (L.get ⟨i, hi⟩) := by
    intro i hi
    by_cases hi0 : i = 0
    · subst hi0
      -- L.get ⟨0, hi⟩ = L.head hNonempty (both return first element)
      have h_head_eq : L.get ⟨0, hi⟩ = L.head hNonempty := by
        cases L
        · exfalso; exact hNonempty rfl
        · simp
      rw [h_head_eq]
      exact hL0_cong
    · exact hRef_cong i (by omega) hi
  -- Step 4: By generalized CRT, a+b = c+d (since both < ∏Lᵢ)
  have heq : a + b = c + d := by
    apply mod_eq_of_coprime_list L hCoprime hPos hAll_cong
    · exact hBound a ha b hb
    · exact hBound c hc d hd
  -- Step 5: By Sidon
  rcases hSidon ha hb hc hd heq with (⟨hac, hbd⟩ | ⟨had, hbc⟩)
  · exact Or.inl ⟨hac, hbd⟩
  · exact Or.inr ⟨had, hbc⟩

end CoreFormalism.CRTSidonN
