#!/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.

Accepts max_depth as a CLI argument (default: 50).
"""

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"

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

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

def pairwise_coprime(moduli):
    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):
    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):
    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):
    return [crt_reconstruct(row, moduli) for row in embedded]

def sidon_check(embedded, moduli):
    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):
    moduli = [identity_base]
    p = first_refinement = 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):
    if start_moduli is None:
        start_moduli = pick_stride_moduli(n_strands, identity_base=3, first_reflection=101, min_gap=10)
    
    embedded, M = embed(label_set, S, start_moduli)
    sidon = sidon_check(embedded, start_moduli)
    
    root = DAGNode("root", "initial", embedded, sidon, start_moduli, 0)
    
    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:
            n = len(parent.moduli)
            for i in range(1, n):
                for j in range(i + 1, n):
                    new_moduli = list(parent.moduli)
                    new_moduli[i], new_moduli[j] = new_moduli[j], new_moduli[i]
                    
                    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():
    max_depth = 50
    if len(sys.argv) > 1:
        max_depth = int(sys.argv[1])
    
    print("=" * 60)
    print(f"  CRT Torus Braid DAG — Capacity Envelope (max_depth={max_depth})")
    print("  Integer-only. No float. No eigenvalue products.")
    print("=" * 60)
    
    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]
    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 (max_depth={max_depth})")
    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}")
    
    # Summary text
    print(f"\n{'='*60}")
    print(f"  KEY METRICS (max_depth={max_depth})")
    print(f"{'='*60}")
    print(f"  {'Label set':20s} {'Strnd':5s} {'States':7s} {'Sidon':7s} {'Stay?':6s} {'Headroom':10s} {'Time':7s}")
    print(f"  {'-'*20} {'-'*5} {'-'*7} {'-'*7} {'-'*6} {'-'*10} {'-'*7}")
    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']:5d} {r['total_states']:7d} {r['sidon_states']:7d} {stays:6s} {cap:10s} {r['elapsed_s']:>5.1f}s")

if __name__ == "__main__":
    main()