#!/usr/bin/env python3
"""chiral_batch_pipeline.py — Six-stage resource-aware search engine.

Implements the full pipeline:
  BraidStorm (256 chiral configs) → TreeBraid (factorize) →
  AngrySphinx (budget) → MultisurfacePacker (spatial) →
  COUCH (geometric) → Sidon (algebraic) → ~4-8 meaningful configs

Pure Python, exact arithmetic (fractions.Fraction). No external deps.
GPU-accelerable: the 256-config Sidon check is embarrassingly parallel.

Usage:
  python3 chiral_batch_pipeline.py [--seed N] [--strands K]
"""

import sys, math, json, time, hashlib, random, argparse
from pathlib import Path
from fractions import Fraction
from itertools import product
from collections import Counter

REPO_ROOT = Path(__file__).resolve().parent.parent
ARTIFACTS_DIR = REPO_ROOT / ".openresearch" / "artifacts"
ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
OUTPUT_PATH = ARTIFACTS_DIR / "chiral_batch_pipeline.json"


# ── Stage 1: BraidStorm — Generate chiral configurations ──────────────

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

def generate_chiral_configs(k):
    """Generate all 2^k chiral configurations for k crossings.
    Each config is a tuple of {0,1} where 0=over, 1=under."""
    return list(product([0, 1], repeat=k))

def embed_chiral(labels, S, moduli, config):
    """CRT embed with chiral configuration.
    config[j] = 0: standard reflection (S-a) mod L_j
    config[j] = 1: flipped reflection (a-S) mod L_j
    """
    embedded = []
    for a in labels:
        row = [a % moduli[0]]  # identity (poloidal)
        for j in range(1, len(moduli)):
            if config[j-1] == 0:
                row.append((S - a) % moduli[j])
            else:
                row.append((a - S) % moduli[j])
        embedded.append(row)
    return embedded

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, m)
    return x % m if g == 1 else None

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 sidon_check(embedded, moduli):
    M = 1
    for m in moduli: M *= m
    n = len(embedded)
    total_pairs = n * (n + 1) // 2
    vals = [crt_reconstruct(row, moduli) for row in embedded]
    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())
    return {
        "total_pairs": total_pairs,
        "collisions": collisions,
        "is_sidon": collisions == 0,
        "sidon_score": round(1.0 - collisions / total_pairs, 6) if total_pairs > 0 else 1.0,
    }


# ── Stage 2: TreeBraid — Factorize via braid relations ─────────────────

def factorize_braid(k, crossing_pairs):
    """Factorize crossings into independent groups.
    crossing_pairs: list of (i, j) strand pairs that cross.
    Two crossings are independent if |i-j| >= 2 (braid relation σ_i σ_j = σ_j σ_i).
    """
    groups = []
    remaining = list(range(k))
    while remaining:
        group = [remaining[0]]
        for idx in remaining[1:]:
            si, sj = crossing_pairs[idx]
            independent = True
            for gidx in group:
                gi, gj = crossing_pairs[gidx]
                if abs(si - gi) < 2 or abs(si - gj) < 2 or abs(sj - gi) < 2 or abs(sj - gj) < 2:
                    independent = False
                    break
            if independent:
                group.append(idx)
        for g in group:
            remaining.remove(g)
        groups.append(group)
    return groups


# ── Stage 3: AngrySphinx — Resource allocation ────────────────────────

def angrysphinx_filter(configs, budget):
    """Filter configs by compute budget.
    Each config's cost = 2^(number of under-crossings).
    Configs exceeding budget are pruned."""
    results = []
    for config in configs:
        under_count = sum(config)
        cost = 1 << under_count  # 2^under_count
        if cost <= budget:
            results.append((config, cost))
    return results


# ── Stage 4: MultisurfacePacker — Spatial fit ─────────────────────────

def multisurface_pack(configs, max_surfaces):
    """Pack configs into available surfaces.
    Simple greedy: each config occupies one surface, up to max_surfaces."""
    if len(configs) <= max_surfaces:
        return [(c, cost) for c, cost in configs]
    # Sort by cost (cheapest first = most efficient packing)
    sorted_configs = sorted(configs, key=lambda x: x[1])
    return sorted_configs[:max_surfaces]


# ── Stage 5: COUCH — Geometric filter ──────────────────────────────────

def couch_filter(configs, sidon_threshold=49152):
    """COUCH gate: check if self-loop (contention) is below threshold.
    Uses the config's under-crossing count as a proxy for contention.
    0 under = ring dispatch (0.0 self-loop, always passes)
    High under = high contention (may fail)"""
    results = []
    for config, cost in configs:
        under_count = sum(config)
        # Approximate self-loop: more under-crossings = more contention
        # 0 under → 0.0 (ring), k/2 under → ~0.823 (SUBLEQ), all under → ~0.885 (AVX)
        self_loop = int(57942 * under_count / max(len(config), 1))
        couch_stable = self_loop < sidon_threshold
        if couch_stable:
            results.append((config, cost, self_loop))
    return results


# ── Stage 6: Sidon filter — Algebraic uniqueness ──────────────────────

def sidon_filter(configs, labels, S, moduli):
    """Check which configs have unique CRT sums (Sidon property)."""
    results = []
    for config, cost, self_loop in configs:
        embedded = embed_chiral(labels, S, moduli, config)
        sidon = sidon_check(embedded, moduli)
        results.append({
            "config": config,
            "cost": cost,
            "self_loop": self_loop,
            "is_sidon": sidon["is_sidon"],
            "sidon_score": sidon["sidon_score"],
            "collisions": sidon["collisions"],
        })
    return results


# ── Main pipeline ──────────────────────────────────────────────────────

def run_pipeline(seed=0, k=8, budget=128, max_surfaces=64):
    rng = random.Random(seed)
    t0 = time.time()

    # Setup: Sidon labels and CRT moduli
    labels = [1, 2, 4, 8, 16, 32, 64, 128]  # powers of 2 (Sidon)
    S = 128  # reflection point
    # Coprime moduli: L0=7 (identity), L1..L8 = small coprimes
    moduli = [7, 3, 5, 11, 13, 17, 19, 23, 29]
    assert pairwise_coprime(moduli), "Moduli not coprime"

    # Crossing pairs (which strands cross)
    crossing_pairs = [(i, i+1) for i in range(k)]

    # Stage 1: BraidStorm — generate all 2^k chiral configs
    stage1_start = time.time()
    all_configs = generate_chiral_configs(k)
    stage1_time = time.time() - stage1_start

    # Stage 2: TreeBraid — factorize
    stage2_start = time.time()
    groups = factorize_braid(k, crossing_pairs)
    stage2_time = time.time() - stage2_start

    # Stage 3: AngrySphinx — budget filter
    stage3_start = time.time()
    budgeted = angrysphinx_filter(all_configs, budget)
    stage3_time = time.time() - stage3_start

    # Stage 4: MultisurfacePacker — spatial fit
    stage4_start = time.time()
    packed = multisurface_pack(budgeted, max_surfaces)
    stage4_time = time.time() - stage4_start

    # Stage 5: COUCH — geometric filter
    stage5_start = time.time()
    couch_passed = couch_filter(packed)
    stage5_time = time.time() - stage5_start

    # Stage 6: Sidon — algebraic filter
    stage6_start = time.time()
    sidon_results = sidon_filter(couch_passed, labels, S, moduli)
    stage6_time = time.time() - stage6_start

    sidon_passed = [r for r in sidon_results if r["is_sidon"]]

    elapsed = time.time() - t0

    result = {
        "experiment": "chiral_batch_pipeline",
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "seed": seed,
        "config": {
            "k": k,
            "labels": labels,
            "S": S,
            "moduli": moduli,
            "crossing_pairs": crossing_pairs,
            "budget": budget,
            "max_surfaces": max_surfaces,
        },
        "pipeline": {
            "stage1_braidstorm": {
                "input": 1,
                "output": len(all_configs),
                "time_s": round(stage1_time, 4),
            },
            "stage2_treebraid": {
                "groups": len(groups),
                "group_sizes": [len(g) for g in groups],
                "time_s": round(stage2_time, 4),
            },
            "stage3_angrysphinx": {
                "input": len(all_configs),
                "output": len(budgeted),
                "budget": budget,
                "time_s": round(stage3_time, 4),
            },
            "stage4_multisurfacepacker": {
                "input": len(budgeted),
                "output": len(packed),
                "max_surfaces": max_surfaces,
                "time_s": round(stage4_time, 4),
            },
            "stage5_couch": {
                "input": len(packed),
                "output": len(couch_passed),
                "time_s": round(stage5_time, 4),
            },
            "stage6_sidon": {
                "input": len(couch_passed),
                "output_sidon": len(sidon_passed),
                "output_total": len(sidon_results),
                "time_s": round(stage6_time, 4),
            },
        },
        "summary": {
            "total_configs": len(all_configs),
            "final_meaningful": len(sidon_passed),
            "reduction_ratio": round(len(all_configs) / max(len(sidon_passed), 1), 2),
            "elapsed_s": round(elapsed, 4),
            "sidon_pass_rate": round(len(sidon_passed) / max(len(sidon_results), 1), 4),
        },
        "sidon_configs": [
            {"config": list(r["config"]), "score": r["sidon_score"], "self_loop": r["self_loop"]}
            for r in sidon_passed
        ],
    }

    content = json.dumps(result, indent=2, sort_keys=True, default=str)
    result["sha256"] = hashlib.sha256(content.encode()).hexdigest()

    # Print summary
    print(f"\n{'='*60}")
    print(f"  SIX-STAGE PIPELINE RESULTS (k={k})")
    print(f"{'='*60}")
    print(f"  Stage 1 (BraidStorm):     {1} → {len(all_configs)} configs")
    print(f"  Stage 2 (TreeBraid):      {len(groups)} groups, sizes={[len(g) for g in groups]}")
    print(f"  Stage 3 (AngrySphinx):    {len(all_configs)} → {len(budgeted)} (budget={budget})")
    print(f"  Stage 4 (Packer):         {len(budgeted)} → {len(packed)} (max={max_surfaces})")
    print(f"  Stage 5 (COUCH):          {len(packed)} → {len(couch_passed)}")
    print(f"  Stage 6 (Sidon):          {len(couch_passed)} → {len(sidon_passed)} Sidon")
    print(f"{'='*60}")
    print(f"  Total → Meaningful:       {len(all_configs)} → {len(sidon_passed)}")
    print(f"  Reduction:                {result['summary']['reduction_ratio']}×")
    print(f"  Sidon pass rate:          {result['summary']['sidon_pass_rate']:.1%}")
    print(f"  Elapsed:                  {elapsed:.2f}s")
    print(f"{'='*60}")

    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Chiral batch pipeline")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--strands", type=int, default=8, help="Number of crossings k (2^k configs)")
    parser.add_argument("--budget", type=int, default=128, help="AngrySphinx compute budget")
    parser.add_argument("--surfaces", type=int, default=64, help="Max surfaces for packer")
    args = parser.parse_args()

    print(f"Chiral Batch Pipeline: k={args.strands}, 2^{args.strands}={1<<args.strands} configs")

    result = run_pipeline(seed=args.seed, k=args.strands, budget=args.budget,
                          max_surfaces=args.surfaces)

    OUTPUT_PATH.write_text(json.dumps(result, indent=2, default=str))
    print(f"\nResults → {OUTPUT_PATH}")
