"""phi_period_sweep.py
Test whether using beta = frac(omega / P) for P in {1, phi, phi^2, golden_angle, 2pi, pi}
gives better LOO than the current period-1 beta (P=1).

Spiral8plus insight: the natural lattice period might be phi^2 = 2.618 (one golden-angle
revolution) rather than 1.0. If so, primes sharing the same phi^2-octave would be better
calibration neighbors, improving the 2D kernel especially in the "plaid" (astro >= omega 7.3) band.
"""
import math, random

PHI = (1+math.sqrt(5))/2
LN_PHI = math.log(PHI)
LN2 = math.log(2)
GOLDEN_ANG = 2*math.pi/PHI   # 3.883 rad — golden angle
PHI2 = PHI*PHI               # 2.618

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
]
OMEGA_MAX = 8.0
PRM16 = [2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53]
KNOWN_SET = set(ALL_P)

def get_omega(p):
    inner = math.log(p * LN2) / LN_PHI
    return math.log(inner) / LN_PHI if inner > 0 else -999.0

def gen_fib(x):
    return (PHI**x - math.cos(math.pi*x)*PHI**(-x)) / math.sqrt(5)

def compute_Dn(n, r, omega_p):
    beta = omega_p - math.floor(omega_p)
    x = n + beta; k = (n+1)/8.0
    Fx = abs(gen_fib(x))
    On = 0.5 + 0.5*math.sin(omega_p + n*GOLDEN_ANG)
    coeff = math.sqrt(max(0.0, PHI*Fx*(2.0**n)*PRM16[n-1]*On*(PHI**(k*x))))
    delta = abs(math.cos(math.pi*beta*PHI))*math.log(PRM16[n-1])/(PHI**x)
    s = coeff*(abs(r)**k)
    pr = (1+delta)*math.cos(math.pi*beta)
    pi2 = (1+delta)*math.sin(math.pi*beta)
    return math.sqrt((s+pr)**2 + pi2**2)

def compute_fp(p):
    om = get_omega(p); r = om / OMEGA_MAX
    return [compute_Dn(n, r, om) for n in range(1, 17)]

print("  Precomputing fingerprints...", flush=True)
fps    = {p: compute_fp(p) for p in ALL_P}
omegas = {p: get_omega(p)  for p in ALL_P}

def frac_period(omega, P):
    """frac(omega / P) — phase within one period, in [0, 1)."""
    return (omega / P) - math.floor(omega / P)

def run_loo_period(P_beta, BW=0.5, BW_B=0.10):
    """LOO with beta = frac(omega / P_beta), 2D Gaussian kernel (omega, beta)."""
    random.seed(42)
    results = []
    for p_i in ALL_P:
        random.seed(p_i ^ 0xDEADBEEFCAFEBABE)
        om_i  = omegas[p_i]
        b_i   = frac_period(om_i, P_beta)

        train  = [p for p in ALL_P if p != p_i]
        tfps   = [fps[p]          for p in train]
        toms   = [omegas[p]       for p in train]
        tbetas = [frac_period(omegas[p], P_beta) for p in train]

        # 2D Gaussian kernel with circular wrap on beta
        wts = []
        for i in range(len(train)):
            db = tbetas[i] - b_i                       # no wrap — raw frac difference
            dw = (toms[i] - om_i) / BW
            wts.append(math.exp(-0.5*(dw*dw + (db/BW_B)**2)))
        tw = sum(wts)
        # Fallback 1: 1D omega-only
        if tw < 1e-15:
            wts = [math.exp(-0.5*((toms[i]-om_i)/BW)**2) for i in range(len(train))]
            tw  = sum(wts)
        # Fallback 2: uniform
        if tw < 1e-15:
            wts = [1.0]*len(train); tw = float(len(train))

        def mahal(target_fp):
            stats = []
            for f in range(16):
                wm = sum(wts[i]*tfps[i][f] for i in range(len(train))) / tw
                wv = sum(wts[i]*(tfps[i][f]-wm)**2 for i in range(len(train))) / tw
                ws = math.sqrt(wv) if wv > 1e-20 else 1.0
                stats.append((wm, ws))
            d2 = sum(((target_fp[f]-stats[f][0])/stats[f][1])**2 for f in range(16))
            return math.sqrt(d2/16)

        loo_d = mahal(fps[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
        span = hi - lo_odd
        nb = 0; nv = 0
        for _ in range(500):
            c = lo_odd + (random.randrange(span)//2)*2
            if c in KNOWN_SET: continue
            nv += 1
            if mahal(compute_fp(c)) < loo_d: nb += 1
        if nv > 0:
            results.append((p_i, 100.0*nb/nv))

    rps   = [r for _, r in results]
    mean  = sum(rps)/len(rps)
    nosig = sum(1 for r in rps if r >= 50)
    strong= sum(1 for r in rps if r <=  5)
    ok33  = sum(1 for r in rps if r <= 33)
    # Show per-prime for failures
    failures = [(p,r) for p,r in results if r >= 50]
    return mean, nosig, strong, ok33, failures

print()
print(f"  {'Period P':>12s}  {'mean%':>6s}  {'no_sig':>7s}  {'strong':>7s}  {'ok<=33':>7s}  {'BW_B':>6s}  failures")
print("  " + "-"*90)

test_cases = [
    (1.0,        "omega mod 1  (current)"),
    (PHI,        "omega mod phi"),
    (PHI2,       "omega mod phi^2"),
    (GOLDEN_ANG, "omega mod G (2pi/phi)"),
    (math.pi,    "omega mod pi"),
    (2*math.pi,  "omega mod 2pi"),
]

test_cases = [
    (1.0,        "omega mod 1  (current, BW_B=0.10)"),
    (PHI,        "omega mod phi, BW_B scaled"),
    (PHI2,       "omega mod phi^2, BW_B scaled"),
    (GOLDEN_ANG, "omega mod G (2pi/phi), BW_B scaled"),
    (math.pi,    "omega mod pi, BW_B scaled"),
]

# Also test a 3D kernel: (omega, beta_P1, beta_P_phi) — adds phi-period as a 3rd axis
# We do this by testing P=1.0 (existing best) as baseline then a combined kernel

for P, note in test_cases:
    # Scale BW_BETA inversely with P so the kernel covers the same omega width
    # BW_B * P = 0.10 * 1.0  =>  BW_B = 0.10 / P  (keeps omega-equivalent width at 0.10)
    bwb = 0.10 / P
    m, ns, st, ok, fails = run_loo_period(P, BW_B=bwb)
    fail_str = ", ".join(f"p={p}({r:.0f}%)" for p,r in fails) if fails else "none"
    marker = " <-- BEST" if ns == 0 else ""
    print(f"  {P:12.4f}  {m:6.1f}%  {ns:7d}  {st:7d}  {ok:7d}  {bwb:.4f}  {fail_str}{marker}")

print()
print("Done.")
