#!/usr/bin/env python3
"""
rrc_photonic_stress_test.py — Convert the chaos-game graph centrality kernel
into a Perceval SLOS photonic circuit, then scale until the tensor network
breaks. Records the absolute limit (memory, NaN, zero-distribution, timeout).

Core idea:
  The Burgers representation graph adjacency A is the Hamiltonian for a
  continuous-time quantum walk  U(t) = e^{-iAt}.  We encode this as a
  linear optical interferometer: each graph node = a photonic mode, each
  edge = a beam-splitter coupling.  The output distribution over modes
  after propagation gives the node centrality (mode occupation probability
  ∝ eigenvector centrality[Quizz Blog 2008]).

  We then scale the graph size (number of modes and/or photons) until
  SLOS returns useless output: all-zero distribution, NaN/inf statevector,
  memory exhaustion, or catastrophic runtime blowup.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

import numpy as np

# ── Perceval — core photonic simulator ──────────────────────────────
try:
    import perceval as pcvl
    from perceval.backends import SLOSBackend
    _HAS_PERVERSE = True
except ImportError:
    _HAS_PERVERSE = False


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


# =========================================================================
# I. Graph → Photonic Interferometer Mapping
# =========================================================================

def adjacency_to_interferometer(
    adj: np.ndarray,
    coupling_phase: float = math.pi / 4,
) -> pcvl.Circuit:
    """Map an N×N adjacency matrix to a linear optical interferometer.

    Builds the unitary U = exp(-i*A*t) at time t = coupling_phase and
    embeds it as a generic N-mode photonic circuit via pcvl.Unitary.

    Single-photon input evolves under the adjacency Hamiltonian — the
    output mode distribution gives the eigenvector centrality (modes
    with higher occupation = nodes with higher centrality).

    Args:
        adj: N×N adjacency matrix (symmetric, zero-diagonal).
        coupling_phase: Evolution time (θ in U = exp(-iAθ)).

    Returns:
        pcvl.Circuit implementing the interferometer.
    """
    N = adj.shape[0]
    H = adj.astype(np.complex128)
    eigenvalues, eigenvectors = np.linalg.eigh(H)
    U = eigenvectors @ np.diag(np.exp(-1j * eigenvalues * coupling_phase)) @ eigenvectors.conj().T

    circuit = pcvl.Circuit(N)
    circuit.add(0, pcvl.Unitary(U))
    return circuit


# =========================================================================
# II. Build a scalable synthetic graph
# =========================================================================

def make_random_graph(N: int, p: float = 0.4, seed: int = 42) -> np.ndarray:
    """Synthetic Erdős–Rényi graph for scaling tests."""
    rng = np.random.RandomState(seed)
    adj = (rng.random((N, N)) < p).astype(np.float64)
    adj = np.triu(adj, 1) + np.triu(adj, 1).T
    return adj


def make_representation_graph() -> np.ndarray:
    """Build the 22-representation graph from burgers_chaos_game.py."""
    SHIM_PATH = Path(__file__).resolve().parent
    sys.path.insert(0, str(SHIM_PATH))
    from burgers_chaos_game import build_representation_graph
    labels, adj = build_representation_graph()
    return adj, labels


# =========================================================================
# III. Run photonic simulation and extract centrality
# =========================================================================

def photonic_centrality(
    adj: np.ndarray,
    n_photons: int = 1,
    coupling_phase: float = math.pi / 4,
    shots: int = 10000,
    timeout_s: float = 60.0,
) -> dict:
    """Run the graph as a photonic quantum walk and extract mode centralities.

    Args:
        adj: N×N adjacency matrix.
        n_photons: Number of indistinguishable photons.
        coupling_phase: Walk evolution time parameter.
        shots: Number of measurement samples.
        timeout_s: Max wall-clock seconds.

    Returns:
        Dict with centralities, output distribution, and diagnostics.
    """
    N = adj.shape[0]
    start = time.time()

    # Build interferometer
    circuit = adjacency_to_interferometer(adj, coupling_phase)
    circuit_t = round(time.time() - start, 4)

    if time.time() - start > timeout_s:
        return {"error": f"timeout after circuit build ({circuit_t}s)", "n": N, "n_photons": n_photons}

    # Input state: first n_photons modes each get 1 photon
    input_list = [1] * n_photons + [0] * (N - n_photons)
    input_state = pcvl.BasicState(input_list)

    # Processor (SLOS — exact tensor network)
    try:
        processor = pcvl.Processor("SLOS")
        processor.set_circuit(circuit)
        processor.with_input(input_state)
    except Exception as exc:
        return {"error": f"Processor init failed: {exc}", "n": N, "n_photons": n_photons}

    proc_t = round(time.time() - start, 4)
    if time.time() - start > timeout_s:
        return {"error": f"timeout after processor init ({proc_t}s)", "n": N, "n_photons": n_photons}

    # Sample
    try:
        sampler = pcvl.algorithm.Sampler(processor)
        res = sampler.sample_count(shots)
    except MemoryError:
        return {"error": "MemoryError (OOM)", "n": N, "n_photons": n_photons}
    except Exception as exc:
        msg = str(exc)
        if "memory" in msg.lower() or "allocation" in msg.lower():
            return {"error": f"Memory exhaustion: {msg[:200]}", "n": N, "n_photons": n_photons}
        return {"error": f"SLOS simulation failed: {msg[:200]}", "n": N, "n_photons": n_photons}

    sample_t = round(time.time() - start, 4)
    if time.time() - start > timeout_s:
        return {"error": f"timeout after sampling ({sample_t}s)", "n": N, "n_photons": n_photons}

    # Extract results
    results = res["results"]
    total_counts = sum(results.values())
    if total_counts == 0:
        return {"error": "Zero-distribution: all counts are zero", "n": N, "n_photons": n_photons}

    # Per-mode occupation probabilities
    mode_probs = [0.0] * N
    for state, count in results.items():
        prob = count / max(total_counts, 1)
        photons_by_mode = list(state)
        for m, pcount in enumerate(photons_by_mode):
            if m < N:
                mode_probs[m] += prob * pcount

    # Centrality from occupation: sum over states weighted by photon count
    total_prob = sum(mode_probs)
    centrality = [p / max(total_prob, 1e-15) for p in mode_probs]

    # Diagnostics
    nonzero_states = sum(1 for c in results.values() if c > 0)
    entropy = 0.0
    for c in results.values():
        p = c / max(total_counts, 1)
        if p > 1e-15:
            entropy -= p * math.log2(p)

    # Check for NaN/Inf
    has_nan = any(math.isnan(c) or math.isinf(c) for c in centrality)

    return {
        "n": N,
        "n_photons": n_photons,
        "centrality": [round(c, 6) for c in centrality],
        "mode_occupations": [round(p, 6) for p in mode_probs],
        "output_entropy": round(entropy, 6),
        "nonzero_output_states": nonzero_states,
        "total_samples": total_counts,
        "has_nan": has_nan,
        "circuit_build_ms": round(circuit_t * 1000, 1),
        "processor_ms": round((proc_t - circuit_t) * 1000, 1),
        "sampling_ms": round((sample_t - proc_t) * 1000, 1),
        "total_ms": round((time.time() - start) * 1000, 1),
        "edge_count": int(np.sum(adj > 0) // 2),
        "edges_successful": True,
    }


# =========================================================================
# IV. Stress-test driver: scale sizes and report breakpoint
# =========================================================================

def stress_test(
    sizes: list[int],
    n_photons_list: list[int],
    shots: int = 10000,
    timeout_s: float = 60.0,
    use_real_graph: bool = True,
    coupling_phase: float = math.pi / 4,
) -> list[dict]:
    """Iterate over sizes, record when SLOS breaks.

    Args:
        sizes: Number of modes (graph nodes) to test.
        n_photons_list: Photon counts to test per size.
        shots: Samples per run.
        timeout_s: Wall-clock timeout per run.
        use_real_graph: If True, use the 22-representation graph for
                        compatible sizes; else synthetic random.
        coupling_phase: Walk time parameter.

    Returns:
        List of result dicts, one per (size, n_photons).
    """
    results: list[dict] = []

    real_adj, real_labels = None, None
    if use_real_graph:
        real_adj, real_labels = make_representation_graph()

    for N in sizes:
        for n_photons in n_photons_list:
            if n_photons > N:
                continue  # can't have more photons than modes (Fock state)

            # Build adjacency
            if use_real_graph and real_adj is not None and N <= real_adj.shape[0]:
                adj = real_adj[:N, :N]
            else:
                adj = make_random_graph(N)

            print(f"  N={N:>4d}  photons={n_photons:>2d}  edges={np.sum(adj>0)//2:>4.0f}", end="")

            t0 = time.time()
            result = photonic_centrality(
                adj, n_photons=n_photons, coupling_phase=coupling_phase,
                shots=shots, timeout_s=timeout_s,
            )
            elapsed = time.time() - t0

            result["size_label"] = f"N={N}_p={n_photons}"
            result["n_modes"] = N
            result["coupling_phase"] = coupling_phase
            result["shots"] = shots
            result["elapsed_s"] = round(elapsed, 2)

            error = result.get("error")
            if error:
                ent_str = "FAIL"
                status = "FAIL"
            else:
                entropy = result.get("output_entropy", 0)
                ent_str = f"H={entropy:.3f}"
                status = "OK"

            print(f"  {status:>4s}  {ent_str:>10s}  {elapsed:>6.1f}s")
            results.append(result)

            if elapsed > timeout_s:
                print(f"    → TIMEOUT at N={N}, photons={n_photons}")
                break

        else:
            continue
        break  # outer break if inner timed out

    return results


# =========================================================================
# V. Hilbert dimension estimation
# =========================================================================

def hilbert_dim(N: int, n_photons: int) -> int:
    """Hilbert space dimension for N modes and n_photons indist. photons."""
    return math.comb(N + n_photons - 1, n_photons)


def estimate_max_size(max_photons: int = 6, max_dim: int = 10_000_000) -> None:
    """Estimate largest (N, n_photons) within a given Hilbert dimension."""
    print(f"\nEstimated max size within dim ≤ {max_dim:,}:")
    for n_photons in range(1, max_photons + 1):
        for N in range(2, 500):
            if hilbert_dim(N, n_photons) > max_dim:
                print(f"  {n_photons} photon(s): N ≤ {N - 1} (dim={hilbert_dim(N - 1, n_photons):,})")
                break


# =========================================================================
# VI. Main
# =========================================================================

def main() -> int:
    if not _HAS_PERVERSE:
        print("Perceval not installed. Install with: pip install perceval-quandela")
        return 1

    parser = argparse.ArgumentParser(
        description="RRC Photonic Stress Test — Push SLOS to its absolute limit"
    )
    parser.add_argument("--sizes", type=int, nargs="*",
                        default=[2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22,
                                 24, 26, 28, 30, 35, 40, 45, 50, 60, 70,
                                 80, 90, 100, 120, 140],
                        help="Number of modes to test")
    parser.add_argument("--photons", type=int, nargs="*", default=[1],
                        help="Photon counts to test per size")
    parser.add_argument("--shots", type=int, default=10000,
                        help="Samples per run")
    parser.add_argument("--timeout", type=float, default=60.0,
                        help="Timeout per run (seconds)")
    parser.add_argument("--synthetic", action="store_true",
                        help="Use synthetic random graphs instead of representation graph")
    parser.add_argument("--phase", type=float, default=math.pi / 4,
                        help="BS coupling phase (walk time)")
    parser.add_argument("--estimate-only", action="store_true",
                        help="Just estimate max sizes, don't run")
    parser.add_argument("--sweep-photons", action="store_true",
                        help="Also sweep 2,3,4 photons (expensive)")

    args = parser.parse_args()

    n_photons_list = list(args.photons)
    if args.sweep_photons:
        n_photons_list = sorted(set(n_photons_list + [1, 2, 3, 4]))

    if args.estimate_only:
        estimate_max_size()
        return 0

    print("=" * 70)
    print("RRC Photonic Stress Test — SLOS Absolute Limit")
    print("=" * 70)
    print(f"  Sizes: {args.sizes[0]}–{args.sizes[-1]} modes")
    print(f"  Photons: {n_photons_list}")
    print(f"  Shots: {args.shots}")
    print(f"  Timeout: {args.timeout}s")
    print(f"  Phase: {args.phase:.4f}")
    print(f"  Graph: {'synthetic random' if args.synthetic else 'representation graph'}")
    print()

    # Pre-compute Hilbert dimensions for context
    print("Hilbert space dimensions:")
    for n_photons in n_photons_list:
        for N in [args.sizes[0], 10, 20, 30, 50, 100]:
            if N <= args.sizes[-1]:
                print(f"  N={N:>3d}, {n_photons} photon(s): dim≈{hilbert_dim(N, n_photons):>12,}")
    print()

    results = stress_test(
        sizes=args.sizes,
        n_photons_list=n_photons_list,
        shots=args.shots,
        timeout_s=args.timeout,
        use_real_graph=not args.synthetic,
        coupling_phase=args.phase,
    )

    # Summarize breakpoint
    failures = [r for r in results if "error" in r]
    successes = [r for r in results if "error" not in r]
    max_ok = max([r["n"] for r in successes], default=0)
    max_ok_photons = max([r["n_photons"] for r in successes], default=0)
    first_fail = failures[0] if failures else None

    print()
    print("=" * 70)
    print("RESULT SUMMARY")
    print("=" * 70)
    print(f"  Total runs: {len(results)}")
    print(f"  Successful: {len(successes)}")
    print(f"  Failed:     {len(failures)}")

    if successes:
        best = max(successes, key=lambda r: (r["n"], r["n_photons"]))
        print(f"\n  Largest successful run:")
        print(f"    N={best['n']} modes, {best['n_photons']} photon(s)")
        print(f"    Output entropy: H={best.get('output_entropy', 'N/A')}")
        print(f"    Total time: {best.get('total_ms', 0):.0f} ms")
        if "centrality" in best:
            top_mode = max(range(len(best["centrality"])),
                          key=lambda i: best["centrality"][i])
            print(f"    Highest centrality: mode {top_mode} ({best['centrality'][top_mode]:.4f})")

    if first_fail:
        print(f"\n  First failure:")
        print(f"    N={first_fail['n']} modes, {first_fail['n_photons']} photon(s)")
        print(f"    Error: {first_fail['error']}")

    # Conservatism estimate: where would 1- and 2-photon regimes diverge?
    print("\n  Hilbert dimension at break boundary:")
    max_ok_dim = hilbert_dim(max_ok, max_ok_photons) if successes else 0
    print(f"    Max SUCCESS: N={max_ok}, p={max_ok_photons}, dim={max_ok_dim:,}")
    if first_fail:
        fail_dim = hilbert_dim(first_fail["n"], first_fail["n_photons"])
        print(f"    First FAIL:  N={first_fail['n']}, p={first_fail['n_photons']}, dim={fail_dim:,}")

    # Receipt
    receipt = dict(
        schema="rrc_photonic_stress_test_v1",
        claim_boundary="absolute-limit-measurement-of-perceval-slos-tensor-network;no-decision-logic",
        parameters=dict(
            sizes=args.sizes,
            n_photons_list=sorted(n_photons_list),
            shots=args.shots,
            timeout_s=args.timeout,
            phase=args.phase,
            synthetic_graph=args.synthetic,
        ),
        hilbert_dims={
            f"N={N}_p={np}": hilbert_dim(N, np)
            for np in n_photons_list
            for N in [args.sizes[0], 10, 20, 30, 50, 100]
            if N <= args.sizes[-1]
        },
        results=results,
        summary=dict(
            total_runs=len(results),
            successful=len(successes),
            failed=len(failures),
            max_successful_n=max_ok,
            max_successful_photons=max_ok_photons,
            first_failure=dict(
                n=first_fail["n"],
                n_photons=first_fail["n_photons"],
                error=first_fail["error"],
            ) if first_fail else None,
        ),
    )
    canonical = json.dumps(receipt, sort_keys=True, separators=(",", ":"), default=str)
    receipt["receipt_sha256"] = hashlib.sha256(canonical.encode()).hexdigest()
    receipt["computed_at"] = datetime.now(timezone.utc).isoformat()
    RECEIPT_PATH.write_text(json.dumps(receipt, indent=2, default=str))
    print(f"\nFull receipt: {RECEIPT_PATH}")
    print(f"SHA256: {receipt['receipt_sha256']}")

    return 0 if not failures else 1


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