#!/usr/bin/env python3
"""
bmcte_extension_v1.py — Extend λ(p)/H(p) invariance to higher photon numbers.

Sweeps N × p × K × seeds grid, measures entropy, variance, runtime.
GPU-accelerated path (PyTorch CUDA) for unitary construction, sampling,
and Ryser permanent.  Falls back to CPU when CUDA unavailable.

Usage:
    python3 bmcte_extension_v1.py                    # full sweep (GPU)
    python3 bmcte_extension_v1.py --dry-run           # show DAG only
    python3 bmcte_extension_v1.py --cpu               # force CPU
    python3 bmcte_extension_v1.py --N 2000 --p 2 4 6  # single N
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import sys
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import numpy as np

# PyTorch for GPU Ryser
try:
    import torch
    HAS_TORCH = True
    HAS_CUDA = torch.cuda.is_available()
except ImportError:
    HAS_TORCH = False
    HAS_CUDA = False

RECEIPT_PATH = Path(__file__).resolve().parent / "extension_v1_receipt.json"

# DB connection (optional)
try:
    sys.path.insert(0, str(Path("/home/allaun/Research Stack/4-Infrastructure/shim")))
    import psycopg2
    import psycopg2.extras
    from rds_connect import connect_rds
    HAS_DB = True
except ImportError:
    HAS_DB = False
    connect_rds = None


# =========================================================================
# DB Ingestion
# =========================================================================

def ingest_results(results: list[dict], experiment_id: str = "bmcte_extension_v1") -> int:
    """Write experiment results to bmcte.experiment_runs table.

    Returns number of rows inserted.
    """
    if not HAS_DB:
        print("  DB: psycopg2 not available, skipping ingestion")
        return 0

    conn = connect_rds()
    cur = conn.cursor()

    # Ensure schema/table exist
    cur.execute("""
        CREATE SCHEMA IF NOT EXISTS bmcte;
        CREATE TABLE IF NOT EXISTS bmcte.experiment_runs (
            id              SERIAL PRIMARY KEY,
            experiment_id   TEXT NOT NULL,
            run_hash        TEXT NOT NULL UNIQUE,
            N               INTEGER NOT NULL,
            p               INTEGER NOT NULL,
            K               INTEGER NOT NULL,
            seed            INTEGER NOT NULL,
            lambda_theory   DOUBLE PRECISION,
            entropy_measured DOUBLE PRECISION,
            nonzero_modes   INTEGER,
            weight_mean     DOUBLE PRECISION,
            weight_variance DOUBLE PRECISION,
            weight_cv       DOUBLE PRECISION,
            ryser_overflow  BOOLEAN DEFAULT FALSE,
            used_shots      INTEGER,
            runtime_ms      DOUBLE PRECISION,
            used_gpu        BOOLEAN DEFAULT FALSE,
            hardware        JSONB,
            created_at      TIMESTAMPTZ DEFAULT NOW(),
            receipt_hash    TEXT
        );
    """)
    conn.commit()

    inserted = 0
    for r in results:
        if "error" in r:
            continue

        run_hash = hashlib.sha256(
            f"{r['N']}:{r['p']}:{r['K']}:{r['seed']}".encode()
        ).hexdigest()[:16]

        try:
            cur.execute("""
                INSERT INTO bmcte.experiment_runs
                    (experiment_id, run_hash, N, p, K, seed,
                     lambda_theory, entropy_measured, nonzero_modes,
                     weight_mean, weight_variance, weight_cv,
                     ryser_overflow, used_shots, runtime_ms, used_gpu, hardware)
                VALUES (%(experiment_id)s, %(run_hash)s, %(N)s, %(p)s, %(K)s, %(seed)s,
                        %(lambda_theory)s, %(entropy_measured)s, %(nonzero_modes)s,
                        %(weight_mean)s, %(weight_variance)s, %(weight_cv)s,
                        %(ryser_overflow)s, %(used_shots)s, %(runtime_ms)s, %(used_gpu)s,
                        %(hardware)s::jsonb)
                ON CONFLICT (run_hash) DO NOTHING
            """, {
                "experiment_id": experiment_id,
                "run_hash": run_hash,
                "N": r["N"], "p": r["p"], "K": r["K"], "seed": r["seed"],
                "lambda_theory": r.get("lambda_theory"),
                "entropy_measured": r.get("entropy_measured"),
                "nonzero_modes": r.get("nonzero_modes"),
                "weight_mean": r.get("weight_mean"),
                "weight_variance": r.get("weight_variance"),
                "weight_cv": r.get("weight_cv"),
                "ryser_overflow": r.get("ryser_overflow", False),
                "used_shots": r.get("used_shots"),
                "runtime_ms": r.get("runtime_ms"),
                "used_gpu": r.get("used_gpu", False),
                "hardware": json.dumps({
                    "cpu": os.cpu_count(),
                    "gpu": torch.cuda.get_device_name(0) if HAS_CUDA else None,
                }),
            })
            inserted += 1
        except Exception as exc:
            print(f"  DB insert error: {exc}")

    conn.commit()
    cur.close()
    conn.close()
    return inserted


# =========================================================================
# DAG Definition
# =========================================================================

@dataclass
class DAGNode:
    """One node in the experiment DAG."""
    id: str
    label: str
    inputs: list[str] = field(default_factory=list)
    outputs: list[str] = field(default_factory=list)
    status: str = "pending"  # pending | running | done | skipped
    duration_ms: float = 0.0
    metadata: dict = field(default_factory=dict)


class ExperimentDAG:
    """Directed acyclic graph for experiment inspection."""

    def __init__(self):
        self.nodes: dict[str, DAGNode] = {}
        self.edges: list[tuple[str, str]] = []

    def add(self, node_id: str, label: str, inputs: list[str] = None,
            outputs: list[str] = None, **meta) -> DAGNode:
        node = DAGNode(
            id=node_id, label=label,
            inputs=inputs or [], outputs=outputs or [],
            metadata=meta,
        )
        self.nodes[node_id] = node
        for inp in (inputs or []):
            self.edges.append((inp, node_id))
        return node

    def mark(self, node_id: str, status: str, duration_ms: float = 0.0):
        if node_id in self.nodes:
            self.nodes[node_id].status = status
            self.nodes[node_id].duration_ms = duration_ms

    def to_mermaid(self) -> str:
        """Render DAG as Mermaid flowchart."""
        lines = ["graph TD"]
        for node in self.nodes.values():
            status_icon = {
                "pending": "⏳", "running": "🔄", "done": "✅", "skipped": "⏭️"
            }.get(node.status, "❓")
            safe_label = node.label.replace('"', "'")
            lines.append(f'    {node.id}["{status_icon} {safe_label}"]')
        for src, dst in self.edges:
            lines.append(f"    {src} --> {dst}")
        return "\n".join(lines)

    def to_ascii(self) -> str:
        """Render DAG as ASCII tree."""
        lines = []
        # Find root nodes (no inputs)
        roots = [n for n in self.nodes.values() if not n.inputs]
        visited = set()

        def render(node: DAGNode, indent: int = 0):
            if node.id in visited:
                return
            visited.add(node.id)
            status_icon = {
                "pending": "[ ]", "running": "[>]", "done": "[x]", "skipped": "[-]"
            }.get(node.status, "[?]")
            prefix = "  " * indent
            dur = f" ({node.duration_ms:.0f}ms)" if node.duration_ms > 0 else ""
            lines.append(f"{prefix}{status_icon} {node.label}{dur}")
            # Find children
            children = [self.nodes[dst] for src, dst in self.edges if src == node.id]
            for child in children:
                render(child, indent + 1)

        for root in roots:
            render(root)
        return "\n".join(lines)

    def summary(self) -> dict:
        return {
            "total_nodes": len(self.nodes),
            "done": sum(1 for n in self.nodes.values() if n.status == "done"),
            "pending": sum(1 for n in self.nodes.values() if n.status == "pending"),
            "total_ms": sum(n.duration_ms for n in self.nodes.values()),
        }


# =========================================================================
# Core: Ryser Permanent
# =========================================================================

def ryser_permanent_cpu(M: np.ndarray) -> complex:
    """Ryser permanent via vectorized bitmask enumeration. CPU (NumPy).

    Uses broadcasting to evaluate all 2^p - 1 non-empty subsets in parallel
    as a single matrix multiply + product.
    """
    p = M.shape[0]
    if p == 0:
        return 1.0 + 0.0j
    if p == 1:
        return complex(M[0, 0])

    # Boolean inclusion matrix: (2^p - 1, p) — True when column j is in mask i
    masks = np.arange(1, 1 << p, dtype=np.uint64)
    cols = np.arange(p, dtype=np.uint64)
    bits = ((masks[:, None] >> cols) & 1).astype(np.bool_)

    # Popcount = number of columns selected per mask
    popcount = bits.sum(axis=1, dtype=np.int32)
    signs = (-1) ** (p - popcount)  # (n_masks,) real

    # Row sums per mask: (n_masks, p) = (n_masks, p) @ (p, p)
    row_sums = bits @ M.T

    # Product across output rows → scalar per mask
    with np.errstate(divide="ignore", invalid="ignore"):
        prods = np.prod(row_sums, axis=1)

    total = np.sum(signs * prods)
    return complex(total)


def ryser_permanent_gpu(M: torch.Tensor) -> torch.Tensor:
    """Ryser permanent via vectorized bitmask enumeration. GPU (PyTorch CUDA).

    Processes all 2^p - 1 non-empty subsets in parallel as tensor ops.
    """
    p = M.shape[0]
    if p == 0:
        return torch.tensor(1.0, dtype=torch.complex64, device=M.device)
    if p == 1:
        return M[0, 0]

    device = M.device
    n_masks = (1 << p) - 1

    # All bitmasks 1 .. 2^p-1 as a vector
    masks = torch.arange(1, 1 << p, device=device)  # (n_masks,)

    # Popcount per mask: (mask >> k) & 1 for all k
    bits = (masks.unsqueeze(1) & (1 << torch.arange(p, device=device)))  # (n_masks, p)
    popcount = bits.count_nonzero(dim=1)  # (n_masks,)
    signs = (-1.0) ** (p - popcount)  # (n_masks,)

    # Mask vectors: 1.0 where bit k is set in mask m
    mask_vec = bits.to(M.dtype)  # (n_masks, p)

    # Row sums = M @ mask_vec^T  → (p, n_masks)
    row_sums = M @ mask_vec.T  # (p, n_masks)

    # Product across rows → (n_masks,)
    prods = torch.prod(row_sums, dim=0)

    return (signs * prods).sum().to(torch.complex64)


# =========================================================================
# Core: Monte Carlo Estimator
# =========================================================================

def run_single(
    N: int, p: int, K: int, seed: int,
    use_gpu: bool = False,
) -> dict:
    """One Monte Carlo run: sample K configurations, compute permanents.

    When ``use_gpu=True`` and CUDA is available, the unitary construction,
    categorical sampling, submatrix gather, and Ryser permanent all run on
    GPU (PyTorch CUDA).  Falls back to CPU otherwise.
    """

    start = time.time()
    eps = 1e-15

    if use_gpu and HAS_CUDA:
        return _run_single_gpu(N, p, K, seed, start, eps)

    # ── CPU path ────────────────────────────────────────────────────────
    rng = np.random.RandomState(seed)

    # Build random N×p isometry (marginal of first p columns of Haar unitary).
    # QR on N×p tall-skinny Gaussian is O(N·p²), vs O(N³) for full N×N QR.
    A = rng.randn(N, p) + 1j * rng.randn(N, p)
    Q, R = np.linalg.qr(A)
    U = Q @ np.diag(np.exp(1j * np.angle(np.diag(R))))  # (N, p) orthogonal columns

    # Column probability distributions
    col_probs = [np.abs(U[:, j]) ** 2 for j in range(p)]

    # Monte Carlo sampling
    mode_counts = np.zeros(N, dtype=np.float64)
    weights = []
    ryser_overflow = False
    used_shots = 0

    for shot in range(K):
        # Sample output modes
        S = [np.searchsorted(np.cumsum(col_probs[j]), rng.random() * col_probs[j].sum())
             for j in range(p)]

        # Gather p×p submatrix
        M = U[np.array(S), :]

        # Compute permanent
        perm = ryser_permanent_cpu(M)

        w = abs(perm) ** 2
        if math.isnan(w) or math.isinf(w):
            ryser_overflow = True
            continue

        weights.append(w)
        for s in S:
            mode_counts[s] += w
        used_shots += 1

    elapsed = time.time() - start

    # Compute metrics
    total_prob = float(np.sum(mode_counts))
    if total_prob > 0:
        mode_probs = mode_counts / total_prob
    else:
        mode_probs = np.ones(N) / N

    # Shannon entropy
    p_vec = np.clip(mode_probs, eps, 1.0)
    entropy = float(-np.sum(p_vec * np.log2(p_vec)))

    # λ(p) theory
    lambda_theory = math.exp(-p * p / N)

    # Weight statistics
    weights = np.array(weights) if weights else np.array([0.0])
    weight_mean = float(np.mean(weights))
    weight_var = float(np.var(weights))
    weight_cv = float(np.std(weights) / max(np.mean(weights), eps))

    nonzero_modes = int(np.sum(mode_probs > eps))

    return {
        "N": N, "p": p, "K": K, "seed": seed,
        "lambda_theory": round(lambda_theory, 6),
        "entropy_measured": round(entropy, 6),
        "nonzero_modes": nonzero_modes,
        "weight_mean": round(weight_mean, 6),
        "weight_variance": round(weight_var, 6),
        "weight_cv": round(weight_cv, 6),
        "ryser_overflow": ryser_overflow,
        "used_shots": used_shots,
        "runtime_ms": round(elapsed * 1000, 1),
        "used_gpu": False,
    }


def _run_single_gpu(
    N: int, p: int, K: int, seed: int,
    start: float, eps: float,
) -> dict:
    """GPU-accelerated Monte Carlo: unitary, sampling, permanent on CUDA."""
    device = "cuda"
    torch.manual_seed(seed)
    torch.cuda.synchronize()

    # ── Build random N×p isometry on GPU (marginal of first p cols of Haar) ─
    A = torch.complex(
        torch.randn(N, p, dtype=torch.float32, device=device),
        torch.randn(N, p, dtype=torch.float32, device=device),
    )
    Q, R = torch.linalg.qr(A)
    phase = torch.exp(1j * torch.angle(torch.diag(R)))
    U = Q @ torch.diag(phase)  # (N, p) complex64, orthogonal columns

    # Precompute column CDFs for categorical sampling
    col_cdfs = []
    for j in range(p):
        probs = (U[:, j].abs() ** 2).to(torch.float64)
        col_cdfs.append(torch.cumsum(probs, dim=0))

    # ── Monte Carlo sampling ─────────────────────────────────────────
    mode_counts = torch.zeros(N, dtype=torch.float64, device=device)
    weights = []
    ryser_overflow = False
    used_shots = 0

    for _ in range(K):
        # Sample p output modes: per-column searchsorted (p small, ≤10)
        rows = torch.stack([
            torch.searchsorted(col_cdfs[j], torch.rand(1, device=device)).squeeze()
            for j in range(p)
        ])

        # Gather p×p submatrix: U[rows, :]
        M = U[rows, :]

        # Permanent on GPU
        perm = ryser_permanent_gpu(M)
        w = (perm.abs() ** 2).item()

        if math.isnan(w) or math.isinf(w):
            ryser_overflow = True
            continue

        weights.append(w)
        mode_counts.scatter_add_(0, rows, torch.full((p,), w, dtype=torch.float64, device=device))
        used_shots += 1

    torch.cuda.synchronize()
    elapsed = time.time() - start

    # ── Compute metrics ───────────────────────────────────────────────
    total_prob = float(mode_counts.sum().item())
    if total_prob > 0:
        mode_probs = mode_counts / total_prob
    else:
        mode_probs = torch.ones(N, device=device) / N

    p_vec = mode_probs.clamp(min=eps)
    entropy = float(-(p_vec * p_vec.log2()).sum().item())

    lambda_theory = math.exp(-p * p / N)

    weight_arr = torch.tensor(weights, dtype=torch.float64) if weights else torch.zeros(1)
    weight_mean = float(weight_arr.mean().item())
    weight_var = float(weight_arr.var().item())
    weight_cv = float((weight_arr.std().item() / max(weight_mean, eps)))

    nonzero_modes = int((mode_probs > eps).sum().item())

    return {
        "N": N, "p": p, "K": K, "seed": seed,
        "lambda_theory": round(lambda_theory, 6),
        "entropy_measured": round(entropy, 6),
        "nonzero_modes": nonzero_modes,
        "weight_mean": round(weight_mean, 6),
        "weight_variance": round(weight_var, 6),
        "weight_cv": round(weight_cv, 6),
        "ryser_overflow": ryser_overflow,
        "used_shots": used_shots,
        "runtime_ms": round(elapsed * 1000, 1),
        "used_gpu": True,
    }


# =========================================================================
# Experiment Grid
# =========================================================================

def build_grid(
    N_values: list[int],
    p_values: list[int],
    K_values: list[int],
    n_seeds: int,
    force_cpu: bool = False,
) -> list[dict]:
    """Build all (N, p, K, seed) combinations."""
    grid = []
    for N in N_values:
        for p in p_values:
            if p >= N:
                continue
            for K in K_values:
                for seed in range(n_seeds):
                    use_gpu = HAS_CUDA and not force_cpu
                    grid.append({
                        "N": N, "p": p, "K": K, "seed": seed,
                        "use_gpu": use_gpu,
                    })
    return grid


def analyze_results(results: list[dict]) -> dict:
    """Aggregate results by (N, p) pair."""
    from collections import defaultdict
    groups = defaultdict(list)
    for r in results:
        groups[(r["N"], r["p"])].append(r)

    analysis = []
    for (N, p), runs in sorted(groups.items()):
        entropies = [r["entropy_measured"] for r in runs]
        runtimes = [r["runtime_ms"] for r in runs]
        cvs = [r["weight_cv"] for r in runs]
        nonzero = [r["nonzero_modes"] for r in runs]
        overflows = sum(1 for r in runs if r["ryser_overflow"])

        analysis.append({
            "N": N, "p": p,
            "lambda_theory": runs[0]["lambda_theory"],
            "entropy_mean": round(float(np.mean(entropies)), 6),
            "entropy_std": round(float(np.std(entropies)), 6),
            "entropy_min": round(float(np.min(entropies)), 6),
            "entropy_max": round(float(np.max(entropies)), 6),
            "nonzero_modes_mean": round(float(np.mean(nonzero)), 1),
            "weight_cv_mean": round(float(np.mean(cvs)), 6),
            "runtime_ms_mean": round(float(np.mean(runtimes)), 1),
            "runtime_ms_std": round(float(np.std(runtimes)), 1),
            "overflow_count": overflows,
            "n_runs": len(runs),
        })

    return {"grid_points": len(analysis), "analysis": analysis}


# =========================================================================
# Main
# =========================================================================

def main() -> int:
    parser = argparse.ArgumentParser(description="BMCTE Extension Experiment v1")
    parser.add_argument("--N", type=int, nargs="*", default=[2000, 5000, 10000],
                        help="Mode counts")
    parser.add_argument("--p", type=int, nargs="*", default=[2, 3, 4, 5, 6, 7, 8, 10],
                        help="Photon numbers")
    parser.add_argument("--K", type=int, nargs="*", default=[1000, 10000],
                        help="Sample counts")
    parser.add_argument("--seeds", type=int, default=5,
                        help="Independent seeds per config")
    parser.add_argument("--cpu", action="store_true",
                        help="Force CPU (default: GPU when CUDA available)")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show DAG only, don't execute")
    parser.add_argument("--output", type=str, default=str(RECEIPT_PATH),
                        help="Output receipt path")
    parser.add_argument("--no-db", action="store_true",
                        help="Skip database ingestion")

    args = parser.parse_args()

    # Build DAG
    dag = ExperimentDAG()

    # Input node
    dag.add("input", "Input Grid", outputs=["grid"])
    dag.nodes["input"].metadata = {
        "N": args.N, "p": args.p, "K": args.K, "seeds": args.seeds,
        "cuda": HAS_CUDA, "torch": HAS_TORCH,
    }

    # Grid expansion
    grid = build_grid(args.N, args.p, args.K, args.seeds, force_cpu=args.cpu)
    dag.add("grid", f"Expand Grid ({len(grid)} runs)", inputs=["input"], outputs=["collect"])

    # Group by (N, p) for DAG visualization
    from collections import Counter
    np_counts = Counter((r["N"], r["p"]) for r in grid)
    run_ids = []
    device = "CPU" if args.cpu or not HAS_CUDA else "GPU"
    for (N, p), count in sorted(np_counts.items()):
        node_id = f"run_N{N}_p{p}"
        run_ids.append(node_id)
        dag.add(node_id, f"N={N}, p={p} ({count} runs, {device})",
                inputs=["grid"], outputs=["collect"])

    # Collect
    dag.add("collect", "Collect Results", inputs=run_ids, outputs=["analyze"])

    # Analyze
    dag.add("analyze", "Analyze (aggregate by N,p)", inputs=["collect"], outputs=["receipt"])

    # Receipt
    dag.add("receipt", "Write Receipt JSON", inputs=["analyze"], outputs=["db"])

    # DB ingestion
    db_label = "Ingest to PostgreSQL" if HAS_DB and not args.no_db else "DB (skipped)"
    dag.add("db", db_label, inputs=["receipt"])

    # Print DAG
    print("=" * 70)
    print("BMCTE Extension Experiment v1 — DAG")
    print("=" * 70)
    print()
    print(dag.to_ascii())
    print()
    print(f"Total runs: {len(grid)}")
    print(f"CUDA available: {HAS_CUDA}")
    if HAS_CUDA:
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"Device: {'CPU (--cpu)' if args.cpu else 'GPU'}")
    else:
        print("Device: CPU")
    print()

    if args.dry_run:
        print("Dry run — exiting without execution.")
        # Write DAG mermaid to file
        dag_path = Path(args.output).with_suffix(".dag.mmd")
        dag_path.write_text(dag.to_mermaid())
        print(f"Mermaid DAG: {dag_path}")
        return 0

    # Execute grid
    dag.mark("input", "done")
    dag.mark("grid", "done")

    results = []
    total_start = time.time()

    for i, run_cfg in enumerate(grid):
        N, p, K, seed = run_cfg["N"], run_cfg["p"], run_cfg["K"], run_cfg["seed"]
        use_gpu = run_cfg["use_gpu"]
        node_id = f"run_N{N}_p{p}"

        if i == 0 or (N, p) != (grid[i-1]["N"], grid[i-1]["p"]):
            dag.mark(node_id, "running")

        try:
            result = run_single(N, p, K, seed, use_gpu=use_gpu)
            results.append(result)

            if result["ryser_overflow"]:
                print(f"  [{i+1}/{len(grid)}] N={N:>5} p={p:>2} K={K:>5} seed={seed} "
                      f"OVERFLOW runtime={result['runtime_ms']:.0f}ms")
            else:
                print(f"  [{i+1}/{len(grid)}] N={N:>5} p={p:>2} K={K:>5} seed={seed} "
                      f"H={result['entropy_measured']:.3f} "
                      f"λ={result['lambda_theory']:.4f} "
                      f"runtime={result['runtime_ms']:.0f}ms")

        except Exception as exc:
            print(f"  [{i+1}/{len(grid)}] N={N:>5} p={p:>2} K={K:>5} seed={seed} "
                  f"ERROR: {exc}")
            results.append({
                "N": N, "p": p, "K": K, "seed": seed,
                "error": str(exc)[:200],
            })

        # Mark node done when all seeds for this (N,p) are complete
        if i + 1 >= len(grid) or (grid[i+1]["N"], grid[i+1]["p"]) != (N, p):
            dag.mark(node_id, "done")

    total_elapsed = time.time() - total_start
    dag.mark("collect", "done", total_elapsed * 1000)

    # Analyze
    t0 = time.time()
    analysis = analyze_results(results)
    dag.mark("analyze", "done", (time.time() - t0) * 1000)

    # Print analysis table
    print()
    print("=" * 70)
    print("RESULTS")
    print("=" * 70)
    print(f"{'N':>6} {'p':>3} {'λ_theory':>10} {'H_mean':>8} {'H_std':>7} "
          f"{'nonzero':>8} {'CV_mean':>8} {'runtime':>9} {'overflow':>8}")
    print("-" * 70)
    for row in analysis["analysis"]:
        print(f"{row['N']:>6} {row['p']:>3} {row['lambda_theory']:>10.4f} "
              f"{row['entropy_mean']:>8.3f} {row['entropy_std']:>7.4f} "
              f"{row['nonzero_modes_mean']:>8.1f} {row['weight_cv_mean']:>8.4f} "
              f"{row['runtime_ms_mean']:>8.0f}ms {row['overflow_count']:>8}")

    # Write receipt
    t0 = time.time()
    receipt = {
        "schema": "bmcte_extension_v1",
        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
        "claim_boundary": "bmcte-regime-extension-empirical",
        "parameters": {
            "N_values": args.N,
            "p_values": args.p,
            "K_values": args.K,
            "seeds": args.seeds,
            "force_cpu": args.cpu,
        },
        "hardware": {
            "cuda": HAS_CUDA,
            "gpu": torch.cuda.get_device_name(0) if HAS_CUDA else None,
            "cpu_threads": os.cpu_count(),
        },
        "results": results,
        "analysis": analysis,
        "dag": {
            "mermaid": dag.to_mermaid(),
            "ascii": dag.to_ascii(),
            "summary": dag.summary(),
        },
        "total_runtime_s": round(total_elapsed, 2),
    }
    canonical = json.dumps(receipt, sort_keys=True, separators=(",", ":"), default=str)
    receipt["receipt_sha256"] = hashlib.sha256(canonical.encode()).hexdigest()

    out_path = Path(args.output)
    out_path.write_text(json.dumps(receipt, indent=2, default=str))
    dag.mark("receipt", "done", (time.time() - t0) * 1000)

    # DB ingestion
    if not args.no_db:
        t0 = time.time()
        n_inserted = ingest_results(results)
        dag.mark("db", "done", (time.time() - t0) * 1000)
        print(f"  DB: {n_inserted} rows inserted into bmcte.experiment_runs")
    else:
        dag.mark("db", "skipped")
        print("  DB: skipped (--no-db)")

    print()
    print(f"Receipt: {out_path}")
    print(f"SHA256: {receipt['receipt_sha256']}")
    print(f"Total: {total_elapsed:.1f}s")

    # Final DAG
    print()
    print("Final DAG:")
    print(dag.to_ascii())

    return 0


if __name__ == "__main__":
    sys.exit(main())
