#!/usr/bin/env python3
"""
N8 — THE LOCKED FRACTION IN 3+1D
==================================
First rung toward the covariant locked fraction (the Omega_Lambda
target of Folio XI). The 1D toy (N5) gave f = 77.5% at contact;
dimension changes entanglement structure qualitatively, so the 3D
number is the first one worth comparing to Omega_Lambda = 0.68.

PROTOCOL: two balls of radius R in the quasi-massless 3D scalar
lattice vacuum (L^3, PBC, m = 0.01), center separation d.
    f(d) = 1 - E_N(A:B)/I(A:B)   (locked fraction of the account)
Sweep d at fixed R; also vary R at contact-like separations.

CAVEATS (named): one scalar, not photon+graviton; lattice cutoff;
locked fraction defined via E_N/MI ratio (a proxy — distillable
entanglement <= E_N, so f is a LOWER bound on the locked share).
"""

import numpy as np
import json
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)

L, m = 16, 0.01
N = L ** 3

idx = lambda x, y, z: (x % L) * L * L + (y % L) * L + (z % L)

def covariances():
    K = np.zeros((N, N))
    for x in range(L):
        for y in range(L):
            for z in range(L):
                i = idx(x, y, z)
                K[i, i] = m * m + 6.0
                for d in ((1, 0, 0), (0, 1, 0), (0, 0, 1)):
                    j = idx(x + d[0], y + d[1], z + d[2])
                    K[i, j] -= 1.0
                    K[j, i] -= 1.0
    w2, U = np.linalg.eigh(K)
    w = np.sqrt(np.clip(w2, 1e-30, None))
    return (U * (0.5 / w)) @ U.T, (U * (0.5 * w)) @ U.T

def ball(cx, R):
    out = []
    for x in range(L):
        for y in range(L):
            for z in range(L):
                dx = min((x - cx[0]) % L, (cx[0] - x) % L)
                dy = min((y - cx[1]) % L, (cx[1] - y) % L)
                dz = min((z - cx[2]) % L, (cx[2] - z) % L)
                if dx * dx + dy * dy + dz * dz <= R * R:
                    out.append(idx(x, y, z))
    return out

def entropy(X, P, R_):
    ix = np.ix_(R_, R_)
    nu = np.sqrt(np.clip(np.linalg.eigvals(X[ix] @ P[ix]).real, 0.25, None))
    up, dn = nu + .5, nu - .5
    return float(np.sum(up * np.log(up)) -
                 np.sum(np.where(dn > 1e-12,
                                 dn * np.log(np.clip(dn, 1e-300, None)), 0)))

def MI(X, P, A, B):
    return entropy(X, P, A) + entropy(X, P, B) - entropy(X, P, A + B)

def logneg(X, P, A, B):
    R_ = A + B
    ix = np.ix_(R_, R_)
    D = np.ones(len(R_)); D[len(A):] = -1.0
    Pt = (P[ix] * D).T * D
    nu = np.sqrt(np.clip(np.linalg.eigvals(X[ix] @ Pt).real, 1e-30, None))
    return float(np.sum(np.where(nu < 0.5, -np.log(2 * nu), 0.0)))


if __name__ == '__main__':
    print(f"N8 — locked fraction in 3+1D. lattice {L}^3, m={m}")
    X, P = covariances()
    c = L // 2
    print(f"\n{'R':>4} {'d':>4} {'|A|':>5} {'MI(A:B)':>10} {'E_N(A:B)':>10} {'locked f':>9}")
    rows = []
    for R in (1.5, 2.0, 2.5):
        for d in (3, 4, 5, 6, 8):
            A = ball((c - (d + 1) // 2, c, c), R)   # centers exactly d apart
            B = ball((c + d // 2, c, c), R)
            if set(A) & set(B):
                continue
            mi, en = MI(X, P, A, B), logneg(X, P, A, B)
            f = 1.0 - en / mi if mi > 1e-12 else float('nan')
            rows.append(dict(R=R, d=d, nA=len(A), MI=mi, EN=en, f=f))
            print(f"{R:4.1f} {d:4d} {len(A):5d} {mi:10.5f} {en:10.6f} {100*f:8.1f}%")
    json.dump(dict(L=L, m=m, rows=rows),
              open('/Users/antoine/agi/ledger/n8_results.json', 'w'), indent=1)
    print("\n-> n8_results.json")
    print("1D contact value was 77.5%; Omega_Lambda observed = 68%.")