#!/usr/bin/env python3
"""
Wrapping Criterion Verification for CRT Torus Embedding.

Tests the Sidon creation condition across random modulus choices
and set configurations for k = 2 and k >= 3.
"""
import math, random, itertools, hashlib, json
from typing import List, Tuple, Set

def egcd(a: int, b: int):
    if b == 0: return a, 1, 0
    g, x, y = egcd(b, a % b)
    return g, y, x - (a // b) * y

def modinv(a: int, m: int) -> int:
    g, x, _ = egcd(a % m, m)
    assert g == 1, f"{a} not invertible mod {m}"
    return x % m

def crt_lift(r1: int, r2: int, L1: int, L2: int) -> int:
    """CRT lift: find x in [0, L1*L2) with x ≡ r1 mod L1, x ≡ r2 mod L2."""
    t = ((r2 - r1) * modinv(L1, L2)) % L2
    return r1 + t * L1

def crt_lift_k(residues: List[int], moduli: List[int]) -> int:
    """CRT lift for k moduli via iterative Garner-like approach."""
    x = residues[0]
    m = moduli[0]
    for i in range(1, len(moduli)):
        t = ((residues[i] - x) * modinv(m, moduli[i])) % moduli[i]
        x += t * m
        m *= moduli[i]
    return x

def f_k(a: int, S: int, moduli: List[int]) -> int:
    """F(a) for k-modulus embedding: axis 1 = a mod L1, others = (S-a) mod Li."""
    residues = [a % moduli[0]] + [(S - a) % Li for Li in moduli[1:]]
    return crt_lift_k(residues, moduli)

def is_sidon(X: List[int]) -> bool:
    """Check Sidon property (all pairwise sums distinct)."""
    sums = set()
    for i in range(len(X)):
        for j in range(i, len(X)):
            s = X[i] + X[j]
            if s in sums: return False
            sums.add(s)
    return True

def sum_collisions(X: List[int]) -> List[Tuple[Tuple[int,int],Tuple[int,int]]]:
    """Return all sum collisions [(a,b),(c,d)] with a+b = c+d, ordered."""
    sum_map = {}
    collisions = []
    for i in range(len(X)):
        for j in range(i, len(X)):
            s = X[i] + X[j]
            if s in sum_map:
                for pair in sum_map[s]:
                    collisions.append((pair, (i, j)))
            sum_map.setdefault(s, []).append((i, j))
    return collisions

def wrapping_criterion(a, b, c, d, S, moduli):
    """Check if two colliding pairs wrap the modulus boundary differently."""
    M = 1
    for Li in moduli: M *= Li
    Fa_sum = f_k(a, S, moduli) + f_k(b, S, moduli)
    Fc_sum = f_k(c, S, moduli) + f_k(d, S, moduli)
    wrap_ab = Fa_sum >= M
    wrap_cd = Fc_sum >= M
    return wrap_ab != wrap_cd, Fa_sum, Fc_sum, M

def test_2_modulus():
    """Test the known Sidon example and random cases for k=2."""
    print("=== k=2 Tests ===")
    tests = [
        # (A, S, L1, L2, description)
        ([1,2,5,6], 7, 3, 4, "Sidon creation example"),
        ([1,2,5,6], 100, 3, 4, "S changed, same A"),
        ([1,3,5,7], 8, 3, 5, "Symmetric set, odd"),
        ([0,2,4,6], 6, 5, 7, "Even set"),
        ([1,4,6,9], 10, 7, 11, "Random set"),
        ([0,1,3,4], 4, 3, 5, "Small set"),
        ([2,5,7,10], 12, 5, 7, "Medium set"),
        ([0,3,5,8,10,13], 13, 5, 8, "6-element set"),
    ]
    for A, S, L1, L2, desc in tests:
        moduli = [L1, L2]
        M = L1 * L2
        A_sidon = is_sidon(A)
        FA = [f_k(a, S, moduli) for a in A]
        FA_sidon = is_sidon(FA)
        collisions = sum_collisions(A)
        wrapped = []
        for (i,j),(p,q) in collisions:
            a,b,c,d = A[i],A[j],A[p],A[q]
            diff, s1, s2, _ = wrapping_criterion(a, b, c, d, S, moduli)
            wrapped.append((a,b,c,d,s1,s2,diff))
        status = "OK" if FA_sidon else "FAIL"
        print(f"  {desc:30s} A_sidon={A_sidon} FA_sidon={FA_sidon} |A|={len(A)} M={M} coll={len(collisions)} wrap={len(wrapped)}")

def test_3_modulus():
    """Test with k=3 moduli."""
    print("\n=== k=3 Tests ===")
    tests = [
        ([1,2,5,6], 7, [3,4,5]),
        ([1,2,5,6], 7, [3,5,7]),
        ([0,1,3,4], 4, [3,5,7]),
        ([0,2,4,6,8,10], 10, [5,7,11]),
        ([1,4,6,9,11,14], 15, [7,11,13]),
    ]
    for A, S, moduli in tests:
        M = 1
        for Li in moduli: M *= Li
        FA = [f_k(a, S, moduli) for a in A]
        FA_sidon = is_sidon(FA)
        collisions = sum_collisions(A)
        wrapped = []
        for (i,j),(p,q) in collisions:
            a,b,c,d = A[i],A[j],A[p],A[q]
            diff, s1, s2, _ = wrapping_criterion(a, b, c, d, S, moduli)
            wrapped.append(diff)
        print(f"  moduli={moduli} |A|={len(A)} M={M} A_sidon={is_sidon(A)} FA_sidon={FA_sidon} coll={len(collisions)} wraps={wrapped.count(True)}")

def test_k_random():
    """Test with randomly generated parameters for various k."""
    print("\n=== Random k >= 2 tests ===")
    random.seed(42)
    primes = [2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53]
    for k in [2,3,4,6,8]:
        for trial in range(20):
            moduli = random.sample(primes, k)
            # Need all coprime — fine with distinct primes
            maxA = random.randint(5, 30)
            A = sorted(random.sample(range(0, maxA), min(maxA, random.randint(4, 8))))
            S = random.randint(maxA, 2*maxA)
            # Quick closure check: ensure A is S-closed (may not be — that's deliberate)
            M = 1
            for Li in moduli: M *= Li
            FA = [f_k(a, S, moduli) for a in A]
            FA_sidon = is_sidon(FA)
            A_sidon = is_sidon(A)
            collisions = sum_collisions(A)
            wrapped_count = 0
            for (i,j),(p,q) in collisions:
                a,b,c,d = A[i],A[j],A[p],A[q]
                diff, _, _, _ = wrapping_criterion(a, b, c, d, S, moduli)
                if diff: wrapped_count += 1
            if collisions or not FA_sidon:
                print(f"  k={k} |A|={len(A)} M={M} A_sidon={A_sidon} FA_sidon={FA_sidon} coll={len(collisions)} wraps={wrapped_count}")

def pairwise_sums(X):
    """Return the set of all pairwise sums of X."""
    sums = {}
    for i in range(len(X)):
        for j in range(i, len(X)):
            s = X[i] + X[j]
            sums.setdefault(s, []).append((i,j))
    return sums

def m_difference_condition(A, M):
    """
    Condition (b): no two distinct pairwise sums of A differ by exactly M.
    Returns (holds: bool, violators: list).
    """
    sums = pairwise_sums(A)
    sum_vals = list(sums.keys())
    violators = []
    for i in range(len(sum_vals)):
        for j in range(i+1, len(sum_vals)):
            if abs(sum_vals[i] - sum_vals[j]) == M:
                violators.append((sum_vals[i], sum_vals[j],
                                  sums[sum_vals[i]], sums[sum_vals[j]]))
    return len(violators) == 0, violators

def wrapping_condition(A, S, moduli):
    """
    Condition (a): for every sum collision in A, the pairs wrap M differently.
    Returns (holds: bool, unresolved: list).
    """
    M = 1
    for Li in moduli: M *= Li
    collisions = sum_collisions(A)
    unresolved = []
    for (i,j),(p,q) in collisions:
        a,b,c,d = A[i],A[j],A[p],A[q]
        diff, s1, s2, _ = wrapping_criterion(a,b,c,d,S,moduli)
        if not diff:
            unresolved.append(((a,b,c,d),(s1,s2)))
    return len(unresolved) == 0, unresolved

def certify_sidon_creation(A, S, moduli, verbose=False):
    """
    Certify whether F(A) is guaranteed Sidon.
    Returns (guaranteed: bool, FA: list, reason: str).
    """
    M = 1
    for Li in moduli: M *= Li
    FA = [f_k(a, S, moduli) for a in A]
    FA_sidon = is_sidon(FA)
    
    # Check injection regime
    if M <= max(A):
        return False, FA, f"Aliasing regime (M={M} <= max(A)={max(A)}), F not injective"
    
    # Check condition (a): wrapping
    wrap_ok, unresolved = wrapping_condition(A, S, moduli)
    
    # Check condition (b): M-difference
    mdiff_ok, violators = m_difference_condition(A, M)
    
    if wrap_ok and mdiff_ok:
        return True, FA, "Guaranteed Sidon (both conditions satisfied)"
    elif not wrap_ok:
        return False, FA, f"Wrapping criterion fails for {len(unresolved)} collision(s)"
    elif not mdiff_ok:
        return False, FA, f"M-difference condition fails ({len(violators)} violator(s))"
    else:
        return False, FA, "Unknown failure"

def verify_complete_theorem():
    """Verify the complete Sidon theorem (both conditions)."""
    print("\n=== Complete Theorem Verification ===")
    random.seed(456)
    primes = [2,3,5,7,11,13,17,19,23,29,31,37]
    passed = 0
    failed = 0
    for trial in range(2000):
        k = random.randint(2, 5)
        moduli = random.sample(primes, k)
        M = 1
        for Li in moduli: M *= Li
        n = random.randint(3, 10)
        maxA = random.randint(3, 20)
        A = sorted(random.sample(range(maxA+1), min(n, maxA+1)))
        S = random.randint(maxA, 2*maxA)
        
        # Only test in the injective regime (M > max(A))
        if M <= max(A):
            continue
        
        guaranteed, FA, reason = certify_sidon_creation(A, S, moduli)
        FA_sidon = is_sidon(FA)
        
        if guaranteed and FA_sidon:
            passed += 1
        elif not guaranteed and not FA_sidon:
            passed += 1
        else:
            print(f"  COUNTEREXAMPLE: guaranteed={guaranteed} FA_sidon={FA_sidon}")
            print(f"    k={k} moduli={moduli} M={M} A={A} S={S} FA={FA}")
            print(f"    reason={reason}")
            failed += 1
            if failed >= 5: break
    print(f"  Passed: {passed} / {passed+failed}")
    
    # Also certify the Sidon creation example
    print()
    A_ex = [1,2,5,6]
    S_ex = 7
    mod_ex = [3,4]
    g, FA, r = certify_sidon_creation(A_ex, S_ex, mod_ex, verbose=True)
    print(f"  Sidon example: guaranteed={g}, FA={FA}")
    print(f"    Reason: {r}")

if __name__ == "__main__":
    test_2_modulus()
    test_3_modulus()
    test_k_random()
    verify_complete_theorem()