#!/usr/bin/env python3
"""
D1 v2 — f*: THE SELF-SIMILAR LOCKED FRACTION (precision build)
================================================================
Locked fraction f = 1 - E_N/MI for two balls (radius R, centers d = 2.5R)
in the massless 3D scalar vacuum (compact boson). L = 64 via FFT
covariances; symmetrized symplectic spectra (eigh on X^1/2 M X^1/2)
to kill the cancellation noise that produced negative MI at L = 32.
Extrapolate f(R) vs 1/R -> f*.
"""
import numpy as np, json, time, warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)

L = 64
n = np.arange(L)
KX, KY, KZ = np.meshgrid(2*np.pi*n/L, 2*np.pi*n/L, 2*np.pi*n/L, indexing='ij')
W = np.sqrt(4*(np.sin(KX/2)**2 + np.sin(KY/2)**2 + np.sin(KZ/2)**2))
invw = np.where(W > 1e-12, 1.0/(2.0*np.clip(W, 1e-12, None)), 0.0)
Xfun = np.real(np.fft.ifftn(invw))
Pfun = np.real(np.fft.ifftn(W/2.0))

def ball(c, R):
    s = []
    rng = range(-int(R)-1, int(R)+2)
    for dx in rng:
        for dy in rng:
            for dz in rng:
                if dx*dx+dy*dy+dz*dz <= R*R:
                    s.append(((c[0]+dx) % L, (c[1]+dy) % L, (c[2]+dz) % L))
    return s

def submat(F, sites):
    P = np.array(sites)
    D0 = (P[:,0][:,None]-P[:,0][None,:]) % L
    D1 = (P[:,1][:,None]-P[:,1][None,:]) % L
    D2 = (P[:,2][:,None]-P[:,2][None,:]) % L
    return F[D0, D1, D2]

def sqrtm_sym(A):
    e, V = np.linalg.eigh(A)
    return (V*np.sqrt(np.clip(e, 1e-300, None))) @ V.T

def S_and_sqrtX(sites):
    Xs, Ps = submat(Xfun, sites), submat(Pfun, sites)
    Xh = sqrtm_sym(Xs)
    nu2 = np.linalg.eigvalsh(Xh @ Ps @ Xh)
    nu = np.sqrt(np.clip(nu2, 0.25, None))
    up, dn = nu+.5, nu-.5
    S = float(np.sum(up*np.log(up)) -
              np.sum(np.where(dn > 1e-12, dn*np.log(np.clip(dn, 1e-300, None)), 0)))
    return S, Xh, Ps

def measure(R, d):
    c = L//2
    A = ball((c-(d+1)//2, c, c), R)
    B = ball((c+d//2, c, c), R)
    SA, _, _ = S_and_sqrtX(A)
    SB, _, _ = S_and_sqrtX(B)
    AB = A + B
    SAB, Xh, Ps = S_and_sqrtX(AB)
    mi = SA + SB - SAB
    Dg = np.ones(len(AB)); Dg[len(A):] = -1.0
    Pt = (Ps * Dg).T * Dg
    nu = np.sqrt(np.clip(np.linalg.eigvalsh(Xh @ Pt @ Xh), 1e-300, None))
    en = float(np.sum(np.where(nu < 0.5, -np.log(2*nu), 0.0)))
    return dict(R=R, d=d, nA=len(A), MI=mi, EN=en,
                f=1.0-en/mi if mi > 1e-12 else float('nan'))

if __name__ == '__main__':
    print(f"D1 v2 — f* on L={L}, massless compact boson, symmetrized spectra")
    print(f"{'R':>4} {'d':>4} {'|A|':>5} {'MI':>11} {'E_N':>11} {'f':>9}")
    out = []
    t0 = time.time()
    for R, d in ((2, 5), (3, 8), (4, 10), (5, 12), (6, 15), (7, 18)):
        r = measure(R, d)
        out.append(r)
        print(f"{R:4d} {d:4d} {r['nA']:5d} {r['MI']:11.7f} {r['EN']:11.7f} "
              f"{100*r['f']:8.2f}%  [{time.time()-t0:.0f}s]", flush=True)
    rows = [q for q in out if np.isfinite(q['f']) and q['R'] >= 3]
    RR = np.array([q['R'] for q in rows], float)
    ff = np.array([q['f'] for q in rows])
    fit = np.polyfit(1/RR, ff, 1)
    print(f"\nf(R) linear in 1/R (R>=3): continuum f* = {fit[1]:.4f} "
          f"(slope {fit[0]:+.3f}/R)")
    json.dump(out, open('/Users/antoine/agi/ledger/d1_results.json', 'w'), indent=1)
    print("-> d1_results.json")