#!/usr/bin/env python3
"""
Iteration DAG for CRT Torus Embedding.

Traces paths through modulus space, searching for a sequence of
modulus choices that transforms A into a Sidon set.
"""
import sys, math, random, itertools
from typing import List, Tuple, Optional, Dict, Set

sys.path.insert(0, '/home/allaun/SilverSight/scripts')
from verify_wrapping import *

# ---------- DAG Node ----------

class DAGNode:
    __slots__ = ('step', 'A', 'moduli', 'S', 'M', 'sidon', 'parent', 'children',
                 'terminal', 'reason', 'id')
    _next_id = 0

    def __init__(self, A, moduli, S, parent=None, step=0):
        self.id = DAGNode._next_id; DAGNode._next_id += 1
        self.step = step
        self.A = sorted(A)
        self.moduli = list(moduli)
        self.S = S
        self.M = 1
        for Li in moduli: self.M *= Li
        self.sidon = is_sidon(self.A)
        self.parent = parent
        self.children = []
        self.terminal = False
        self.reason = ""

    def key(self):
        return (tuple(self.A), tuple(self.moduli), self.S)

    def __repr__(self):
        return f"Node#{self.id}(step={self.step}, |A|={len(self.A)}, M={self.M}, sidon={self.sidon})"

# ---------- Regeneration Rules ----------

class GeometricRule:
    """Geometric growth: L1' = alpha * L1, L2' = beta * L2, ensuring coprimality."""
    def __init__(self, alpha=2, beta=3):
        self.alpha = alpha
        self.beta = beta

    def next_moduli(self, current_moduli, maxA=None):
        # Ensure next moduli remain coprime by using distinct growth factors
        results = []
        L1, L2 = current_moduli
        for a in [1, 2, 3]:
            for b in [1, 2, 3]:
                if a == b: continue  # keep moduli distinct
                nL1 = a * L1 if a > 0 else L1
                nL2 = b * L2 if b > 0 else L2
                if math.gcd(nL1, nL2) == 1:
                    results.append([nL1, nL2])
        return results[:5]  # limit branching

class AdaptiveRule:
    """Try all coprime modulus pairs with M in (maxA, 2*maxA]."""
    def __init__(self, max_val=16):
        self.numbers = [n for n in range(2, max_val + 1)]

    def next_moduli(self, current_moduli, maxA):
        candidates = []
        for L1 in self.numbers:
            for L2 in self.numbers:
                if L1 == L2:
                    continue
                if math.gcd(L1, L2) != 1:
                    continue
                M = L1 * L2
                if maxA < M <= 2 * maxA:
                    candidates.append([L1, L2])
        return candidates

class ExhaustiveRule:
    """Try all coprime k-modulus tuples within a bound."""
    def __init__(self, max_val=16):
        self.numbers = [n for n in range(2, max_val + 1)]

    def next_moduli(self, current_moduli, maxA):
        candidates = []
        for k in range(2, 5):
            for combo in itertools.permutations(self.numbers, k):
                # Check pairwise coprimality
                ok = True
                for i in range(k):
                    for j in range(i+1, k):
                        if math.gcd(combo[i], combo[j]) != 1:
                            ok = False
                            break
                    if not ok: break
                if not ok: continue
                M = 1
                for n in combo: M *= n
                if maxA < M <= 2 * maxA:
                    candidates.append(list(combo))
        return candidates[:self.max_branch] if hasattr(self, 'max_branch') else candidates

# ---------- DAG Builder ----------

class IterationDAG:
    def __init__(self, A0, S, regen_rule, max_steps=5, max_branch=100):
        self.root = DAGNode(A0, [3, 4], S)  # default initial moduli
        self.regen_rule = regen_rule
        self.max_steps = max_steps
        self.max_branch = max_branch
        self.all_nodes: Dict[str, DAGNode] = {self.root.key(): self.root}
        self.sidon_paths: List[List[DAGNode]] = []
        self.stats = {"explored": 0, "sidon_found": 0, "terminal": 0}

    def apply_F(self, node, new_moduli):
        """Apply F with new moduli to node.A, return child node or None."""
        new_A = [f_k(a, node.S, new_moduli) for a in node.A]
        child = DAGNode(new_A, new_moduli, node.S, parent=node, step=node.step + 1)
        return child

    def should_terminate(self, node):
        """Check if a node is terminal."""
        if node.sidon:
            node.terminal = True
            node.reason = "Sidon (goal reached)"
            return True
        if node.M > 2 * max(node.A):
            node.terminal = True
            node.reason = f"Preservation regime (M={node.M} > 2*maxA)"
            return True
        if node.step >= self.max_steps:
            node.terminal = True
            node.reason = f"Max steps ({self.max_steps}) reached"
            return True
        return False

    def build(self):
        """BFS build of the DAG."""
        queue = [self.root]
        visited = set()

        while queue:
            node = queue.pop(0)

            if node.key() in visited:
                continue
            visited.add(node.key())

            self.stats["explored"] += 1

            if self.should_terminate(node):
                self.stats["terminal"] += 1
                if node.sidon:
                    # Trace path to root
                    path = []
                    n = node
                    while n:
                        path.append(n)
                        n = n.parent
                    path.reverse()
                    self.sidon_paths.append(path)
                    self.stats["sidon_found"] += 1
                continue

            maxA = max(node.A)
            candidates = self.regen_rule.next_moduli(node.moduli, maxA)

            # Limit branching
            if len(candidates) > self.max_branch:
                candidates = candidates[:self.max_branch]

            for new_moduli in candidates:
                child = self.apply_F(node, new_moduli)
                if child.key() not in self.all_nodes:
                    self.all_nodes[child.key()] = child
                    node.children.append(child)
                    queue.append(child)

        return self

    def print_path(self, path):
        """Pretty-print a path from root to Sidon."""
        for i, node in enumerate(path):
            sidon = "★ SIDON" if node.sidon else ""
            term = " ⚑" if node.terminal else ""
            print(f"  Step {i}: A={node.A} M={node.M} {sidon}{term}")
            if node.parent and i > 0:
                print(f"           moduli={node.moduli}")

    def to_dot(self, filename=None):
        """Export DAG as DOT graph for visualization."""
        lines = ["digraph IterationDAG {"]
        lines.append("  rankdir=TB;")
        lines.append("  node [shape=record];")
        for key, node in self.all_nodes.items():
            sidon_style = "style=filled, fillcolor=lightgreen" if node.sidon else ""
            term_style = "style=filled, fillcolor=lightyellow" if node.terminal else ""
            style = sidon_style or term_style or ""
            label = f"A={node.A}\\nM={node.M} step={node.step}"
            if node.sidon: label += " ★SIDON"
            if style:
                lines.append(f"  n{node.id} [{style}, label=\"{label}\"];")
            else:
                lines.append(f"  n{node.id} [label=\"{label}\"];")
        for key, node in self.all_nodes.items():
            for child in node.children:
                lines.append(f"  n{node.id} -> n{child.id} [label=\"{child.moduli}\"];")
        lines.append("}")
        dot = "\n".join(lines)
        if filename:
            with open(filename, 'w') as f:
                f.write(dot)
            print(f"  DOT written to {filename}")
        return dot

    def summary(self):
        """Print DAG statistics."""
        print(f"DAG Statistics:")
        print(f"  Nodes explored: {self.stats['explored']}")
        print(f"  Sidon paths found: {self.stats['sidon_found']}")
        print(f"  Terminal nodes: {self.stats['terminal']}")
        print(f"  Total nodes: {len(self.all_nodes)}")
        if self.sidon_paths:
            print(f"  Shortest path length: {len(min(self.sidon_paths, key=len))}")
            print(f"\n  Shortest path:")
            self.print_path(min(self.sidon_paths, key=len))

# ---------- Main ----------

def test_sidon_example():
    """Trace the known Sidon creation example."""
    print("=== Sidon Creation Example ===")
    A0, S0 = [1, 2, 5, 6], 7
    rule = AdaptiveRule()
    dag = IterationDAG(A0, S0, rule, max_steps=3)
    dag.build()
    dag.summary()

def test_evolution():
    """Trace evolution of a non-Sidon set through modulus choices."""
    print("\n=== Evolution of A={0,1,3,8,13} ===")
    A0, S0 = [0, 1, 3, 8, 13], 27
    rule = AdaptiveRule()
    dag = IterationDAG(A0, S0, rule, max_steps=3)
    dag.build()
    dag.summary()

def test_geometric_cascade():
    """Trace a deterministic geometric cascade."""
    print("\n=== Geometric Cascade ===")
    A0, S0 = [1, 2, 5, 6], 7
    rule = GeometricRule(alpha=2, beta=2)
    dag = IterationDAG(A0, S0, rule, max_steps=5)
    dag.build()
    dag.summary()

if __name__ == "__main__":
    test_sidon_example()
    test_evolution()
    test_geometric_cascade()
