#!/usr/bin/env python3
"""crt_capacity_envelope.py — CRT Torus Braid DAG capacity envelope.

Explores how many Sidon states are reachable under the CRT Torus Embedding
for different strand counts (8, 12, 16) and label sets.

Core mechanics (all integer arithmetic):
  CRT embedding: F(a)₁ = a mod L₁ (identity), F(a)ᵢ = S-a mod Lᵢ (reflection)
  Axis-swap: permutes reflection moduli, ±2 increments on identity only
  Coprimality Guard: pairwise gcd(Lᵢ, Lⱼ) == 1 enforced after every step
  Sidon check: wrapping criterion — F(a)+F(b) vs F(c)+F(d) mod M
  Capacity: log₂(M) - log₂(max_label) as bits of headroom
"""

import sys
import math
import json
import time
import hashlib
from pathlib import Path
from collections import Counter

REPO_ROOT = Path(__file__).resolve().parent.parent
ARTIFACTS_DIR = REPO_ROOT / ".openresearch" / "artifacts"
OUTPUT_PATH = ARTIFACTS_DIR / "crt_capacity_envelope.json"

# ── Coprimality ──────────────────────────────────────────────────────────────

def gcd(a, b):
    while b:
        a, b = b, a % b
    return a

def pairwise_coprime(moduli):
    """Check all moduli are pairwise coprime."""
    for i in range(len(moduli)):
        for j in range(i + 1, len(moduli)):
            if gcd(moduli[i], moduli[j]) != 1:
                return False
    return True

# ── CRT Torus Embedding ─────────────────────────────────────────────────────

def embed(labels, S, moduli):
    """CRT Torus Embedding: F(a)₁ = a mod L₁, F(a)ᵢ = S - a mod Lᵢ."""
    n = len(moduli)
    M = math.prod(moduli)
    embedded = []
    for a in labels:
        row = []
        for i in range(n):
            if i == 0:
                row.append(a % moduli[0])
            else:
                row.append((S - a) % moduli[i])
        embedded.append(row)
    return embedded, M

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

def modinv(a, m):
    g, x, _ = egcd(a, m)
    if g != 1:
        return None
    return x % m

def crt_reconstruct(residues, moduli):
    """CRT reconstruction: find x mod M such that x ≡ residues[i] (mod moduli[i])."""
    M = 1
    for m in moduli:
        M *= m
    x = 0
    for r, m in zip(residues, moduli):
        Mi = M // m
        inv = modinv(Mi % m, m)
        if inv is None:
            return None
        x = (x + r * Mi * inv) % M
    return x

def crt_values(embedded, moduli):
    """CRT-reconstruct each embedded vector: x(a) ≡ F(a)_k (mod moduli[k])."""
    return [crt_reconstruct(row, moduli) for row in embedded]

def sidon_check(embedded, moduli):
    """Sidon check via CRT reconstruction. Uses unordered pairs (i ≤ j).
    Perfect Sidon set has n(n+1)/2 distinct sums (one per unordered pair)."""
    M = math.prod(moduli)
    n = len(embedded)
    total_pairs = n * (n + 1) // 2
    vals = crt_values(embedded, moduli)
    sums = []
    for i in range(n):
        for j in range(i, n):
            sums.append((vals[i] + vals[j]) % M)
    counts = Counter(sums)
    collisions = sum(c - 1 for c in counts.values())
    score = 1.0 - collisions / total_pairs
    return {
        "total_pairs": total_pairs,
        "distinct_residues": len(counts),
        "collisions": collisions,
        "sidon_score": round(score, 6),
    }

# ── Prime-aware modulus selection ────────────────────────────────────────────

def pick_stride_moduli(n, identity_base=3, first_reflection=101, min_gap=10):
    """Pick n moduli: identity at identity_base, then reflection moduli starting
    at first_reflection with at least min_gap between them.
    The large gap between identity and first reflection gives room for L_id ±2
    steps to explore before colliding with a reflection modulus."""
    moduli = [identity_base]
    p = first_reflection
    while len(moduli) < n:
        if all(p % d != 0 for d in range(2, int(p ** 0.5) + 1)):
            if p - moduli[-1] >= min_gap:
                moduli.append(p)
        p += 1
    return moduli

# ── DAG Node ─────────────────────────────────────────────────────────────────

class DAGNode:
    def __init__(self, node_id, node_type, state, sidon_result, moduli, depth):
        self.node_id = node_id
        self.node_type = node_type
        self.state = state
        self.sidon_result = sidon_result
        self.moduli = moduli
        self.depth = depth
        self.children = []
        self.max_modulus = max(moduli) if moduli else 0
        M = math.prod(moduli) if moduli else 1
        self.capacity = math.log2(max(M, 1)) - math.log2(max(self.max_modulus, 1))
    
    def is_sidon(self):
        return self.sidon_result["collisions"] == 0

# ── DAG traversal ───────────────────────────────────────────────────────────

def explore_component(label_set, S, n_strands, max_depth, start_moduli=None):
    """Explore the DAG starting from the initial CRT embedding.
    
    Args:
        label_set: list of labels to embed
        S: sum parameter for reflection embedding
        n_strands: number of CRT moduli to use
        max_depth: maximum number of axis-swaps to explore
        start_moduli: optional initial moduli (auto-generated if None)
    """
    if start_moduli is None:
        start_moduli = pick_stride_moduli(n_strands, identity_base=3, first_reflection=101, min_gap=10)
    
    # Initial embedding
    embedded, M = embed(label_set, S, start_moduli)
    sidon = sidon_check(embedded, start_moduli)
    
    root = DAGNode("root", "initial", embedded, sidon, start_moduli, 0)
    
    # BFS traversal
    frontier = [root]
    visited_signatures = set()
    sidon_nodes = 0
    sidon_paths = 0
    max_modulus_seen = max(start_moduli)
    
    sig = tuple(sorted(start_moduli))
    visited_signatures.add(sig)
    if root.is_sidon():
        sidon_nodes += 1
        sidon_paths += 1
    
    depth = 0
    while frontier and depth < max_depth:
        new_frontier = []
        depth += 1
        for parent in frontier:
            # Generate axis-swaps: swap each pair of reflection moduli
            n = len(parent.moduli)
            for i in range(1, n):
                for j in range(i + 1, n):
                    # Axis-swap: permute moduli i and j
                    new_moduli = list(parent.moduli)
                    new_moduli[i], new_moduli[j] = new_moduli[j], new_moduli[i]
                    
                    # L_id-only adjustment: ±2 on identity modulus
                    for delta in [2, -2]:
                        adj_moduli = list(new_moduli)
                        adj_moduli[0] += delta
                        if adj_moduli[0] <= 0:
                            continue
                        if not pairwise_coprime(adj_moduli):
                            continue
                        
                        sig = tuple(sorted(adj_moduli))
                        if sig in visited_signatures:
                            continue
                        visited_signatures.add(sig)
                        
                        embedded_n, M_n = embed(label_set, S, adj_moduli)
                        sidon_n = sidon_check(embedded_n, adj_moduli)
                        
                        node = DAGNode(f"d{depth}_{i}_{j}_{delta}",
                                       "swap", embedded_n, sidon_n,
                                       adj_moduli, depth)
                        parent.children.append(node)
                        new_frontier.append(node)
                        
                        max_modulus_seen = max(max_modulus_seen, max(adj_moduli))
                        
                        if node.is_sidon():
                            sidon_nodes += 1
                            sidon_paths += 1
        
        frontier = new_frontier
        if not frontier:
            break
    
    return {
        "label_set": label_set,
        "S": S,
        "n_strands": n_strands,
        "max_depth": max_depth,
        "total_states": len(visited_signatures),
        "sidon_states": sidon_nodes,
        "sidon_paths": sidon_paths,
        "max_modulus": max_modulus_seen,
        "M": M,
        "capacity_headroom_bits": round(math.log2(M) - math.log2(max_modulus_seen), 2) if max_modulus_seen > 0 else 0,
    }

# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    print("=" * 60)
    print("  CRT Torus Braid DAG — Capacity Envelope")
    print("  Integer-only. No float. No eigenvalue products.")
    print("=" * 60)
    
    # Test configurations (true Sidon sets only for capacity envelope)
    label_sets = [
        ("sidon_pow2", [1, 2, 4, 8, 16], 32),
        ("sidon_singer5", [0, 1, 4, 14, 16], 30),
        ("sidon_pow6", [1, 2, 4, 8, 16, 32], 64),
        ("nonsidon_seq5", [0, 1, 2, 3, 4], 5),
        ("nonsidon_seq6", [0, 1, 2, 3, 4, 5], 5),
    ]
    
    strand_counts = [8, 12, 16]
    max_depth = 50
    results = []
    
    for desc, labels, S in label_sets:
        print(f"\n{'='*60}")
        print(f"  Label set: {desc} (n={len(labels)}, S={S})")
        print(f"{'='*60}")
        
        for n_strands in strand_counts:
            print(f"\n  --- {n_strands} strands ---")
            t0 = time.time()
            
            result = explore_component(labels, S, n_strands, max_depth)
            elapsed = time.time() - t0
            
            result["desc"] = desc
            result["elapsed_s"] = round(elapsed, 2)
            is_sidon_label = "sidon" in desc
            result["collapsed"] = result["total_states"] <= n_strands * 2
            result["stays_sidon"] = is_sidon_label and result["sidon_states"] == result["total_states"]
            
            results.append(result)
            
            sys_s = "✓" if result["stays_sidon"] else ("✗" if is_sidon_label else "n/a")
            print(f"    Total states:  {result['total_states']}")
            print(f"    Sidon states:  {result['sidon_states']}")
            print(f"    Sidon paths:   {result['sidon_paths']}")
            print(f"    Stays Sidon:   {sys_s}")
            print(f"    Max modulus:   {result['max_modulus']}")
            print(f"    Capacity:      {result['capacity_headroom_bits']} bits")
            print(f"    Collapsed:     {result['collapsed']}")
            print(f"    Elapsed:       {result['elapsed_s']}s")
    
    # Summary table
    print(f"\n{'='*60}")
    print(f"  Summary — Capacity Envelope")
    print(f"{'='*60}")
    print(f"  {'Label set':20s} {'Strands':8s} {'States':8s} {'Sidon':8s} {'Stays?':6s} {'Max mod':8s} {'Capacity':10s}")
    print(f"  {'-'*20} {'-'*8} {'-'*8} {'-'*8} {'-'*6} {'-'*8} {'-'*10}")
    for r in results:
        cap = f"{r['capacity_headroom_bits']} bits" if not r['collapsed'] else "COLLAPSED"
        stays = "✓" if r.get("stays_sidon") else ("✗" if "sidon" in r["desc"] else "-")
        print(f"  {r['desc']:20s} {r['n_strands']:8d} {r['total_states']:8d} "
              f"{r['sidon_states']:8d} {stays:6s} {r['max_modulus']:8d} {cap:10s}")
    
    # Save
    ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
    output = {
        "schema": "crt_capacity_envelope_v1",
        "claim_boundary": "crt-torus-braid-dag:capacity-envelope:integer-only",
        "config": {"label_sets": label_sets, "strand_counts": strand_counts, "max_depth": max_depth},
        "results": results,
    }
    with open(OUTPUT_PATH, "w") as f:
        json.dump(output, f, indent=2, default=str)
    print(f"\n  Saved to {OUTPUT_PATH}")
    print("=" * 60)

if __name__ == "__main__":
    main()