"""loo_validator.py
Leave-one-out (LOO) cross-validation for BigG overfit and PLN scorers.
Uses the 15 known Mersenne prime digit files for extended feature extraction.

Steps:
1. Load all 52 known exponents + 15 digit files.
2. Extract features (exponent-based + digit-file-based where available).
3. LOO for BigG: for each known prime p_i, train on 51, score p_i vs 1000
   random neighbor primes. Report rank percentile.
4. Feature attribution: which features are most stable / most discriminating.
5. Calibration: convert nearest-neighbor similarity to interpretable confidence.
"""
import os
import re
import sys
import math
import time
import struct
import random
import subprocess

# ================================================================
# ALL 52 KNOWN MERSENNE EXPONENTS
# ================================================================
ALL_P = [
    2,3,5,7,13,17,19,31,61,89,107,127,521,607,1279,2203,2281,3217,4253,
    4423,9689,9941,11213,19937,21701,23209,44497,86243,110503,132049,
    216091,756839,859433,1257787,1398269,2976221,3021377,6972593,13466917,
    20996011,24036583,25964951,30402457,32582657,37156667,42643801,43112609,
    57885161,74207281,77232917,82589933,136279841
]

# ================================================================
# DIGIT FILE FEATURES
# ================================================================
def load_digit_features(exponent, directory='.'):
    """Load digit file for a known prime and extract statistical features.
    Returns dict of features, or None if file not available."""
    # Try both capitalisations: M<p>.txt and m<p>.txt
    for prefix in ('M', 'm'):
        fn = os.path.join(directory, f'{prefix}{exponent}.txt')
        if os.path.exists(fn):
            with open(fn) as f:
                raw = f.read().replace('\n', '').replace('\r', '')
            total = len(raw)
            if total < 100:
                return None
            expected = total / 10.0
            hist = [raw.count(str(d)) for d in range(10)]
            chi_sq = sum((h - expected) ** 2 / expected for h in hist)
            entropy = -sum((h / total) * math.log2(h / total) for h in hist if h > 0)
            # Last 16 digits as a float fingerprint
            last16 = int(raw[-16:]) / 1e16
            # First 16 digits (leading digit distribution follows Benford's law)
            first_digit_hist = [raw.count(str(d), 0, min(len(raw), 100000)) for d in range(1, 10)]
            first_d_total = sum(first_digit_hist)
            benford_chisq = 0.0
            if first_d_total > 0:
                for d in range(1, 10):
                    expected_b = first_d_total * math.log10(1 + 1/d)
                    benford_chisq += (first_digit_hist[d-1] - expected_b) ** 2 / max(expected_b, 1)
            return {
                'chi_sq': chi_sq,
                'entropy': entropy,
                'last16': last16,
                'benford_chisq': benford_chisq,
                'num_digits': total,
            }
    return None

# ================================================================
# PLN COORDINATE (mirrors mersenne_sieve.c)
# ================================================================
PHI    = (1.0 + math.sqrt(5.0)) / 2.0
LN_PHI = math.log(PHI)
LN2    = math.log(2.0)
BW      = 0.5   # omega-space bandwidth = Re(s) = 1/2 on the Riemann critical line
BW_BETA = 0.10  # beta-space bandwidth: tightens calibration to primes with same omega-phase

# Per-dimension signal weights for the Mahalanobis scorer.
# Measured per-dim LOO rank (fraction of randoms with lower |z| per dim):
#   Best: D15=28.8%, D12=29.7%, D4=30.0%, D10=31.5%, D5=32.6%, D1=33.3%
#   Weak: D3=44.2%, D16=43.1%, D2=40.4%
# Signal-squared weighting tested (18.5% vs 18.6%) — marginal, caused regression
# on p=31 (26.1%→55.6%). Equal weights remain optimal overall.
# Preserved here for future per-band adaptive weighting experiments.
_DIM_RANKS = [33.3, 40.4, 44.2, 30.0, 32.6, 39.2, 37.1, 35.0,
              35.2, 31.5, 37.3, 29.7, 35.5, 36.1, 28.8, 43.1]
DIM_WEIGHTS = [1.0] * 16   # equal weights — optimal at current training size

KNOWN_SET = set(ALL_P)   # for contamination filter

def compute_pln(p):
    omega = math.log(math.log(p * math.log(2)) / math.log(PHI)) / math.log(PHI)
    D = int(omega % 8) + 1
    theta = omega % 1.0
    dw = omega % 1.0
    r_raw = dw / (1.0 + dw)
    r = r_raw
    # orbit_r refinement (same as mersenne_sieve.c)
    for _ in range(8):
        r = r_raw * (1.0 + math.exp(-r * r / (2.0 * PHI)))
    return {'omega': omega, 'D': D, 'theta': theta, 'r': r}

# ================================================================
# C-EXACT 16-DIM FINGERPRINT (mirrors dn_resonance_checker.c exactly)
# ================================================================
GOLDEN_ANG  = 3.88322207745300172950  # 2*pi/phi — irrational dimension spacing
OMEGA_MAX   = 8.0                     # prism range: all 52 known in [-0.81, 7.57]
PRM16       = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
NUM_DIMS_C  = 16

def gen_fib_c(x):
    """Generalised Fibonacci via Binet: (phi^x - cos(pi*x)*phi^{-x}) / sqrt(5).
    Matches C gen_fib() exactly."""
    phi_x     = PHI ** x
    phi_neg_x = PHI ** (-x)
    return (phi_x - math.cos(math.pi * x) * phi_neg_x) / math.sqrt(5.0)

def prismatic_omega_c(omega_p, n):
    """Per-dimension field tension.  Matches C prismatic_omega()."""
    return 0.5 + 0.5 * math.sin(omega_p + float(n) * GOLDEN_ANG)

def compute_Dn_c(n, r, omega_p):
    """Full unified lattice operator |L_n|.  Exact port of C compute_Dn().

    |L_n| = sqrt( (coeff*|r|^k + 1_eff*cos(pi*beta))^2
                 + (1_eff*sin(pi*beta))^2 )

    where:
      beta   = frac(omega_p)
      x      = n + beta          (non-integer Fibonacci index)
      F_x    = gen_fib_c(x)      (Binet generalisation)
      Omega_n = 0.5+0.5*sin(omega_p + n*GOLDEN_ANG)
      coeff  = sqrt(phi * F_x * 2^n * P_n * Omega_n * phi^{k*x})
      1_eff  = 1 + |cos(pi*beta*phi)| * ln(P_n) / phi^x
    """
    beta    = omega_p - math.floor(omega_p)
    x       = float(n) + beta
    k       = float(n + 1) / 8.0
    F_x     = abs(gen_fib_c(x))
    Omega_n = prismatic_omega_c(omega_p, n)
    P_n     = float(PRM16[n - 1])
    two_n   = 2.0 ** float(n)
    phi_kx  = PHI ** (k * x)
    coeff   = math.sqrt(max(0.0, PHI * F_x * two_n * P_n * Omega_n * phi_kx))
    delta   = abs(math.cos(math.pi * beta * PHI)) * math.log(P_n) / (PHI ** x)
    one_eff = 1.0 + delta
    scale   = coeff * (abs(r) ** k)
    phase_re = one_eff * math.cos(math.pi * beta)
    phase_im = one_eff * math.sin(math.pi * beta)
    re = scale + phase_re
    return math.sqrt(re * re + phase_im * phase_im)

def compute_fingerprint(p):
    """16-dim fingerprint matching dn_resonance_checker.c exactly.
    r = omega / OMEGA_MAX  (C-exact: NOT fractional-part sigmoid).
    """
    om = get_omega(p)
    r  = om / OMEGA_MAX
    return [compute_Dn_c(n, r, om) for n in range(1, NUM_DIMS_C + 1)]

FP_LEN = NUM_DIMS_C   # 16


# ================================================================
# LOO MODEL: Gaussian-weighted local Mahalanobis (BW=0.5 in omega space)
# mirrors dn_resonance_checker.c with LOCAL_BW=0.5
# ================================================================
def get_omega(p):
    """PLN omega for any prime p (fast, no full pln needed)."""
    inner = math.log(p * LN2) / LN_PHI
    if inner <= 0.0:
        return -999.0
    return math.log(inner) / LN_PHI

def get_beta(p):
    """beta = frac(omega) — the omega-phase of p in the PLN spiral."""
    om = get_omega(p)
    return om - math.floor(om)

def local_mahal(fp_target, fps_train, omegas_train, betas_train,
                omega_query, beta_query, bw=BW, bw_beta=BW_BETA):
    """2D Gaussian-weighted Mahalanobis distance.
    Kernel: exp(-0.5 * (dOmega/BW)^2) * exp(-0.5 * (dBeta/BW_BETA)^2)

    The 2D kernel ensures the cluster mean/std used for scoring is calibrated
    against training primes that match BOTH the omega position AND the
    beta = frac(omega) phase of the query prime.  This fixes the 'edge-beta'
    failure mode where primes with beta near 0 or 1 appeared as outliers from
    the average-beta centroid computed by the 1D (omega-only) kernel.
    """
    n = len(fps_train)
    weights = [
        math.exp(-0.5 * (
            ((omegas_train[i] - omega_query) / bw) ** 2 +
            ((betas_train[i]  - beta_query)  / bw_beta) ** 2
        ))
        for i in range(n)
    ]
    total_w = sum(weights)
    if total_w < 1e-15:
        # Fallback: 1D omega-only kernel (maintains coverage for isolated primes)
        weights = [math.exp(-0.5 * ((omegas_train[i] - omega_query) / bw) ** 2)
                   for i in range(n)]
        total_w = sum(weights)
    if total_w < 1e-15:
        weights = [1.0] * n
        total_w = float(n)

    # Gaussian-weighted mean and std per feature
    feat_stat = []
    for f in range(FP_LEN):
        wm = sum(weights[i] * fps_train[i][f] for i in range(n)) / total_w
        wv = sum(weights[i] * (fps_train[i][f] - wm) ** 2 for i in range(n)) / total_w
        ws = math.sqrt(wv) if wv > 1e-20 else 1.0
        feat_stat.append((wm, ws))

    # Mahalanobis distance (normalised by # features so scales don't depend on FP_LEN).
    # DIM_WEIGHTS are equal (1.0) — signal-weighting tested but caused regression on p=31.
    dist2 = sum(((fp_target[f] - feat_stat[f][0]) / feat_stat[f][1]) ** 2
                for f in range(FP_LEN))
    return math.sqrt(dist2 / FP_LEN)

# Keep legacy score_nn for comparison if needed
def fit_loo_model(fps):
    n = len(fps)
    feat_std = []
    for f in range(FP_LEN):
        vals = [fps[i][f] for i in range(n)]
        mean = sum(vals) / n
        var = sum((v - mean) ** 2 for v in vals) / n
        feat_std.append(math.sqrt(var) if var > 1e-20 else 1.0)
    return feat_std

# ================================================================
# PRIMALITY (simple trial division up to 2^20)
# ================================================================
def is_probably_prime(n):
    if n < 2: return False
    if n == 2: return True
    if n % 2 == 0: return False
    if n < 9: return True
    if n % 3 == 0: return False
    r = int(math.isqrt(n))
    d = 5
    while d <= min(r, 1 << 20):
        if n % d == 0 or n % (d + 2) == 0:
            return False
        d += 6
    return True

# ================================================================
# MAIN LOO VALIDATION
# ================================================================
def run_loo():
    print("=== LOO CROSS-VALIDATION ===")
    print(f"  {len(ALL_P)} known Mersenne prime exponents")
    print()

    # Precompute all fingerprints
    print("  Computing fingerprints...")
    t0 = time.time()
    fps = {p: compute_fingerprint(p) for p in ALL_P}
    print(f"  Done in {time.time()-t0:.2f}s")
    print()

    # Load digit-file features where available
    digit_feats = {}
    for p in ALL_P:
        df = load_digit_features(p)
        if df:
            digit_feats[p] = df
    print(f"  Digit files available for {len(digit_feats)} of {len(ALL_P)} primes:")
    print(f"    {sorted(digit_feats.keys())}")
    print()

    # LOO results per held-out prime
    loo_results = []  # (p, held_out_score, rank_pct, n_sampled)
    K_SAMPLES = 500   # random neighbors to compare against

    random.seed(42)

    # Precompute omega AND beta for all known primes
    fps_omega = {p: get_omega(p) for p in ALL_P}
    fps_beta  = {p: get_beta(p)  for p in ALL_P}

    print(f"  Running {len(ALL_P)} LOO trials ({K_SAMPLES} random samples each)...")
    print(f"  Calibration: BW_omega={BW}, BW_beta={BW_BETA} (2D omega×beta kernel)")
    print()
    print(f"  {'p':>10}  {'dist':>8}  {'rank%':>6}  n_valid  n_eff   notes")
    print("  " + "-"*70)

    for p_i in ALL_P:
        # XorShift-64 seed matching C: p ^ 0xDEADBEEFCAFEBABE
        random.seed(p_i ^ 0xDEADBEEFCAFEBABE)
        omega_i = fps_omega[p_i]
        beta_i  = fps_beta[p_i]

        # Build training set (all except p_i)
        train       = [p for p in ALL_P if p != p_i]
        train_fps   = [fps[p] for p in train]
        train_oms   = [fps_omega[p] for p in train]
        train_betas = [fps_beta[p]  for p in train]

        # Effective neighbourhood size (1D omega weights, for display)
        n_eff = sum(math.exp(-0.5 * ((om - omega_i) / BW) ** 2) for om in train_oms)

        # 2D Gaussian-weighted Mahalanobis distance for held-out prime
        loo_dist = local_mahal(fps[p_i], train_fps, train_oms, train_betas, omega_i, beta_i)

        # C-exact window: [p//2, p*2] — log-symmetric around p_i
        lo  = max(3, p_i // 2)
        hi  = p_i * 2
        if hi <= lo + 1000:
            hi = lo + 1001
        lo_odd = lo if lo % 2 == 1 else lo + 1
        if lo_odd >= hi:
            hi = lo_odd + 2
        span = hi - lo_odd

        # Sample raw odd integers (matches C: no primality filter, just skip known)
        n_better = 0
        n_valid  = 0
        for _ in range(K_SAMPLES):
            cand = lo_odd + (random.randrange(span) // 2) * 2
            if cand in KNOWN_SET:
                continue    # contamination filter
            n_valid += 1
            fp_cand   = compute_fingerprint(cand)
            # random integer scored against the SAME 2D calibration centered at (omega_i, beta_i)
            dist_cand = local_mahal(fp_cand, train_fps, train_oms, train_betas, omega_i, beta_i)
            if dist_cand < loo_dist:
                n_better += 1

        rank_pct = 100.0 * n_better / n_valid if n_valid > 0 else 50.0

        if n_valid == 0:
            continue

        notes = ""
        if rank_pct <= 5.0:   notes = "*** STRONG"
        elif rank_pct <= 15.0: notes = "** GOOD"
        elif rank_pct <= 33.0: notes = "* OK"
        elif rank_pct >= 50.0: notes = "(no signal)"

        loo_results.append((p_i, loo_dist, rank_pct, n_valid))
        print(f"  {p_i:>10d}  {loo_dist:8.4f}  {rank_pct:6.1f}%  {n_valid:>7}  {n_eff:5.1f}  {notes}")

    # Summary statistics
    print()
    print("=== LOO SUMMARY ===")
    # LOW rank_pct = good signal (matching dn_resonance_checker.c convention)
    strong = sum(1 for _, _, rp, _ in loo_results if rp <=  5.0)
    good   = sum(1 for _, _, rp, _ in loo_results if rp <= 15.0)
    ok     = sum(1 for _, _, rp, _ in loo_results if rp <= 33.0)
    no_sig = sum(1 for _, _, rp, _ in loo_results if rp >= 50.0)
    n = len(loo_results)
    mean_rp = sum(rp for _, _, rp, _ in loo_results) / n if n else 0
    print(f"  Trials: {n}")
    print(f"  Mean rank percentile: {mean_rp:.1f}%  (random baseline = 50%, target ~33%)")
    print(f"  STRONG (<= 5th pct):  {strong}/{n} = {100*strong/n:.1f}%")
    print(f"  GOOD   (<= 15th pct): {good}/{n} = {100*good/n:.1f}%")
    print(f"  OK     (<= 33rd pct): {ok}/{n} = {100*ok/n:.1f}%")
    print(f"  No sig (>= 50th pct): {no_sig}/{n} = {100*no_sig/n:.1f}%")
    verdict = ("STRONG discriminative power" if mean_rp < 35.0 else
               "Moderate signal above random" if mean_rp < 45.0 else
               "Marginal signal above random" if mean_rp < 50.0 else
               "No discriminative power (at/below random baseline)")
    print(f"  Verdict: {verdict}")
    print()

    # Stratified by omega band (mirrors dn_resonance_checker.c --loo output)
    print("=== STRATIFIED BY OMEGA BAND ===")
    bands = [
        ("primitive  p<=127      (omega<4.65)",   lambda om: om < 4.65),
        ("middle     521-132049  (4.65<=omega<6.0)", lambda om: 4.65 <= om < 6.0),
        ("astron.    p>=216091   (omega>=6.0)",   lambda om: om >= 6.0),
    ]
    for label, cond in bands:
        subset = [(rp, p) for (p, _, rp, _) in loo_results if cond(fps_omega[p])]
        if subset:
            bmean = sum(rp for rp, _ in subset) / len(subset)
            byes  = sum(1 for rp, _ in subset if rp < 50.0)
            print(f"  {label}  n={len(subset):2d}  mean={bmean:5.1f}%  YES={byes}/{len(subset)}")
    print()

    # Digit-file stats
    print("=== DIGIT FILE ANALYSIS ===")
    print(f"  {'p':>10}  {'chi_sq':>7}  {'entropy':>8}  {'benford':>8}")
    print("  " + "-"*45)
    for p in sorted(digit_feats):
        df = digit_feats[p]
        print(f"  {p:>10d}  {df['chi_sq']:7.3f}  {df['entropy']:8.6f}  {df['benford_chisq']:8.3f}")
    print()

    # Feature attribution: which features have highest variance-per-std in LOO models
    print("=== FEATURE STABILITY (LOO model variance) ===")
    all_fps = list(fps.values())
    feat_std_full = fit_loo_model(all_fps)
    feat_names = [f"D{n}" for n in range(1, NUM_DIMS_C + 1)]
    # Compute coefficient of variation for each feature across known primes
    rows = []
    for f in range(FP_LEN):
        vals = [fps[p][f] for p in ALL_P]
        mean = sum(vals) / len(vals)
        std = feat_std_full[f]
        cv = std / max(abs(mean), 1e-10)
        rows.append((f, feat_names[f], mean, std, cv))
    rows.sort(key=lambda x: -x[4])
    print(f"  {'#':>2}  {'feature':18s}  {'mean':>8}  {'std':>8}  {'c.v.':>6}  status")
    print("  " + "-"*62)
    for rank, (f, name, mean, std, cv) in enumerate(rows[:20], 1):
        status = "DEAD (no variance)" if std < 1e-10 else ""
        print(f"  {rank:2d}  {name:18s}  {mean:8.4f}  {std:8.4f}  {cv:6.3f}  {status}")

    print()
    print(f"Total wall time: {time.time()-t0:.1f}s")
    return loo_results


if __name__ == '__main__':
    run_loo()
