#!/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.
Uses PyTorch CUDA for p ≥ 9 (Ryser permanent on GPU).

Usage:
    python3 bmcte_extension_v1.py                    # full sweep
    python3 bmcte_extension_v1.py --dry-run           # show DAG only
    python3 bmcte_extension_v1.py --p-max 8           # CPU only
    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 bitmask subset enumeration. CPU (NumPy)."""
    p = M.shape[0]
    if p == 0:
        return 1.0 + 0.0j
    if p == 1:
        return complex(M[0, 0])

    total = 0.0 + 0.0j
    for mask in range(1, 1 << p):
        sign = (-1) ** (p - bin(mask).count("1"))
        row_sums = np.zeros(p, dtype=np.complex128)
        for j in range(p):
            if mask & (1 << j):
                row_sums += M[:, j]
        total += sign * np.prod(row_sums)
    return total


def ryser_permanent_gpu(M: torch.Tensor) -> torch.Tensor:
    """Ryser permanent via bitmask enumeration. GPU (PyTorch CUDA)."""
    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
    total = torch.tensor(0.0, dtype=torch.complex64, device=device)

    # Enumerate all non-empty subsets as bitmasks
    for mask in range(1, 1 << p):
        sign = (-1) ** (p - bin(mask).count("1"))
        # Build column mask
        cols = [j for j in range(p) if mask & (1 << j)]
        row_sums = M[:, cols].sum(dim=1)  # (p,)
        total = total + sign * row_sums.prod()

    return total


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

def categorical_sample(probs: np.ndarray, rng: np.random.RandomState) -> int:
    """Sample from categorical distribution."""
    cdf = np.cumsum(probs)
    r = rng.random() * cdf[-1]
    return int(np.searchsorted(cdf, r))


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."""

    rng = np.random.RandomState(seed)
    start = time.time()

    # Build random unitary (Haar-ish via QR)
    A = rng.randn(N, N) + 1j * rng.randn(N, N)
    Q, R = np.linalg.qr(A)
    U = Q @ np.diag(np.exp(1j * np.angle(np.diag(R))))

    # 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 _ in range(K):
        # Sample output modes
        S = [categorical_sample(col_probs[j], rng) for j in range(p)]

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

        # Compute permanent
        if use_gpu and HAS_CUDA:
            M_gpu = torch.tensor(M, dtype=torch.complex64, device="cuda")
            perm = ryser_permanent_gpu(M_gpu).cpu().item()
        else:
            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
    eps = 1e-15
    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 > 1e-15))

    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": use_gpu,
    }


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

def build_grid(
    N_values: list[int],
    p_values: list[int],
    K_values: list[int],
    n_seeds: int,
    p_gpu_threshold: int = 9,
) -> list[dict]:
    """Build all (N, p, K, seed) combinations."""
    grid = []
    for N in N_values:
        for p in p_values:
            if p >= N:
                continue  # skip impossible configs
            for K in K_values:
                for seed in range(n_seeds):
                    use_gpu = HAS_CUDA and p >= p_gpu_threshold
                    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("--p-gpu-threshold", type=int, default=9,
                        help="Photon count at which to use GPU")
    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, args.p_gpu_threshold)
    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 = []
    for (N, p), count in sorted(np_counts.items()):
        node_id = f"run_N{N}_p{p}"
        run_ids.append(node_id)
        use_gpu = HAS_CUDA and p >= args.p_gpu_threshold
        device = "GPU" if use_gpu else "CPU"
        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"GPU threshold: p >= {args.p_gpu_threshold}")
    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,
            "p_gpu_threshold": args.p_gpu_threshold,
        },
        "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())