#!/usr/bin/env python3
"""prime_slos_explore.py — Spectral signature of primes in SLOS concentration.

Tests whether prime-based label sets consistently sit between Sidon
(pure, maximal concentration) and dense non-Sidon (minimal concentration)
in eigenvalue product and SLOS KL divergence — across scales from small
primes (2,3,5,7,11,13,17) up to primes near 1 billion.

Hypothesis:
  Primes have partial Sidon-like additive structure: sums of primes are
  more constrained than sums of consecutive integers (non-Sidon) but less
  constrained than sums of powers of 2 (Sidon). If true, the ordering
  Sidon < primes < non-Sidon (by KL divergence from uniform) should hold
  at every scale and every set size.

Method:
  1. Generate label sets at 4 scales (small, 10^3, 10^6, 10^9)
  2. For each: compute eigenvalue product distinct_ratio (all sets)
  3. For each: compute exact tensor network entropy (N ≤ 8)
  4. For a subset: run sampled SLOS (100k shots) to cross-validate
  5. Report whether the ordering Sidon < primes < non-Sidon holds
"""

import sys
import os
import math
import json
import time
import hashlib
import itertools
from pathlib import Path

import numpy as np

REPO_ROOT = Path(__file__).resolve().parent.parent
ARTIFACTS_DIR = REPO_ROOT / ".openresearch" / "artifacts"
OUTPUT_PATH = ARTIFACTS_DIR / "prime_slos_explore.json"

# ── Prime generation ────────────────────────────────────────────────────────

def primes_upto(limit):
    """Segmented sieve for primes up to limit. Returns list."""
    if limit < 2:
        return []
    sieve = bytearray(b'\x01') * (limit + 1)
    sieve[0:2] = b'\x00\x00'
    for i in range(2, int(limit ** 0.5) + 1):
        if sieve[i]:
            step = i
            start = i * i
            sieve[start:limit+1:step] = b'\x00' * ((limit - start) // step + 1)
    return [i for i, is_prime in enumerate(sieve) if is_prime]

def prime_clusters(count, offset_start=0):
    """Return `count` consecutive primes starting from offset_start-th prime."""
    # For small and medium limits, use sieve
    if offset_start < 1_000_000:
        limit = max(offset_start * 2 + 100, 10_000_000)
        all_primes = primes_upto(limit)
        if offset_start + count > len(all_primes):
            # Extend with nextprime for large offsets
            import sympy
            p = sympy.prime(offset_start + 1)
            cluster = []
            for _ in range(count):
                cluster.append(p)
                p = sympy.nextprime(p)
            return cluster
        return all_primes[offset_start:offset_start + count]
    else:
        import sympy
        p = sympy.prime(offset_start + 1)
        cluster = []
        for _ in range(count):
            cluster.append(p)
            p = sympy.nextprime(p)
        return cluster

# ── Sidon/non-Sidon reference sets ──────────────────────────────────────────

def sidon_pow2_set(size):
    return [1 << i for i in range(size)]

def nonsidon_consecutive_set(size, start=1):
    return list(range(start, start + size))

# ── Matrix and eigenvalue product ───────────────────────────────────────────

def build_sum_matrix(labels):
    n = len(labels)
    S = np.zeros((n, n), dtype=np.float64)
    for i in range(n):
        for j in range(n):
            S[i, j] = float(labels[i] + labels[j])
    max_val = S.max()
    if max_val > 0:
        S /= max_val
    return S

def matrix_to_unitary(S):
    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

def eigenvalue_products(eigenvalues, k):
    products = []
    for indices in itertools.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 = {}
    for r in rounded:
        counts[r] = counts.get(r, 0) + 1
    total = len(products)
    distinct = len(counts)
    max_degen = max(counts.values()) if counts else 0
    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,
        "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,
    }

# ── Tensor network entropy (exact, no sampling) ────────────────────────────

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

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

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

def _fock_probs_2(U):
    N = U.shape[0]
    col0, col1 = U[:, 0], U[:, 1]
    T = _symmetrize_2(np.outer(col0, col1))
    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):
    N = U.shape[0]
    col0, col1, col2 = U[:, 0], U[:, 1], 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):
    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

# ── Run SLOS via Perceval (sampled) ─────────────────────────────────────────

def run_slos(U, k, n_shots=100000):
    """Run Perceval SLOS for a given unitary and photon number."""
    try:
        import perceval as pcvl
    except ImportError:
        return None
    n_modes = U.shape[0]
    circuit = pcvl.Unitary(U[:n_modes, :n_modes])
    input_state = pcvl.BasicState([1] * k + [0] * (n_modes - k))
    processor = pcvl.Processor("SLOS", circuit)
    processor.with_input(input_state)
    sampler = pcvl.algorithm.Sampler(processor)
    results = sampler.sample_count(n_shots)
    probs = [c / n_shots for c in results["results"].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_ent = math.log2(max(nonzero, 1))
    return {
        "entropy": round(entropy, 6),
        "entropy_ratio": round(entropy / max_ent, 4) if max_ent > 0 else 1.0,
        "nonzero": nonzero,
        "max_prob": round(max(probs), 6) if probs else 0,
    }

# ── Main exploration ────────────────────────────────────────────────────────

def explore():
    print("=" * 60)
    print("  Prime SLOS Spectral Exploration")
    print("=" * 60)
    
    scales = {
        "small": 0,        # first primes: 2,3,5,7,11,13,17,19
        "kilo": 168,       # primes near 1000 (~168th prime = 997)
        "million": 78498,  # primes near 10^6 (~78498 primes below 1M)
        "billion": None,   # primes near 10^9
        "trillion": None,  # primes near 10^12
        "quadrillion": None, # primes near 10^15
        "quintillion": None, # primes near 10^18
    }
    
    sizes = [5, 6, 7, 8, 10]
    k_values = [2]
    results = []
    
    scales_list = []
    
    # Generate Sidon references
    print("\n--- Generating reference sets ---")
    for size in sizes:
        scales_list.append({
            "desc": f"sidon_pow2_s{size}",
            "category": "sidon",
            "labels": sidon_pow2_set(size),
            "scale": "reference",
        })
        scales_list.append({
            "desc": f"nonsidon_seq_s{size}",
            "category": "nonsidon",
            "labels": nonsidon_consecutive_set(size),
            "scale": "reference",
        })
    
    # Map scale names to starting values for sympy.nextprime
    scale_starts = {
        "small": None,      # Use sieve (offset-based)
        "kilo": None,       # Use sieve (offset-based)
        "million": None,    # Use sieve (offset-based)
        "billion": 10**9,
        "trillion": 10**12,
        "quadrillion": 10**15,
        "quintillion": 10**18,
    }
    
    # Generate prime sets at each scale
    print("\n--- Generating prime sets at all scales ---")
    for scale_name, offset in scales.items():
        start_val = scale_starts[scale_name]
        
        if start_val is not None:
            # Large scale: use sympy.nextprime from starting value
            try:
                import sympy
                for size in sizes:
                    p = sympy.nextprime(start_val)
                    cluster = []
                    for _ in range(size):
                        cluster.append(p)
                        p = sympy.nextprime(p)
                    scales_list.append({
                        "desc": f"primes_{scale_name}_s{size}",
                        "category": "prime",
                        "labels": cluster,
                        "scale": scale_name,
                    })
            except ImportError:
                print(f"  sympy not available, skipping {scale_name} scale")
            continue
        
        for size in sizes:
            try:
                cluster = prime_clusters(size, offset)
                scales_list.append({
                    "desc": f"primes_{scale_name}_s{size}",
                    "category": "prime",
                    "labels": cluster[:size],
                    "scale": scale_name,
                })
            except Exception as e:
                print(f"  Error at {scale_name} size={size}: {e}")
    
    print(f"\n--- Computing eigenvalue products for {len(scales_list)} sets ---")
    t0 = time.time()
    
    for entry in scales_list:
        labels = entry["labels"]
        
        # Eigenvalue products (fast)
        S = build_sum_matrix(labels)
        U, eigenvalues = matrix_to_unitary(S)
        
        prods = {}
        for k in [2]:
            products = eigenvalue_products(eigenvalues, k)
            prods[k] = product_distribution(products)
        
        # Tensor network entropy (exact, N ≤ 8)
        tensor = {}
        for k in [1, 2, 3]:
            if len(labels) <= 8:
                ent, ratio = tensor_entropy(U, k)
                if ent is not None:
                    tensor[k] = {"entropy": round(ent, 4), "entropy_ratio": round(ratio, 4)}
        
        # SLOS (sampled, N ≤ 8, K=2 only)
        slos = {}
        if len(labels) <= 8:
            slos_result = run_slos(U, 2, 100000)
            if slos_result:
                slos[2] = slos_result
        
        result = {
            "desc": entry["desc"],
            "category": entry["category"],
            "scale": entry["scale"],
            "n_labels": len(labels),
            "labels": labels[:5] if len(labels) > 5 else labels,  # Truncate for readability
            "labels_full": labels,
            "eigenvalue_products": {str(k): v for k, v in prods.items()},
            "tensor_entropy": {str(k): v for k, v in tensor.items()},
            "slos": {str(k): v for k, v in slos.items()},
        }
        results.append(result)
        
        prod_r = prods[2]["distinct_ratio"]
        tensor_str = f" tensor_K2_ER={tensor[2]['entropy_ratio']:.3f}" if 2 in tensor else ""
        slos_str = f" slos_K2_ER={slos[2]['entropy_ratio']:.3f}" if 2 in slos else ""
        print(f"  {entry['desc']:30s} K=2 prod_ratio={prod_r:.4f}{tensor_str}{slos_str}")
    
    elapsed = time.time() - t0
    print(f"\n  Computed {len(scales_list)} sets in {elapsed:.1f}s")
    
    # ── Analysis: does the ordering hold? ─────────────────────────────────
    print("\n--- Analysis: ordering test ---")
    
    analysis = []
    for size in sizes:
        # Collect by category
        sidon_vals = [r for r in results if r["category"] == "sidon" and r["n_labels"] == size]
        nonsidon_vals = [r for r in results if r["category"] == "nonsidon" and r["n_labels"] == size]
        prime_vals = [r for r in results if r["category"] == "prime" and r["n_labels"] == size]
        
        for k in k_values:
            sk = str(k)
            for pv in prime_vals:
                prod_r = pv["eigenvalue_products"].get(sk, {}).get("distinct_ratio", None)
                tensor_er = pv.get("tensor_entropy", {}).get(sk, {}).get("entropy_ratio", None)
                
                if prod_r is None:
                    continue
                
                # Compare with Sidon and non-Sidon baselines
                s_prod = sidon_vals[0]["eigenvalue_products"][sk]["distinct_ratio"] if sidon_vals else None
                n_prod = nonsidon_vals[0]["eigenvalue_products"][sk]["distinct_ratio"] if nonsidon_vals else None
                
                # For product ratio: lower = more concentrated
                # Expected: Sidon < primes < non-Sidon
                if s_prod is not None and n_prod is not None:
                    if s_prod <= prod_r <= n_prod:
                        product_ordering = "CONFIRMED"
                    elif prod_r < s_prod:
                        product_ordering = "ABOVE_SIDON"
                    elif prod_r > n_prod:
                        product_ordering = "BELOW_NONSIDON"
                    else:
                        product_ordering = "UNORDERED"
                else:
                    product_ordering = "NO_BASELINE"
                
                # For tensor entropy ratio: lower = more concentrated
                # Expected: Sidon < primes < non-Sidon
                if tensor_er is not None and s_prod is not None and n_prod is not None:
                    s_tensor = sidon_vals[0].get("tensor_entropy", {}).get(sk, {}).get("entropy_ratio", None)
                    n_tensor = nonsidon_vals[0].get("tensor_entropy", {}).get(sk, {}).get("entropy_ratio", None)
                    if s_tensor is not None and n_tensor is not None:
                        if s_tensor <= tensor_er <= n_tensor:
                            tensor_ordering = "CONFIRMED"
                        elif tensor_er < s_tensor:
                            tensor_ordering = "ABOVE_SIDON"
                        elif tensor_er > n_tensor:
                            tensor_ordering = "BELOW_NONSIDON"
                        else:
                            tensor_ordering = "UNORDERED"
                    else:
                        tensor_ordering = "NO_BASELINE"
                else:
                    tensor_ordering = "NO_TENSOR"
                
                entry = {
                    "size": size,
                    "k": k,
                    "desc": pv["desc"],
                    "scale": pv["scale"],
                    "prod_ratio": prod_r,
                    "tensor_entropy_ratio": tensor_er,
                    "sidon_prod_ratio": s_prod,
                    "nonsidon_prod_ratio": n_prod,
                    "product_ordering": product_ordering,
                    "tensor_ordering": tensor_ordering,
                }
                analysis.append(entry)
                
                print(f"  {pv['desc']:35s} K={k} prod_r={prod_r:.4f} [{product_ordering}] "
                      f"tensor_ER={tensor_er} [{tensor_ordering}]")
    
    # ── Summary ───────────────────────────────────────────────────────────
    print("\n--- Summary ---")
    confirmed = sum(1 for a in analysis if a["product_ordering"] == "CONFIRMED")
    above = sum(1 for a in analysis if a["product_ordering"] == "ABOVE_SIDON")
    below = sum(1 for a in analysis if a["product_ordering"] == "BELOW_NONSIDON")
    total = len(analysis)
    print(f"  Product ordering Sidon < primes < non-Sidon:")
    print(f"    CONFIRMED: {confirmed}/{total}")
    print(f"    Above Sidon (even more concentrated): {above}/{total}")
    print(f"    Below non-Sidon (less concentrated): {below}/{total}")
    
    t_confirmed = sum(1 for a in analysis if a["tensor_ordering"] == "CONFIRMED")
    t_above = sum(1 for a in analysis if a["tensor_ordering"] == "ABOVE_SIDON")
    t_below = sum(1 for a in analysis if a["tensor_ordering"] == "BELOW_NONSIDON")
    t_total = sum(1 for a in analysis if a["tensor_ordering"] != "NO_TENSOR" and a["tensor_ordering"] != "NO_BASELINE")
    print(f"  Tensor entropy ordering Sidon < primes < non-Sidon:")
    print(f"    CONFIRMED: {t_confirmed}/{t_total}")
    print(f"    Above Sidon: {t_above}/{t_total}")
    print(f"    Below non-Sidon: {t_below}/{t_total}")
    
    # ── Save ──────────────────────────────────────────────────────────────
    output = {
        "schema": "prime_slos_explore_v1",
        "claim_boundary": "prime-slos-spectral-signature:ordering-test",
        "config": {"sizes": sizes, "k_values": k_values, "scales": list(scales.keys())},
        "results": results,
        "analysis": analysis,
        "summary": {
            "confirmed": confirmed,
            "above_sidon": above,
            "below_nonsidon": below,
            "total": total,
            "tensor_confirmed": t_confirmed,
            "tensor_total": t_total,
        },
    }
    ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
    with open(OUTPUT_PATH, "w") as f:
        json.dump(output, f, indent=2, default=str)
    print(f"\n  Saved to {OUTPUT_PATH}")
    print(f"  To run with SLOS (if Perceval available): needs perceval installed")
    print("=" * 60)

if __name__ == "__main__":
    explore()