"""n46 — the information-metric cost, 1D stage. Per n46_preregistration.md
(committed before this file). Cost = fidelity susceptibility of the Gaussian
vacuum; drainage = its minimum under the monogamy sum rule; sign = output.
"""
import numpy as np

N = 200; MU2 = 1e-6; F = 0.3

def K_of(dk):
    k = 1.0 + dk
    K = MU2 * np.eye(N)
    for e in range(N):
        i, j = e, (e+1) % N
        K[i,i] += k[e]; K[j,j] += k[e]; K[i,j] -= k[e]; K[j,i] -= k[e]
    return K

def Omega_of(dk):
    w2, V = np.linalg.eigh(K_of(dk))
    return V @ np.diag(np.sqrt(np.maximum(w2, 1e-30))) @ V.T

# --- exact metric G_ef = (1/8) Tr[Om^-1 dOm_e Om^-1 dOm_f] ---------------
w2, V = np.linalg.eigh(K_of(np.zeros(N)))
w = np.sqrt(np.maximum(w2, 1e-30))
U = np.zeros((N, N))                 # U[e,n] = v_n(i)-v_n(j), edge e=(i,i+1)
for e in range(N):
    U[e] = V[e] - V[(e+1) % N]
# Frechet derivative of sqrt: (dOm_e)_{nm} = U[e,n]U[e,m] * s_nm,
# s_nm = 1/(w_n + w_m)  (divided difference of sqrt on eigenvalues)
W1, W2 = np.meshgrid(w, w, indexing='ij')
s_dd = 1.0/(W1 + W2)
# G_ef = (1/8) sum_{nm} (dOm_e)_{nm} (dOm_f)_{mn} / (w_n w_m)
#      = (1/8) sum_{nm} U[e,n]U[e,m]U[f,n]U[f,m] * s_nm^2 / (w_n w_m)
kern = (s_dd**2) / (W1*W2)
B = np.einsum('en,em->enm', U, U).reshape(N, N*N)
G = 0.125 * (B * kern.reshape(1, N*N)) @ B.T

# S0: validate against the exact Gaussian overlap
def ln_overlap(dk):
    O1 = Omega_of(np.zeros(N)); O2 = Omega_of(dk)
    s1 = np.linalg.slogdet(O1)[1]; s2 = np.linalg.slogdet(O2)[1]
    sm = np.linalg.slogdet(0.5*(O1+O2))[1]
    return 0.25*s1 + 0.25*s2 - 0.5*sm
eps = 3e-3; maxerr = 0.0   # truncation measured to scale as eps^2 (exponent 2.02)
def d2(x, e):
    return -(ln_overlap(e*x) + ln_overlap(-e*x))/(e*e)
for b in (0, 7, 63, 128):
    for a in (b, (b+1) % N, (b+50) % N):
        x = np.zeros(N); x[a] += 1; x[b] += 1     # probe direction e_a+e_b
        f2 = (4*d2(x, eps) - d2(x, 2*eps))/3       # Richardson for eps^2 error
        quad = x @ G @ x   # -d2 ln|<0(0)|0(dk)>| = x^T g_FS x, Fubini-Study
        maxerr = max(maxerr, abs(f2 - quad)/np.abs(G).max())
ok = maxerr < 1e-5
print(f"S0 metric check: max |FD - quadratic| / max|G| = {maxerr:.2e} -> {'PASS' if ok else 'FAIL'}")

# S1: structure
d0 = np.mean(np.diag(G)); d1 = np.mean([G[e,(e+1)%N] for e in range(N)])
d2m = np.mean([G[e,(e+2)%N] for e in range(N)]); d5 = np.mean([G[e,(e+5)%N] for e in range(N)])
d20 = np.mean([G[e,(e+20)%N] for e in range(N)])
print(f"S1 structure: G diag {d0:+.5f}, nn {d1:+.5f}, 2nd {d2m:+.5f}, 5th {d5:+.5f}, 20th {d20:+.5f}")
ev = np.linalg.eigvalsh(G)
print(f"   spectrum: min {ev.min():+.3e} max {ev.max():+.3e} (PSD expected)")

# S2: the Gauss-sector direction
s = np.array([(-1)**e for e in range(N)], float)
sGs = s @ G @ s / N
print(f"S2 kernel check: s^T G s / N = {sGs:+.6e} " +
      ("(metric BLIND to Gauss sector — declared distinct outcome)" if abs(sGs) < 1e-10*np.abs(G).max() else "(metric prices the Gauss sector)"))

# --- pair interaction --------------------------------------------------
A = np.zeros((N, N))
for i in range(N): A[i,(i-1)%N] = 1; A[i,i] += 1

ds = [11,21,31,41,51,61,71,81]
g_ref = np.array([d*(N-d)/(2*N) for d in ds])

def V_of(d, f):
    c = np.zeros(N); c[40] = f; c[40+d] = f
    dk_p, *_ = np.linalg.lstsq(A, -c, rcond=None)
    t = -(s @ G @ dk_p)/(s @ G @ s)
    dk = dk_p + t*s
    return 0.5 * dk @ G @ dk

for f in (0.3, 0.15):
    Vv = np.array([V_of(d, f) for d in ds])
    Ad = np.vstack([np.ones_like(g_ref), g_ref]).T
    cf, *_ = np.linalg.lstsq(Ad, Vv, rcond=None)
    fr = np.sqrt(np.mean((Vv-Ad@cf)**2))/np.ptp(Vv)
    print(f"F={f}: V(d) = {list(np.round(Vv,8))}")
    print(f"   slope b = {cf[1]:+.6e} ({'ATTRACTIVE' if cf[1]>0 else 'REPULSIVE'}), Green-fit fracRMS = {fr:.3%}")
    if f == 0.3: b03 = cf[1]
    else: print(f"   F-scaling b(0.3)/b(0.15) = {b03/cf[1]:.4f} (4.00 expected)")
np.save("n46_results.npy", dict(G_diag=d0, G_nn=d1, sGs=sGs, ds=ds), allow_pickle=True)