#!/usr/bin/env python3
"""perceval_slos_verify.py — SLOS eigenvalue product verification with recoverable DAG.

5-minute time limit per shot on Quandela cloud. Each SLOS call:
- Checkpoints results to disk before and after
- Records the computation DAG (which steps, in what order)
- If a shot times out, resume from the last checkpoint
- The DAG itself is data — how SLOS computes is informative

Usage:
  # Local (needs perceval installed)
  python3 scripts/perceval_slos_verify.py
  
  # Quandela cloud (needs PERCEVAL_TOKEN)
  PERCEVAL_TOKEN='your_token' python3 scripts/perceval_slos_verify.py --cloud
  
  # Resume from checkpoint after timeout
  python3 scripts/perceval_slos_verify.py --resume
"""
import sys
import os
import math
import json
import time
import hashlib
from pathlib import Path
from typing import Optional
from collections import Counter

REPO_ROOT = Path(__file__).resolve().parent.parent
ARTIFACTS_DIR = REPO_ROOT / ".openresearch" / "artifacts"
CHECKPOINT_DIR = ARTIFACTS_DIR / "slos_checkpoints"
DAG_PATH = ARTIFACTS_DIR / "slos_computation_dag.json"

_findings = []
def finding(m, s, c, v, d=None):
    _findings.append({"module": m, "severity": s, "claim": c, "verdict": v, "details": d or {}})

# ── Recoverable DAG ───────────────────────────────────────────────────

class ComputationDAG:
    """Records each computation step as a DAG node.
    
    Each node records:
    - what was computed (labels, k, method)
    - when it started/ended (timestamp)
    - the result (SLOS output or eigenvalue products)
    - whether it succeeded or timed out
    
    The DAG is checkpointed to disk after each node, so if a shot
    times out (5-min limit), the DAG preserves what was computed
    and how."""
    
    def __init__(self):
        self.nodes = []
        self.edges = []
        CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
    
    def add_node(self, node_id, node_type, inputs, result, status, elapsed):
        """Add a computation node to the DAG."""
        node = {
            "id": node_id,
            "type": node_type,  # "eigenvalue_products" or "slos" or "compare"
            "inputs": inputs,    # labels, k, n_modes, etc.
            "result": result,    # the computation output
            "status": status,    # "success", "timeout", "error"
            "elapsed_s": round(elapsed, 2),
            "timestamp": time.time(),
        }
        self.nodes.append(node)
        self._checkpoint(node)
        return node
    
    def add_edge(self, from_id, to_id, edge_type="depends_on"):
        """Add an edge: to_id depends on from_id."""
        self.edges.append({"from": from_id, "to": to_id, "type": edge_type})
    
    def _checkpoint(self, node):
        """Save this node to its own checkpoint file."""
        path = CHECKPOINT_DIR / f"node_{node['id']}.json"
        path.write_text(json.dumps(node, default=str, indent=2))
        # Also save the full DAG state
        dag_state = {"nodes": self.nodes, "edges": self.edges}
        DAG_PATH.write_text(json.dumps(dag_state, default=str, indent=2))
    
    def load_checkpoint(self):
        """Load DAG state from disk (for --resume)."""
        if DAG_PATH.exists():
            state = json.loads(DAG_PATH.read_text())
            self.nodes = state.get("nodes", [])
            self.edges = state.get("edges", [])
            print(f"  Resumed from checkpoint: {len(self.nodes)} nodes, {len(self.edges)} edges")
            return True
        return False
    
    def has_node(self, node_id):
        """Check if a node already exists (skip recompute on resume)."""
        return any(n["id"] == node_id for n in self.nodes)
    
    def get_node(self, node_id):
        """Get a node's result (for dependent computations)."""
        for n in self.nodes:
            if n["id"] == node_id:
                return n
        return None
    
    def save_dag_report(self):
        """Save the full DAG as a report."""
        report_path = ARTIFACTS_DIR / "slos_computation_dag.md"
        lines = [
            "# SLOS Computation DAG\n",
            f"**Nodes:** {len(self.nodes)}",
            f"**Edges:** {len(self.edges)}\n",
            "## Nodes\n",
            "| ID | Type | Status | Elapsed | Inputs |",
            "|----|------|--------|---------|--------|",
        ]
        for n in self.nodes:
            inputs_str = json.dumps(n["inputs"], default=str)[:60]
            lines.append(f"| {n['id']} | {n['type']} | {n['status']} | {n['elapsed_s']}s | {inputs_str} |")
        
        lines.append("\n## Edges\n")
        lines.append("| From | To | Type |")
        lines.append("|------|----|------|")
        for e in self.edges:
            lines.append(f"| {e['from']} | {e['to']} | {e['type']} |")
        
        lines.append("\n## How SLOS Computes (the informative part)\n")
        lines.append("The DAG records the exact computation path:")
        for n in self.nodes:
            lines.append(f"\n### Node {n['id']}: {n['type']}")
            lines.append(f"- Status: {n['status']}")
            lines.append(f"- Elapsed: {n['elapsed_s']}s")
            lines.append(f"- Inputs: {json.dumps(n['inputs'], default=str)}")
            if n['status'] == 'success':
                result = n['result']
                if isinstance(result, dict):
                    for k, v in list(result.items())[:10]:
                        lines.append(f"- {k}: {v}")
        
        report_path.write_text("\n".join(lines))
        return report_path

# ── Sidon crossing matrix ─────────────────────────────────────────────

def build_sum_matrix(labels):
    n = len(labels)
    S = [[0.0]*n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            S[i][j] = float(labels[i] + labels[j])
    max_val = max(max(row) for row in S)
    if max_val > 0:
        for i in range(n):
            for j in range(n):
                S[i][j] /= max_val
    return S

def matrix_to_unitary(S):
    try:
        import numpy as np
    except ImportError:
        return None, None, None
    A = np.array(S, dtype=np.float64)
    eigenvalues, eigenvectors = np.linalg.eigh(A)
    U = eigenvectors @ np.diag(np.exp(-1j * eigenvalues * math.pi / 4)) @ eigenvectors.conj().T
    return U, eigenvalues, eigenvectors

# ── Eigenvalue products ───────────────────────────────────────────────

def eigenvalue_products(eigenvalues, k):
    from itertools import combinations_with_replacement
    products = []
    for indices in combinations_with_replacement(range(len(eigenvalues)), k):
        prod = 1.0 + 0j
        for idx in indices:
            prod *= eigenvalues[idx]
        products.append(prod)
    return products

def product_distribution(products):
    rounded = [round(p.real, 8) + round(p.imag, 8) * 1j for p in products]
    counts = Counter(rounded)
    total = len(products)
    distinct = len(counts)
    max_degen = max(counts.values()) if counts else 0
    # Entropy of the degeneracy distribution
    probs = [c / total for c in counts.values()]
    prod_entropy = -sum(p * math.log2(p) for p in probs if p > 1e-15) if probs else 0.0
    max_possible_entropy = math.log2(distinct) if distinct > 1 else 1.0
    return {
        "distinct": distinct,
        "total": total,
        "max_degeneracy": max_degen,
        "distinct_ratio": distinct / max(total, 1),
        "max_degen_ratio": max_degen / max(total, 1),
        "product_entropy": round(prod_entropy, 4),
        "entropy_ratio": round(prod_entropy / max_possible_entropy, 4) if distinct > 1 else 1.0,
        "is_concentrated": distinct < total * 0.5,
    }

# ── SLOS via Perceval (with timeout + checkpoint) ────────────────────

def run_slos(U, n_photons, n_modes, n_shots, cloud, dag, node_id, labels, k):
    """Run SLOS with 5-minute timeout and DAG checkpointing.
    
    The DAG node records:
    - The circuit (n_modes, n_photons)
    - The SLOS computation path (how it processes)
    - The output distribution
    - Whether it timed out
    """
    try:
        import perceval as pcvl
        import numpy as np
    except ImportError:
        dag.add_node(node_id, "slos",
                     {"labels": labels, "k": k, "n_modes": n_modes, "n_shots": n_shots},
                     {"error": "perceval not installed"}, "error", 0)
        return {"error": "perceval not installed"}
    
    t0 = time.time()
    
    try:
        # Build circuit
        try:
            circuit = pcvl.Unitary(U[:n_modes, :n_modes])
        except (AttributeError, TypeError):
            circuit = pcvl.Circuit(n_modes)
            for i in range(min(n_modes, 8)):
                circuit.add(i, pcvl.PS(float(np.angle(U[i, i]))))
            for i in range(n_modes - 1):
                circuit.add((i, i+1), pcvl.BS())
        
        input_state = pcvl.BasicState([1] * n_photons + [0] * (n_modes - n_photons))
        
        # Choose backend
        if cloud:
            try:
                processor = pcvl.RemoteProcessor(
                    "sim:ascella",
                    token=os.environ.get("PERCEVAL_TOKEN", "")
                )
                processor.set_circuit(circuit)
                processor.with_input(input_state)
            except Exception as e:
                # Fall back to local
                dag.add_node(node_id + "_cloud_fail", "slos",
                             {"labels": labels, "k": k, "error": str(e)},
                             {"fallback": "local"}, "error", time.time() - t0)
                processor = pcvl.Processor("SLOS", circuit)
                processor.with_input(input_state)
        else:
            processor = pcvl.Processor("SLOS", circuit)
            processor.with_input(input_state)
        
        # Run with timeout awareness
        # Quandela cloud: 5 min per shot
        # Local: no limit but we log elapsed
        sampler = pcvl.algorithm.Sampler(processor)
        
        # Checkpoint: circuit built, about to sample
        dag.add_node(node_id + "_circuit", "slos",
                     {"labels": labels, "k": k, "n_modes": n_modes,
                      "n_photons": n_photons, "cloud": cloud},
                     {"circuit_built": True, "input_state": str(input_state)},
                     "success", time.time() - t0)
        
        results = sampler.sample_count(n_shots)
        elapsed = time.time() - t0
        
        # Extract output distribution
        output_dist = {}
        for state, count in results["results"].items():
            prob = count / n_shots
            output_dist[str(tuple(state))] = prob
        
        probs = list(output_dist.values())
        entropy = -sum(p * math.log2(p) for p in probs if p > 1e-15)
        nonzero = sum(1 for p in probs if p > 1e-10)
        max_entropy = math.log2(max(nonzero, 1))
        entropy_ratio = entropy / max_entropy if max_entropy > 0 else 1.0
        kl_uniform = max_entropy - entropy  # D_KL(P||U) = log2(N) - H(P)
        
        result = {
            "n_modes": n_modes,
            "n_photons": n_photons,
            "n_shots": n_shots,
            "n_output_states": len(output_dist),
            "n_nonzero": nonzero,
            "entropy": round(entropy, 6),
            "max_entropy": round(max_entropy, 4),
            "entropy_ratio": round(entropy_ratio, 4),
            "kl_divergence": round(kl_uniform, 4),
            "max_prob": max(probs) if probs else 0,
            "elapsed_s": round(elapsed, 2),
            "cloud": cloud,
            "n_shots_actual": sum(results["results"].values()),
        }
        
        # Checkpoint: SLOS completed
        dag.add_node(node_id, "slos",
                     {"labels": labels, "k": k, "n_modes": n_modes},
                     result, "success", elapsed)
        
        return result
        
    except Exception as e:
        elapsed = time.time() - t0
        # Checkpoint: SLOS failed/timed out
        dag.add_node(node_id, "slos",
                     {"labels": labels, "k": k, "n_modes": n_modes},
                     {"error": str(e), "elapsed_before_timeout": round(elapsed, 2)},
                     "timeout" if elapsed > 290 else "error", elapsed)
        return {"error": str(e), "elapsed": round(elapsed, 2)}

# ── Bosonic tensor network (exact entropy, no sampling) ──────────────

def _symmetrize_2(T):
    return (T + T.swapaxes(0, 1)) / math.sqrt(2)

def _symmetrize_3(T):
    import itertools
    s = sum(T.transpose(p) for p in itertools.permutations([0, 1, 2]))
    return s / math.sqrt(6)

def _fock_probs_1(U):
    import numpy as np
    col0 = U[:, 0]
    probs = np.abs(col0) ** 2
    probs = probs[probs > 1e-15]
    return probs

def _fock_probs_2(U):
    import numpy as np
    N = U.shape[0]
    col0 = U[:, 0]
    col1 = U[:, 1]
    T_dist = np.outer(col0, col1)
    T = _symmetrize_2(T_dist)
    output_density = np.abs(T) ** 2
    fock = []
    for i in range(N):
        for j in range(i, N):
            p = float(output_density[i, j] + output_density[j, i]) if i != j else float(output_density[i, i])
            if p > 1e-15:
                fock.append(p)
    return np.array(fock)

def _fock_probs_3(U):
    import numpy as np
    N = U.shape[0]
    col0 = U[:, 0]
    col1 = U[:, 1]
    col2 = U[:, 2]
    T_dist = np.einsum('i,j,k->ijk', col0, col1, col2)
    T = _symmetrize_3(T_dist)
    output_density = np.abs(T) ** 2
    fock = []
    for i in range(N):
        for j in range(i, N):
            for k in range(j, N):
                if i == j == k:
                    p = float(output_density[i, i, i])
                elif i == j:
                    p = float(output_density[i, i, k] + output_density[i, k, i] + output_density[k, i, i])
                elif j == k:
                    p = float(output_density[i, j, j] + output_density[j, i, j] + output_density[j, j, i])
                else:
                    p = float(output_density[i, j, k] + output_density[i, k, j] +
                              output_density[j, i, k] + output_density[j, k, i] +
                              output_density[k, i, j] + output_density[k, j, i])
                if p > 1e-15:
                    fock.append(p)
    return np.array(fock)

def tensor_entropy(U, k):
    import numpy as np
    """Exact entropy via tensor network (no sampling noise)."""
    if k == 1:
        probs = _fock_probs_1(U)
    elif k == 2:
        probs = _fock_probs_2(U)
    elif k == 3:
        probs = _fock_probs_3(U)
    else:
        return None, None
    total = np.sum(probs)
    if total <= 0:
        return 0.0, 0
    p = probs / total
    entropy = float(-np.sum(p * np.log2(p + 1e-15)))
    nonzero = int(np.sum(p > 1e-15))
    max_ent = math.log2(max(nonzero, 1))
    return entropy, entropy / max_ent if max_ent > 0 else 1.0

# ── Tests with DAG ────────────────────────────────────────────────────

def spearman_rank(xs, ys):
    """Compute Spearman rank correlation between two lists of numbers."""
    n = len(xs)
    if n < 2:
        return 0.0, n
    # Rank both lists (lower value = rank 1)
    def rank(vals):
        sorted_vals = sorted(set(vals))
        return [sorted_vals.index(v) + 1 for v in vals]
    rx = rank(xs)
    ry = rank(ys)
    # Pearson on ranks
    mx = sum(rx) / n
    my = sum(ry) / n
    num = sum((rx[i] - mx) * (ry[i] - my) for i in range(n))
    dx = sum((rx[i] - mx) ** 2 for i in range(n))
    dy = sum((ry[i] - my) ** 2 for i in range(n))
    denom = (dx * dy) ** 0.5
    return num / denom if denom > 0 else 0.0, n

def test_sidon_vs_nonsidon(cloud, dag):
    """Compare SLOS for Sidon vs non-Sidon with rank correlation."""
    test_cases = [
        # Existing cases
        ([1, 2, 4, 8], "sidon_pow2"),
        ([1, 2, 3, 4], "nonsidon_seq"),
        ([1, 2, 5, 7], "sidon_h8"),
        ([1, 2, 5, 10, 16], "sidon_h16"),
        # Larger Sidon sets
        ([1, 2, 4, 8, 16, 32, 64], "sidon_pow7"),
        ([1, 2, 5, 11, 16, 19], "sidon_s6"),
        ([1, 2, 5, 11, 19, 28, 35], "sidon_s7"),
        # Dense non-Sidon (consecutive blocks of same size)
        ([5, 6, 7, 8, 9, 10], "nonsidon_dense6"),
        ([10, 11, 12, 13, 14, 15, 16], "nonsidon_dense7"),
        # Prime-based
        ([2, 3, 5, 7, 11, 13, 17], "primes_s7"),
        ([2, 3, 5, 7, 11, 13, 17, 19], "primes_s8"),
        # Sidon powers of 2 up to 128 (8 labels)
        ([1, 2, 4, 8, 16, 32, 64, 128], "sidon_pow8"),
    ]
    
    for k in [1, 2, 3, 4]:
        results = []
        for labels, desc in test_cases:
            node_prefix = f"{desc}_k{k}"
            
            # Eigenvalue products
            prod_id = f"{node_prefix}_products"
            if not dag.has_node(prod_id):
                S = build_sum_matrix(labels)
                U, eigenvalues, _ = matrix_to_unitary(S)
                if U is None:
                    finding("T1", "CRITICAL", f"{desc} K={k}: numpy available", "FAIL", {})
                    continue
                t0 = time.time()
                products = eigenvalue_products(eigenvalues, k)
                prod_dist = product_distribution(products)
                dag.add_node(prod_id, "eigenvalue_products",
                             {"labels": labels, "k": k},
                             prod_dist, "success", time.time() - t0)
            else:
                prod_dist = dag.get_node(prod_id)["result"]
            
            # SLOS
            slos_id = f"{node_prefix}_slos"
            if not dag.has_node(slos_id):
                S = build_sum_matrix(labels)
                U, _, _ = matrix_to_unitary(S)
                if U is None:
                    continue
                slos = run_slos(U, k, len(labels), 100000, cloud, dag, slos_id, labels, k)
            else:
                slos = dag.get_node(slos_id)["result"]
            
            if "error" in slos:
                continue
            
            results.append({
                "desc": desc,
                "labels": labels,
                "prod_distinct_ratio": prod_dist["distinct_ratio"],
                "slos_kl": slos.get("kl_divergence", 0.0),
                "slos_entropy_ratio": slos.get("entropy_ratio", 1.0),
            })
        
        # Rank correlation: do product distinct_ratios predict SLOS KL ordering?
        if len(results) >= 2:
            prod_vals = [r["prod_distinct_ratio"] for r in results]
            kl_vals = [r["slos_kl"] for r in results]
            rho, n = spearman_rank(prod_vals, kl_vals)
            
            # Negative rho = lower distinct_ratio → higher KL (confirmed prediction)
            # Strong rank agreement = rho < -0.5
            rank_confirmed = rho < -0.5
            
            for r in results:
                compare_id = f"{r['desc']}_k{k}_compare"
                if not dag.has_node(compare_id):
                    comparison = {
                        "k": k, **r,
                        "spearman_rho": round(rho, 4),
                        "rank_n": n,
                        "rank_confirmed": rank_confirmed,
                    }
                    dag.add_node(compare_id, "compare",
                                 {"labels": r["labels"], "k": k},
                                 comparison, "success", 0.0)
                    prod_id = f"{r['desc']}_k{k}_products"
                    slos_id = f"{r['desc']}_k{k}_slos"
                    dag.add_edge(prod_id, compare_id)
                    dag.add_edge(slos_id, compare_id)
            
            finding("T1", "CRITICAL",
                   f"K={k}: Spearman ρ={rho:.3f} (n={n}), "
                   f"{'lower distinct_ratio→higher KL CONFIRMED' if rank_confirmed else 'WEAK rank agreement'}"
                   f" — Sidon structure predicts SLOS concentration ordering",
                   "PASS" if rank_confirmed else "FAIL",
                   {"k": k, "spearman_rho": round(rho, 4), "n": n})
        # Also compute exact entropy via tensor network and compare with SLOS
        tensor_results = []
        for r in results:
            labels = r["labels"]
            S = build_sum_matrix(labels)
            U, _, _ = matrix_to_unitary(S)
            if U is None:
                continue
            ten_ent, ten_ratio = tensor_entropy(U, k)
            if ten_ent is not None:
                tensor_results.append({
                    "desc": r["desc"],
                    "slos_kl": r["slos_kl"],
                    "tensor_entropy": ten_ent,
                    "tensor_entropy_ratio": ten_ratio,
                })
        
        if len(tensor_results) >= 2:
            slos_kl_vals = [tr["slos_kl"] for tr in tensor_results]
            tensor_ent_vals = [tr["tensor_entropy"] for tr in tensor_results]
            rho_tensor, n_tensor = spearman_rank(slos_kl_vals, tensor_ent_vals)
            
            # Inverse: higher KL (SLOS) should correlate with lower entropy (tensor)
            # Negative rho = SLOS higher KL → tensor lower entropy_ratio (confirmed)
            tensor_ratio_vals = [tr["tensor_entropy_ratio"] for tr in tensor_results]
            rho_ratio, _ = spearman_rank(slos_kl_vals, tensor_ratio_vals)
            
            tensor_confirmed = rho_ratio < -0.5
            
            finding("T1", "CRITICAL",
                   f"K={k}: SLOS vs tensor ρ={rho_ratio:.3f} (n={n_tensor}), "
                   f"{'METHODS AGREE on ordering' if tensor_confirmed else 'WEAK agreement between methods'}",
                   "PASS" if tensor_confirmed else "FAIL",
                   {"k": k, "spearman_rho_tensor_vs_slos": round(rho_tensor, 4),
                    "spearman_rho_tensor_ratio_vs_slos_kl": round(rho_ratio, 4), "n": n_tensor})
        else:
            finding("T1", "CRITICAL", f"K={k}: insufficient data for tensor vs SLOS",
                   "FAIL", {"k": k, "reason": "insufficient tensor results"})

# ── Main ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    cloud = "--cloud" in sys.argv
    resume = "--resume" in sys.argv
    
    if cloud and not os.environ.get("PERCEVAL_TOKEN"):
        print("WARNING: --cloud but PERCEVAL_TOKEN not set. Use local.")
        cloud = False
    
    print("=" * 60)
    print("  SLOS Eigenvalue Product Verification")
    print(f"  Mode: {'QUANDELA CLOUD (5-min/shot)' if cloud else 'LOCAL'}")
    print(f"  Resume: {'YES (from checkpoint)' if resume else 'NO'}")
    print("=" * 60)
    
    dag = ComputationDAG()
    if resume:
        dag.load_checkpoint()
    
    print("\n[T1] Sidon vs non-Sidon: do products predict SLOS?")
    test_sidon_vs_nonsidon(cloud, dag)
    
    # Save DAG report
    report_path = dag.save_dag_report()
    
    # Write EVAL
    ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
    eval_path = ARTIFACTS_DIR / "EVAL.md"
    total = len(_findings)
    passed = sum(1 for f in _findings if f["verdict"] == "PASS")
    failed = sum(1 for f in _findings if f["verdict"] == "FAIL")
    
    lines = [
        "# EVAL.md — SLOS Eigenvalue Product Verification\n",
        f"**Mode:** {'Quandela Cloud' if cloud else 'Local SLOS'}",
        f"**Overall:** {'PASS' if failed==0 else 'FAIL'}",
        f"**Checks:** {total} total, {passed} PASS, {failed} FAIL",
        f"**DAG nodes:** {len(dag.nodes)}",
        f"**DAG report:** {report_path}\n",
        "## Results\n",
        "| Test | Claim | Verdict |\n|------|-------|---------|",
    ]
    for f in _findings:
        lines.append(f"| {f['module']} | {f['claim'][:70]} | {f['verdict']} |")
    lines.append(f"\n## Computation DAG\nSee: {report_path}")
    lines.append(f"## Checkpoints\nSee: {CHECKPOINT_DIR}/")
    
    eval_path.write_text("\n".join(lines))
    
    # Save evidence
    evidence_path = ARTIFACTS_DIR / "slos_verify_evidence.jsonl"
    with open(evidence_path, "w") as f:
        for obj in _findings:
            f.write(json.dumps(obj, default=str) + "\n")
    
    print(f"\n{'='*60}")
    print(f"  Results: {passed} PASS, {failed} FAIL out of {total}")
    print(f"  DAG: {len(dag.nodes)} nodes, {len(dag.edges)} edges")
    print(f"  Report: {report_path}")
    print(f"  Checkpoints: {CHECKPOINT_DIR}/")
    print(f"  To resume: python3 scripts/perceval_slos_verify.py --resume")
    print(f"{'='*60}")