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

open Finset

namespace CoreFormalism.CRTSidon

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

/-- CRT Torus Embedding of label a with sum parameter S and moduli L₀::Ls. -/
def crtEmbed (a S : ℕ) : List ℕ → List ℕ
  | [] => []
  | L₀ :: Ls => (a % L₀) :: (Ls.map fun Lᵢ => (S - a) % Lᵢ)

/-- Project the first element of a vector (0 for empty). -/
def vecFirst (v : List ℕ) : ℕ :=
  match v with
  | [] => 0
  | x :: _ => x

/-- Componentwise vector sum. -/
def vecAdd (v w : List ℕ) : List ℕ :=
  List.zipWith (· + ·) v w

lemma vecFirst_crtEmbed (a S L₀ : ℕ) (Ls : List ℕ) (ha : a < L₀) :
    vecFirst (crtEmbed a S (L₀ :: Ls)) = a := by
  simp [crtEmbed, vecFirst, Nat.mod_eq_of_lt ha]

lemma vecFirst_vecAdd_crtEmbed (a b S L₀ : ℕ) (Ls : List ℕ) (ha : a < L₀) (hb : b < L₀) :
    vecFirst (vecAdd (crtEmbed a S (L₀ :: Ls)) (crtEmbed b S (L₀ :: Ls))) = a + b := by
  simp [vecAdd, crtEmbed, vecFirst, Nat.mod_eq_of_lt ha, Nat.mod_eq_of_lt hb]

/-- Total modulus M = ∏ Lᵢ. -/
def totalMod (L : List ℕ) : ℕ := L.prod

/-- **Main Theorem (Componentwise)**: CRT Torus embedding preserves the Sidon property
    under componentwise vector addition. -/
theorem sidon_preserved (A : Finset ℕ) (hSidon : IsSidon A) (S L₀ : ℕ) (Ls : List ℕ)
    (hBound : ∀ a ∈ A, a < L₀) :
    ∀ ⦃a b c d : ℕ⦄, a ∈ A → b ∈ A → c ∈ A → d ∈ A →
      vecAdd (crtEmbed a S (L₀ :: Ls)) (crtEmbed b S (L₀ :: Ls)) =
      vecAdd (crtEmbed c S (L₀ :: Ls)) (crtEmbed d S (L₀ :: Ls)) →
        (a = c ∧ b = d) ∨ (a = d ∧ b = c) := by
  intro a b c d ha hb hc hd hvec
  have ha_lt : a < L₀ := hBound a ha
  have hb_lt : b < L₀ := hBound b hb
  have hc_lt : c < L₀ := hBound c hc
  have hd_lt : d < L₀ := hBound d hd
  have hfirst : a + b = c + d := by
    calc
      a + b = vecFirst (vecAdd (crtEmbed a S (L₀ :: Ls)) (crtEmbed b S (L₀ :: Ls))) := by
        symm; exact vecFirst_vecAdd_crtEmbed a b S L₀ Ls ha_lt hb_lt
      _ = vecFirst (vecAdd (crtEmbed c S (L₀ :: Ls)) (crtEmbed d S (L₀ :: Ls))) := by rw [hvec]
      _ = c + d := vecFirst_vecAdd_crtEmbed c d S L₀ Ls hc_lt hd_lt
  rcases hSidon ha hb hc hd hfirst with (⟨hac, hbd⟩ | ⟨had, hbc⟩)
  · exact Or.inl ⟨hac, hbd⟩
  · exact Or.inr ⟨had, hbc⟩

/-- 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: n ∣ (a-b) in ℤ → a%n = b%n in ℕ. -/
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

/-- Two-modulus CRT injectivity: if a ≡ b (mod m) and a ≡ b (mod n) and
    gcd(m,n) = 1, and a,b < m*n, then a = b.
    
    Uses Bezout's identity via Nat.gcdA/Nat.gcdB to prove m*n ∣ (a-b) in ℤ,
    then the bound |a-b| < m*n forces a = b. -/
lemma mod_eq_of_coprime {a b m n : ℕ} (hmn : Nat.Coprime m n) (hm : 0 < m) (hn : 0 < n)
    (ham : a % m = b % m) (han : a % n = b % n) (ha : a < m * n) (hb : b < m * n) :
    a = b := by
  have hm_dvd : (m : ℤ) ∣ ((a : ℤ) - (b : ℤ)) := mod_eq_dvd a b m ham
  have hn_dvd : (n : ℤ) ∣ ((a : ℤ) - (b : ℤ)) := mod_eq_dvd a b n han
  -- Bezout: gcd(m,n)=1 → (Nat.gcdA m n) * m + (Nat.gcdB m n) * n = 1 in ℤ
  have h_gcd_one : Nat.gcd m n = 1 := hmn
  have hbez : (m : ℤ) * Nat.gcdA m n + (n : ℤ) * Nat.gcdB m n = (1 : ℤ) := by
    calc
      (m : ℤ) * Nat.gcdA m n + (n : ℤ) * Nat.gcdB m n = (Nat.gcd m n : ℤ) :=
        (Nat.gcd_eq_gcd_ab m n).symm
      _ = (1 : ℤ) := by exact_mod_cast h_gcd_one
  set d := (a : ℤ) - (b : ℤ) with hd
  -- Show m*n ∣ s*(m*d) (since n ∣ d) and m*n ∣ t*(n*d) (since m ∣ d)
  have h_mn_dvd_md : (m * n : ℤ) ∣ (m : ℤ) * d := by
    obtain ⟨q, hq⟩ := hn_dvd
    refine ⟨q, ?_⟩
    calc
      (m : ℤ) * d = (m : ℤ) * ((n : ℤ) * q) := by rw [hq]
      _ = (m * n : ℤ) * q := by ring
  have h_mn_dvd_nd : (m * n : ℤ) ∣ (n : ℤ) * d := by
    obtain ⟨q, hq⟩ := hm_dvd
    refine ⟨q, ?_⟩
    calc
      (n : ℤ) * d = (n : ℤ) * ((m : ℤ) * q) := by rw [hq]
      _ = (m * n : ℤ) * q := by ring
  -- d = s*(m*d) + t*(n*d), so m*n ∣ d
  have h_d_eq : d = (Nat.gcdA m n : ℤ) * ((m : ℤ) * d) + (Nat.gcdB m n : ℤ) * ((n : ℤ) * d) := by
    calc
      d = (1 : ℤ) * d := by ring
      _ = ((m : ℤ) * Nat.gcdA m n + (n : ℤ) * Nat.gcdB m n) * d := by rw [hbez]
      _ = (Nat.gcdA m n : ℤ) * ((m : ℤ) * d) + (Nat.gcdB m n : ℤ) * ((n : ℤ) * d) := by ring
  have h_mn_dvd_d : (m * n : ℤ) ∣ d := by
    rw [h_d_eq]
    exact dvd_add (dvd_mul_of_dvd_right h_mn_dvd_md _) (dvd_mul_of_dvd_right h_mn_dvd_nd _)
  -- |d| < m*n (since 0 ≤ a,b < m*n), so d = 0
  have ha_z : (a : ℤ) < (m * n : ℤ) := by exact_mod_cast ha
  have hb_z : (b : ℤ) < (m * n : ℤ) := by exact_mod_cast hb
  have ha_nonneg : 0 ≤ (a : ℤ) := by exact_mod_cast (Nat.zero_le _)
  have hb_nonneg : 0 ≤ (b : ℤ) := by exact_mod_cast (Nat.zero_le _)
  have hd_upper : d < (m * n : ℤ) := by
    dsimp [d]
    omega
  have hd_lower : -(m * n : ℤ) < d := by
    dsimp [d]
    omega
  have hd_abs : |d| < (m * n : ℤ) :=
    abs_lt.mpr ⟨hd_lower, hd_upper⟩
  obtain ⟨k, hk⟩ := h_mn_dvd_d
  have h_abs_mul : |(m * n : ℤ) * k| = (m * n : ℤ) * |k| := by simp
  by_cases hk0 : k = 0
  · dsimp [d] at hk
    rw [hk0, mul_zero] at hk
    omega
  · have hk_abs_ge_one : 1 ≤ |k| := by
      have hk_pos : 0 < |k| := abs_pos.mpr hk0
      omega
    have h_abs_ge_mn : |(m * n : ℤ) * k| ≥ (m * n : ℤ) := by
      rw [h_abs_mul]
      have h_nonneg_mn : 0 ≤ (m * n : ℤ) := by positivity
      nlinarith
    have h_abs_d : |d| = |(m * n : ℤ) * k| := by rw [hk]
    rw [h_abs_d] at hd_abs
    omega

/-- **Modular Sidon Preservation Theorem (two-modulus CRT)**.
    If A is a Sidon set, L₀ and L₁ are coprime, all pairwise sums < L₀·L₁,
    S ≥ all labels, and the componentwise modular sums are congruent:
      (a%L₀ + b%L₀) % L₀ = (c%L₀ + d%L₀) % L₀
      ((S-a)%L₁ + (S-b)%L₁) % L₁ = ((S-c)%L₁ + (S-d)%L₁) % L₁
    then {a,b} = {c,d}.

    This matches the `crt_capacity_envelope.py` verification, which checks
    CRT-reconstructed mod-M sums via componentwise congruences.
    The componentwise `sidon_preserved` theorem proves a stronger property. -/
theorem sidon_preserved_mod (A : Finset ℕ) (hSidon : IsSidon A) (S L₀ L₁ : ℕ)
    (hCoprime : Nat.Coprime L₀ L₁) (hL₀ : 0 < L₀) (hL₁ : 0 < L₁)
    (hS : ∀ a ∈ A, a ≤ S)
    (hBound : ∀ a ∈ A, ∀ b ∈ A, a + b < L₀ * L₁) :
    ∀ ⦃a b c d : ℕ⦄, a ∈ A → b ∈ A → c ∈ A → d ∈ A →
      (a % L₀ + b % L₀) % L₀ = (c % L₀ + d % L₀) % L₀ →
      ((S - a) % L₁ + (S - b) % L₁) % L₁ = ((S - c) % L₁ + (S - d) % L₁) % L₁ →
        (a = c ∧ b = d) ∨ (a = d ∧ b = c) := by
  intro a b c d ha hb hc hd h0 h1
  -- Step 1: identity component → (a+b) % L₀ = (c+d) % L₀
  have habL₀ : (a + b) % L₀ = (c + d) % L₀ := by
    rw [Nat.add_mod a b L₀, Nat.add_mod c d L₀]
    exact h0
  -- Step 2: reflection component → (a+b) % L₁ = (c+d) % L₁
  have habL₁ : (a + b) % L₁ = (c + d) % L₁ := by
    have haS : a ≤ S := hS a ha
    have hbS : b ≤ S := hS b hb
    have hcS : c ≤ S := hS c hc
    have hdS : d ≤ S := hS d hd
    rw [← Nat.add_mod (S-a) (S-b) L₁, ← Nat.add_mod (S-c) (S-d) L₁] at h1
    rw [show (S - a) + (S - b) = 2 * S - (a + b) from by omega,
        show (S - c) + (S - d) = 2 * S - (c + d) from by omega] at h1
    -- h1: (2*S - (a+b)) % L₁ = (2*S - (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 h1
    have h_dvd : (L₁ : ℤ) ∣ ((a + b : ℤ) - (c + d : ℤ)) := by
      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 (Int.modEq_zero_iff_dvd.mp h_sub)
    exact dvd_mod_eq (a + b) (c + d) L₁ h_dvd
  -- Step 3: By CRT, a+b = c+d (since both < L₀*L₁)
  have heq : a + b = c + d :=
    mod_eq_of_coprime hCoprime hL₀ hL₁ habL₀ habL₁ (hBound a ha b hb) (hBound c hc d hd)
  -- Step 4: 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.CRTSidon