#!/usr/bin/env python3
"""Stress test: push CRT Torus DAG until model collapse.

Measures:
  - Max depth reached before all paths die
  - What kills the last frontier (coprimality, energy, exhaustion, dedup)
  - Depth vs alive count profile
  - Sidon paths found before collapse
"""

import sys, math, json, time
sys.path.insert(0, '/home/allaun/SilverSight/python')
from full_chiral_dag import *

def run_stress_test(
    A0: list,
    S: int,
    n_strands: int = 8,
    max_steps: int = 100,
    max_nodes: int = 200000,
    base_prime_offset: int = 10,
) -> dict:
    """Run DAG until collapse or resource limit."""
    pairs = chiral_pairs(
        n_strands=n_strands,
        band_gap=0,
        base_prime_offset=base_prime_offset,
    )

    root = DAGNode(pairs, A0, S, depth=0)
    visited = {root.modulus_hash(): root}
    frontier = [root]

    stats = {
        'nodes_created': 1,
        'sidon_nodes': 1 if root.is_sidon else 0,
        'pruned_coprimality': 0,
        'pruned_energy': 0,
        'pruned_exhausted': 0,
        'deduped': 0,
        'alive_by_depth': {0: 1},
        'sidon_by_depth': {0: 1 if root.is_sidon else 0},
        'collapse_depth': None,
        'collapse_cause': None,
        'final_frontier_count': 0,
        'runtime_s': 0,
        'sidon_paths': 0,
        'max_modulus': max(m for pair in pairs for m in pair),
    }

    start = time.time()

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

        if node.depth >= max_steps:
            continue
        if node.is_sidon:
            stats['sidon_paths'] += 1
            continue

        any_alive = False
        for s in range(n_strands):
            for direction in ['over', 'under']:
                cap = dag_capacity(node.capacities, direction)
                if cap <= 0:
                    stats['pruned_exhausted'] += 1
                    continue

                result = coupled_crossing(
                    node.pairs, s, direction, A0, S
                )
                if result is None:
                    stats['pruned_coprimality'] += 1
                    continue

                new_pairs, new_energy = result

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

                # Energy monotonicity gate
                if child.energy > node.energy + 1e-9:
                    stats['pruned_energy'] += 1
                    continue

                # Capacity check
                for d in ['over', 'under']:
                    if dag_capacity(child.capacities, d) <= 0:
                        stats['pruned_exhausted'] += 1
                        any_alive = True
                        break
                else:
                    any_alive = True

                h = child.modulus_hash()
                if h in visited:
                    existing = visited[h]
                    if existing.energy <= child.energy:
                        stats['deduped'] += 1
                        continue

                visited[h] = child
                stats['nodes_created'] += 1
                if child.is_sidon:
                    stats['sidon_nodes'] += 1
                    stats['sidon_paths'] += 1
                    # Don't expand Sidon nodes further

                d = child.depth
                stats['alive_by_depth'][d] = stats['alive_by_depth'].get(d, 0) + 1
                if child.is_sidon:
                    stats['sidon_by_depth'][d] = stats['sidon_by_depth'].get(d, 0) + 1

                node.children.append(child)
                if not child.is_sidon:
                    frontier.append(child)

                if stats['nodes_created'] >= max_nodes:
                    break
            if stats['nodes_created'] >= max_nodes:
                break

        # Check for collapse at this node's depth
        if not any_alive and not node.is_sidon:
            stats['collapse_depth'] = node.depth
            stats['collapse_cause'] = 'no_children'
            stats['final_frontier_count'] = len(frontier)

        # Periodic reporting
        if stats['nodes_created'] % 10000 == 0:
            elapsed = time.time() - start
            alive = len(frontier)
            print(
                f"  [{elapsed:.0f}s] depth={node.depth} "
                f"nodes={stats['nodes_created']} "
                f"alive={alive} "
                f"sidon={stats['sidon_nodes']} "
                f"coprimality={stats['pruned_coprimality']} "
                f"energy={stats['pruned_energy']} "
                f"exhausted={stats['pruned_exhausted']} "
                f"deduped={stats['deduped']} "
                f"max_mod={max(m for n in visited.values() for m in n.moduli)}"
            )

    stats['runtime_s'] = round(time.time() - start, 2)
    stats['total_visited'] = len(visited)
    stats['max_frontier_depth'] = max(frontier, key=lambda n: n.depth).depth if frontier else stats['collapse_depth']

    if not frontier and stats['nodes_created'] < max_nodes:
        stats['collapse_cause'] = 'full_collapse'
        stats['final_frontier_count'] = 0

    # Final pruning breakdown
    total_prune = (stats['pruned_coprimality'] + stats['pruned_energy']
                   + stats['pruned_exhausted'] + stats['deduped'])
    stats['total_pruned'] = total_prune

    # Modulus range analysis
    all_mods = [m for n in visited.values() for m in n.moduli]
    stats['max_modulus'] = max(all_mods) if all_mods else 0
    stats['min_modulus'] = min(all_mods) if all_mods else 0

    return stats


def main():
    test_sets = [
        ([0, 1, 3, 8, 13], 27, "working_Sidon"),
        ([0, 1, 2, 3, 4, 5], 5, "consecutive"),
        ([0, 1, 4, 6, 9, 14, 16, 21], 21, "sparse"),
    ]

    for A0, S, label in test_sets:
        coll = find_collisions(A0)
        print(f"\n{'=' * 60}")
        print(f"  Test: {label}")
        print(f"  A0={A0}, S={S}, collisions={len(coll)}")
        if len(coll) <= 5:
            for a, b, c, d, T in coll:
                print(f"    {a}+{b} = {c}+{d} = {T}")

        for n_strands in [2, 3, 4, 6, 8, 16]:
            print(f"\n  >> {n_strands} strands <<")
            stats = run_stress_test(
                A0=A0, S=S, n_strands=n_strands,
                max_steps=80, max_nodes=100000,
            )

            print(
                f"  ├── {stats['nodes_created']:>5} nodes "
                f"│ {stats['sidon_paths']:>4} Sidon "
                f"│ coll={stats['collapse_depth']} "
                f"│ copr={stats['pruned_coprimality']:>6} "
                f"│ en={stats['pruned_energy']:>5} "
                f"│ dedup={stats['deduped']:>5} "
                f"│ max_m={stats['max_modulus']} "
                f"│ {stats['runtime_s']:.1f}s"
            )

            depths = sorted(stats['alive_by_depth'].keys())
            alive_list = [(d, stats['alive_by_depth'].get(d, 0),
                          stats['sidon_by_depth'].get(d, 0)) for d in depths]
            print(f"  └── depths {depths[0]}–{depths[-1]} "
                  f"frontier @{stats['max_frontier_depth']}: "
                  + ", ".join(f"d{d}={cnt}" for d, cnt, _ in alive_list[-5:]))

            sys.stdout.flush()


if __name__ == '__main__':
    main()
