"""full_chiral_dag.py — CRT Torus Braid DAG with Coprimality Guard.

Combines axis-swap (topology, YB ✓) and adjustment (resource, FA-changing)
into a unified DAG traversal. The Coprimality Guard ensures all 16 moduli
remain pairwise coprime after every crossing.

References:
  - docs/crt-torus-embedding.md (core CRT Torus embedding)
  - docs/research/unified_crt_torus_dag.md (graded Sidon energy + hierarchy)
  - docs/research/braid_group_action.md (dual-model framework)
"""

import math
from typing import List, Optional, Tuple

# ═══════════════════════════════════════════════════════════════
# §1  SIDON CHECK VIA WRAPPING CRITERION
# ═══════════════════════════════════════════════════════════════

# The Sidon creation theorem (docs/research/sidon_preservation_creation.md):
# For A0 ⊆ ℤ and CRT embedding F, F(A0) is Sidon iff for EVERY sum
# collision a+b = c+d in A0, the CRT lifts wrap M differently:
#   F(a)+F(b) = T + r₁·M,  F(c)+F(d) = T + r₂·M,  r₁ ≠ r₂
# where wrap indicator r = 1 if F(x)+F(y) ≥ M, else 0.
#
# With L_id-only adjustment and M > 2·max(A0), new collisions cannot
# form (M-difference condition is vacuous). The only question is
# whether the existing collisions break.

def crt_sum(a: int, b: int, pairs: List[Tuple[int, int]], S: int) -> Tuple[int, int]:
    """Compute F(a)+F(b) and its wrap indicator.

    Returns (sum, wrap) where wrap = 1 if sum ≥ M, 0 otherwise.
    """
    Fa = crt_embed(a, pairs, S)
    Fb = crt_embed(b, pairs, S)
    total = Fa + Fb
    M = math.prod(m for pair in pairs for m in pair)
    return (total, 1 if total >= M else 0)


def find_collisions(A0: List[int]) -> List[Tuple[int, int, int, int]]:
    """Find all sum collisions in A0.

    Returns list of ((a,b), (c,d), T) where a+b = c+d = T and (a,b) ≠ (c,d).
    """
    n = len(A0)
    sum_map = {}
    collisions = []
    for i in range(n):
        for j in range(i, n):
            s = A0[i] + A0[j]
            if s in sum_map:
                ci, cj = sum_map[s]
                if ci != i or cj != j:
                    collisions.append((A0[ci], A0[cj], A0[i], A0[j], s))
            else:
                sum_map[s] = (i, j)
    return collisions


def sidon_check(
    A0: List[int],
    pairs: List[Tuple[int, int]],
    S: int,
) -> Tuple[bool, int, float]:
    """Check if F(A0) is Sidon under current moduli.

    Returns (is_sidon, broken_count, score).
    - is_sidon: True if all collisions broken
    - broken_count: how many collisions are broken
    - score: 0 if Sidon, else graded residual (lower = closer to Sidon)
    """
    collisions = find_collisions(A0)
    if not collisions:
        return (True, 0, 0.0)

    broken = 0
    for a, b, c, d, T in collisions:
        _, wrap1 = crt_sum(a, b, pairs, S)
        _, wrap2 = crt_sum(c, d, pairs, S)
        if wrap1 != wrap2:
            broken += 1

    if broken == len(collisions):
        return (True, broken, 0.0)

    # Score: fraction of unbroken collisions, scaled to (0, 4].
    total = len(collisions)
    score = 4.0 * (1.0 - broken / total)
    return (False, broken, max(0.0, score))


def sidon_energy(
    A0: List[int],
    pairs: List[Tuple[int, int]],
    S: int,
) -> float:
    """Graded Sidon energy: 0 if Sidon, else ℰ ∈ (0, 4] for non-Sidon."""
    _, _, score = sidon_check(A0, pairs, S)
    return score


# ═══════════════════════════════════════════════════════════════
# §2  CRT EMBEDDING
# ═══════════════════════════════════════════════════════════════

def crt_embed(
    a: int,
    pairs: List[Tuple[int, int]],
    S: int,
) -> int:
    """CRT Torus Embedding F: ℤ → ℤ/Mℤ.

    Axis 1: a ↦ a mod L₁ (identity)
    Axes 2…k: a ↦ S − a mod Lᵢ (reflection)

    Reconstructs via CRT to produce a unique integer lift in [0, M).
    """
    residues = []
    moduli = []
    for i, (L_id, L_ref) in enumerate(pairs):
        moduli.append(L_id)
        if i == 0:
            residues.append(a % L_id)
        else:
            residues.append((S - a) % L_id)
        moduli.append(L_ref)
        residues.append((S - a) % L_ref)

    # Iterative CRT
    x = residues[0]
    M = moduli[0]
    for i in range(1, len(moduli)):
        m_i = moduli[i]
        r_i = residues[i]
        # Find k such that x + k·M ≡ r_i (mod m_i)
        # k ≡ (r_i − x) · M⁻¹ (mod m_i)
        inv = pow(M, -1, m_i)
        k = ((r_i - x) * inv) % m_i
        x = x + k * M
        M = M * m_i
    return x


def crt_embed_set(
    A: List[int],
    pairs: List[Tuple[int, int]],
    S: int,
) -> List[int]:
    """Apply CRT Torus Embedding F to every element of A."""
    return sorted([crt_embed(a, pairs, S) for a in A])


# ═══════════════════════════════════════════════════════════════
# §3  COPRIMALITY GUARD
# ═══════════════════════════════════════════════════════════════

def pairwise_coprime(moduli: List[int]) -> bool:
    """Coprimality Guard: check all moduli are pairwise coprime.

    Returns True iff gcd(m_i, m_j) = 1 for all i ≠ j.
    This is the CRITICAL invariant: CRT requires pairwise coprime moduli
    to guarantee injectivity of the torus embedding F.

    Failure mode: adjusting a modulus by ±2 can make it share a factor
    with another modulus (e.g., one hits 7, another was already 14).
    The guard catches this before it corrupts the node.
    """
    n = len(moduli)
    for i in range(n):
        for j in range(i + 1, n):
            if math.gcd(moduli[i], moduli[j]) != 1:
                return False
    return True


# ═══════════════════════════════════════════════════════════════
# §3  CHIRAL PAIRS — INITIALIZATION
# ═══════════════════════════════════════════════════════════════

def _nth_prime(n: int) -> int:
    """Return the n-th prime (0-indexed), generating on the fly."""
    known = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53,
             59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113]
    while len(known) <= n:
        candidate = known[-1] + 2
        while any(candidate % p == 0 for p in known):
            candidate += 2
        known.append(candidate)
    return known[n]


def chiral_pairs(
    n_strands: int = 8,
    band_gap: int = 30,
    base_prime_offset: int = 0,
) -> List[Tuple[int, int]]:
    """Initialize chiral pairs with distinct primes.

    Each strand gets an (L_id, L_ref) pair where both are prime.
    All 2·n_strands moduli are pairwise coprime by construction.

    With L_id-only adjustment (L_ref fixed), the spacing between
    L_id and L_ref doesn't restrict capacity — only the Q16_16
    bound (32767) and L_id > 1 matter.

    Args:
        n_strands:   number of braid strands
        band_gap:    (unused with L_id-only adjustment, kept for API compat)
        base_prime_offset: starting index into prime sequence

    Returns:
        List of (L_id, L_ref) pairs, one per strand
    """
    pairs = []
    idx = base_prime_offset
    for _ in range(n_strands):
        L_id = _nth_prime(idx); idx += 1
        L_ref = _nth_prime(idx); idx += 1
        pairs.append((L_id, L_ref))
    return pairs


# ═══════════════════════════════════════════════════════════════
# §4  AXIS-SWAP (TOPOLOGY, YB-COMPLIANT)
# ═══════════════════════════════════════════════════════════════

def axis_swap(
    pairs: List[Tuple[int, int]],
    s: int,
) -> List[Tuple[int, int]]:
    """Swap reflection moduli of adjacent strands s and s+1.

    This is the braid generator σ_s acting on the reflection axis only.
    The identity moduli are untouched. This is FA-invariant (CRT symmetry)
    and satisfies YB, σ²=id, and far commutativity.

    Args:
        pairs: current list of (L_id, L_ref) per strand
        s:     strand index (0 ≤ s < len(pairs) − 1)

    Returns a NEW list with the reflection moduli swapped.
    """
    if s < 0 or s >= len(pairs) - 1:
        return pairs[:]
    new_pairs = list(pairs)
    L_id_s, L_ref_s = new_pairs[s]
    L_id_s1, L_ref_s1 = new_pairs[s + 1]
    new_pairs[s] = (L_id_s, L_ref_s1)
    new_pairs[s + 1] = (L_id_s1, L_ref_s)
    return new_pairs


# ═══════════════════════════════════════════════════════════════
# §5  ADJUSTMENT (RESOURCE, FA-CHANGING)
# ═══════════════════════════════════════════════════════════════

def adjust(
    pairs: List[Tuple[int, int]],
    s: int,
    direction: str,
) -> Optional[List[Tuple[int, int]]]:
    """Adjust modulus values for strand s — L_id only.

    Design finding from Coprimality Guard (full_chiral_dag.py §3):
    The original ±2/∓1 adjustment on BOTH moduli breaks within-pair
    coprimality after ≤1 crossing (e.g., (13,41)→(15,40) shares
    factor 5). Fix: adjust only L_id, keeping L_ref fixed at its
    initial prime. This guarantees within-pair coprimality since
    gcd(L_id ± 2k, L_ref) = 1 when L_ref is a distinct prime and
    doesn't divide the adjusted L_id.

    Over:   L_id += 2
    Under:  L_id −= 2

    Returns a new list of pairs, or None if:
    - New modulus ≤ 1 (invalid for CRT)
    - Fails the Coprimality Guard (shares factor with another modulus)
    """
    if direction not in ('over', 'under'):
        raise ValueError(f"Invalid direction: {direction}")

    new_pairs = [(a, b) for a, b in pairs]
    L_id, L_ref = new_pairs[s]

    if direction == 'over':
        new_id = L_id + 2
    else:
        new_id = L_id - 2

    if new_id <= 1:
        return None  # modulus invalid

    new_pairs[s] = (new_id, L_ref)

    moduli = [m for pair in new_pairs for m in pair]
    if not pairwise_coprime(moduli):
        return None

    return new_pairs


# ═══════════════════════════════════════════════════════════════
# §6  COUPLED CROSSING
# ═══════════════════════════════════════════════════════════════

def coupled_crossing(
    pairs: List[Tuple[int, int]],
    s: int,
    direction: str,
    A0: List[int],
    S: int,
) -> Optional[Tuple[List[Tuple[int, int]], float]]:
    """One coupled crossing: axis-swap → adjust → verify.

    Returns (new_pairs, new_energy) if successful, None if coprimality fails.
    """
    swapped = axis_swap(pairs, s)
    adjusted = adjust(swapped, s, direction)
    if adjusted is None:
        return None
    E_new = sidon_energy(A0, adjusted, S)
    if not pairwise_coprime([m for pair in adjusted for m in pair]):
        return None
    return (adjusted, E_new)


# ═══════════════════════════════════════════════════════════════
# §7  DIRECTIONAL CAPACITY
# ═══════════════════════════════════════════════════════════════

def compute_directional_capacities(
    pairs: List[Tuple[int, int]],
    max_across: int = 15,
) -> List[int]:
    """Compute directional capacities per strand, packed into 4-bit word.

    L_id-only adjustment: 2 directions (over/under).
    Bits 0-1: cap_over (over-crossings, limited by Q16_16 bound 32767)
    Bits 2-3: cap_under (under-crossings, limited by L_id > 1)

    Each capacity capped at 3 (2-bit range).

    Args:
        pairs:      current chiral pairs
        max_across: maximum crossings used for normalization

    Returns:
        Packed capacities per strand, as list of ints
    """
    Q16_BOUND = 32767
    capacities = []
    for L_id, L_ref in pairs:
        cap_over = min(3, (Q16_BOUND - L_id) // 2)
        cap_under = min(3, (L_id - 3) // 2) if L_id > 3 else 0

        packed = cap_over | (cap_under << 2)
        capacities.append(packed)
    return capacities


def dag_capacity(capacities: List[int], direction: str) -> int:
    """DAG-level capacity: min of strand capacities in this direction."""
    shift = 0 if direction == 'over' else 2
    vals = [(c >> shift) & 3 for c in capacities]
    return min(vals)


# ═══════════════════════════════════════════════════════════════
# §8  DAG NODE
# ═══════════════════════════════════════════════════════════════

class DAGNode:
    """A node in the CRT torus braid DAG.

    Attributes:
        pairs: chiral pairs (L_id, L_ref) per strand
        A:     current integer set (CRT lifts)
        M:     product of all moduli
        energy: SidonEnergy ℰ of this state
        capacities: 4-directional capacities (8-bit per strand)
        braid_word: list of (strand, direction) crossings from root
        depth: number of crossings from root
        children: child node references (by moduli hash)
    """
    __slots__ = (
        'pairs', 'A0', 'S', 'M', 'energy', 'capacities',
        'braid_word', 'depth', 'children', 'is_sidon', 'broken',
    )

    def __init__(
        self,
        pairs: List[Tuple[int, int]],
        A0: List[int],
        S: int,
        braid_word: Optional[List[Tuple[int, str]]] = None,
        depth: int = 0,
    ):
        self.pairs = pairs
        self.A0 = A0
        self.S = S
        self.M = math.prod(m for pair in pairs for m in pair)
        sidon_ok, self.broken, self.energy = sidon_check(A0, pairs, S)
        self.is_sidon = sidon_ok
        self.capacities = compute_directional_capacities(pairs)
        self.braid_word = braid_word or []
        self.depth = depth
        self.children = []

    @property
    def moduli(self) -> List[int]:
        return [m for pair in self.pairs for m in pair]

    def modulus_hash(self) -> int:
        h = 0
        for m in self.moduli:
            h = h * 31 + m
        return h

    def __repr__(self) -> str:
        return (
            f"DAGNode(depth={self.depth}, M={self.M}, "
            f"ℰ={self.energy:.4f}, "
            f"braid={self.braid_word})"
        )


# ═══════════════════════════════════════════════════════════════
# §9  CHIRAL DAG TRAVERSAL
# ═══════════════════════════════════════════════════════════════

EPSILON = 1e-9


class ChiralDAG:
    """CRT Torus Braid DAG with unified axis-swap × adjustment traversal.

    Usage:
        dag = ChiralDAG(A0=[1, 2, 5, 6], S=7, n_strands=3)
        dag.build(max_steps=8)
        print(dag.summary())
    """

    def __init__(
        self,
        A0: List[int],
        S: int,
        n_strands: int = 8,
        band_gap: int = 60,
    ):
        self.A0 = sorted(A0)
        self.S = S
        self.n_strands = n_strands
        self.band_gap = band_gap
        self.root: Optional[DAGNode] = None
        self.visited: dict = {}
        self.stats = {
            'nodes_created': 0,
            'sidon_nodes': 0,
            'pruned_coprimality': 0,
            'pruned_energy': 0,
            'pruned_exhausted': 0,
            'deduped': 0,
            'phases': [0, 0, 0],
        }

    def _make_root(self) -> DAGNode:
        pairs = chiral_pairs(
            n_strands=self.n_strands,
            band_gap=self.band_gap,
            base_prime_offset=10,
        )
        node = DAGNode(pairs, self.A0, self.S, depth=0)
        self.visited[node.modulus_hash()] = node
        self.stats['nodes_created'] += 1
        return node

    def _maybe_prune(
        self,
        parent: DAGNode,
        child: DAGNode,
    ) -> bool:
        """Check if child should be pruned. Returns True if pruned."""
        # 1. Coprimality invariant: already checked in coupled_crossing,
        #    but re-check for safety.
        moduli = child.moduli
        if not pairwise_coprime(moduli):
            self.stats['pruned_coprimality'] += 1
            return True

        # 2. Monotonicity: SidonEnergy must not increase
        if child.energy > parent.energy + EPSILON:
            self.stats['pruned_energy'] += 1
            return True

        # 3. Sidon reached: accept but don't expand further
        if child.is_sidon:
            self.stats['sidon_nodes'] += 1
            return False  # accept, mark as terminal

        # 4. Capacity exhaustion: DAG-level check
        for d in ['over', 'under']:
            if dag_capacity(child.capacities, d) <= 0:
                self.stats['pruned_exhausted'] += 1
                return True

        return False

    def build(
        self,
        max_steps: int = 30,
        max_nodes: int = 10000,
        use_axis_swap: bool = True,
        use_adjustment: bool = True,
    ) -> None:
        """Build the DAG using BFS with three-phase traversal.

        Phase 1: Graded Sidon search (small bands)
        Phase 2: DAG topology expansion (wide bands)
        Phase 3: Content-addressable dedup (hash-based)
        """
        self.root = self._make_root()
        queue = [self.root]
        self.stats['phases'][0] += 1

        while queue and self.stats['nodes_created'] < max_nodes:
            node = queue.pop(0)

            if node.depth >= max_steps:
                continue

            # Phase transition: when energy is low, widen bands
            if node.energy < 0.5 and self.stats['phases'][1] == 0:
                self.stats['phases'][1] = 1
                self.band_gap = 500

            if node.is_sidon:
                continue  # terminal

            for s in range(self.n_strands):
                for direction in ['over', 'under']:
                    # DAG-level capacity check (fast prune)
                    if dag_capacity(node.capacities, direction) <= 0:
                        self.stats['pruned_exhausted'] += 1
                        continue

                    result = None
                    if use_axis_swap and use_adjustment:
                        result = coupled_crossing(
                            node.pairs, s, direction, self.A0, self.S,
                        )
                    elif use_axis_swap:
                        new_pairs = axis_swap(node.pairs, s)
                        new_E = sidon_energy(self.A0, new_pairs, self.S)
                        if pairwise_coprime([m for pair in new_pairs for m in pair]):
                            result = (new_pairs, new_E)
                    elif use_adjustment:
                        new_pairs = adjust(node.pairs, s, direction)
                        if new_pairs is not None:
                            new_E = sidon_energy(self.A0, new_pairs, self.S)
                            result = (new_pairs, new_E)
                        else:
                            self.stats['pruned_coprimality'] += 1
                    else:
                        continue

                    if result is None:
                        self.stats['pruned_coprimality'] += 1
                        continue

                    new_pairs, new_energy = result

                    child = DAGNode(
                        pairs=new_pairs,
                        A0=self.A0,
                        S=self.S,
                        braid_word=node.braid_word + [(s, direction)],
                        depth=node.depth + 1,
                    )

                    child.energy = new_energy

                    # Pruning gates
                    if self._maybe_prune(node, child):
                        continue

                    # Content-addressable dedup
                    h = child.modulus_hash()
                    if h in self.visited:
                        existing = self.visited[h]
                        if existing.energy <= child.energy:
                            self.stats['deduped'] += 1
                            node.children.append(existing)
                            continue

                    self.visited[h] = child
                    self.stats['nodes_created'] += 1
                    node.children.append(child)
                    queue.append(child)

        self.stats['phases'][2] = 1

    def summary(self) -> str:
        """Return a text summary of the DAG build."""
        sidon_nodes = [
            n for n in self.visited.values()
            if n.is_sidon
        ]
        if sidon_nodes:
            shortest = min(sidon_nodes, key=lambda n: n.depth)
            sidon_str = (
                f"Sidon paths found: {len(sidon_nodes)}\n"
                f"Shortest path: depth={shortest.depth}, "
                f"braid={shortest.braid_word}, "
                f"ℰ={shortest.energy:.4f}\n"
                f"Final moduli: {shortest.moduli}"
            )
        else:
            sidon_str = "No Sidon paths found."

        return (
            f"── ChiralDAG Summary ──\n"
            f"Strands: {self.n_strands}, Band gap: {self.band_gap}\n"
            f"Nodes created: {self.stats['nodes_created']}\n"
            f"Deduped: {self.stats['deduped']}\n"
            f"Pruned — coprimality: {self.stats['pruned_coprimality']}\n"
            f"Pruned — energy: {self.stats['pruned_energy']}\n"
            f"Pruned — exhausted: {self.stats['pruned_exhausted']}\n"
            f"Phases: {self.stats['phases']}\n"
            f"Sidon nodes: {self.stats['sidon_nodes']}\n"
            f"{sidon_str}"
        )

    def to_json(self, path: str) -> None:
        """Export DAG to JSON for visualization."""
        import json
        def _node_to_dict(n: DAGNode) -> dict:
            return {
                'depth': n.depth,
                'pairs': n.pairs,
                'moduli': n.moduli,
                'M': n.M,
                'energy': round(n.energy, 6),
                'is_sidon': n.is_sidon,
                'braid_word': n.braid_word,
                'children': [
                    c.modulus_hash() for c in n.children
                ],
            }
        data = {
            'n_strands': self.n_strands,
            'band_gap': self.band_gap,
            'A0': self.A0,
            'stats': self.stats,
            'nodes': {str(h): _node_to_dict(n)
                      for h, n in self.visited.items()},
        }
        with open(path, 'w') as f:
            json.dump(data, f, indent=2)


# ═══════════════════════════════════════════════════════════════
# §10  MAIN / SELF-TEST
# ═══════════════════════════════════════════════════════════════

if __name__ == '__main__':
    # Working test case (collision 3+13=8+8=16 is breakable via CRT wrapping)
    A0 = [0, 1, 3, 8, 13]
    S = 27

    print("=== 2-strand test ===")
    dag = ChiralDAG(A0=A0, S=S, n_strands=2)
    dag.build(max_steps=15, max_nodes=500)
    print(dag.summary())
    print()

    print("=== 8-strand test ===")
    dag8 = ChiralDAG(A0=A0, S=S, n_strands=8)
    dag8.build(max_steps=15, max_nodes=5000)
    print(dag8.summary())
    print()

    # Q16_16 bound check
    all_mods = [m for n in dag8.visited.values() for m in n.moduli]
    max_m = max(all_mods) if all_mods else 0
    print(f"Max modulus (8-strand): {max_m} {'✓' if max_m < 32767 else '✗ > 32767!'}")

    # Depth distribution of Sidon paths
    sidon_nodes = [n for n in dag8.visited.values() if n.is_sidon]
    if sidon_nodes:
        depths = {}
        for n in sidon_nodes:
            depths[n.depth] = depths.get(n.depth, 0) + 1
        print(f"Sidon depth distribution: {dict(sorted(depths.items()))}")