"""
n21 — The 3D rung: Newton's inverse square, or not.

The Gauss ladder so far, all in Newton's source-and-test-mass geometry:
    1D: F ~ const        MEASURED (n18, +-5%)
    2D: F ~ 1/d          MEASURED (n20, p = 1.085, gate 0.7-1.3)
    3D: F ~ 1/d^2        THIS SCRIPT
A monitored grain (the source) and a passive stiffness bump (the test
mass) in a 3D lattice scalar vacuum with uniform matched damping; force on
the test mass by Hellmann-Feynman; Hartree dressing exactly as in n18/n20.

PRE-REGISTERED GATES
 V0   undriven NESS == ground state (exact by construction).
 V1   HF force on a static bump vs exact dE0/dx_B: agree to ~10%.
 G0   g = 0: the bare record-flux pushes the test mass (Maxwell).
 MAIN dF(d) = F(g=1) - F(0) attractive at every d, fitted to
      C * d^-p * exp(-d/lambda):  SUCCESS if p in [1.6, 2.4]
      (the 3D Gauss law is p = 2). FAILURE reported as such.

Geometry: 28 x 10 x 10 = 2800 sites; grains/bumps on the long axis.
Runtime: ~1.5-2 h (eigh(2800) per Hartree iteration; BLAS).
"""

import time
import numpy as np

np.seterr(divide='ignore', over='ignore', invalid='ignore')

LX, LY, LZ = 30, 26, 26   # v3 — the definitive box (20 280 sites).
# v1 (28x10x10): tube crossover at d~8, p = 0.312 — rejected.
# v2 (30x16x16): clean local exponents 1.79-2.33 at d = 5-8 (bracketing the
# Newtonian 2), crossover at d~9, global p = 1.45 — rejected by geometry.
# v3 uses a 3-step undamped Hartree, validated against the fully converged
# solver on the 2D reference: force ratio 0.998 (2 steps) / 1.000 (3 steps).
N      = LX * LY * LZ
MU2    = 4e-4
ETA    = 0.05
GAMMA  = 0.05
GBIG   = 1.0
SMEAR  = 1.0
ALPHA  = 1.0    # undamped (Jacobi), per the validated fast scheme
TOL    = 1e-8
MAXIT  = 3      # 3-step Jacobi Hartree (validated on 2D: ratio 1.000)
CY, CZ = LY // 2, LZ // 2
T0     = time.time()

def say(m):
    print(f"[{time.time()-T0:7.1f}s] {m}", flush=True)

def idx(ix, iy, iz):
    return (ix * LY + iy) * LZ + iz

# ---------------------------------------------------------------- lattice ---

def K_matrix(mu2_vec):
    K = np.zeros((N, N))
    for ix in range(LX):
        for iy in range(LY):
            for iz in range(LZ):
                i = idx(ix, iy, iz)
                nb = 0
                for dx, dy, dz in ((1,0,0), (-1,0,0), (0,1,0), (0,-1,0),
                                   (0,0,1), (0,0,-1)):
                    jx, jy, jz = ix + dx, iy + dy, iz + dz
                    if 0 <= jx < LX and 0 <= jy < LY and 0 <= jz < LZ:
                        K[i, idx(jx, jy, jz)] = -1.0
                        nb += 1
                K[i, i] = nb + mu2_vec[i]
    return K

def weight(cx):
    """Sum-zero isotropic grain (Laplacian of Gaussian), on the long axis."""
    w = np.zeros(N)
    s2 = SMEAR ** 2
    for ix in range(LX):
        for iy in range(LY):
            for iz in range(LZ):
                r2 = (ix - cx) ** 2 + (iy - CY) ** 2 + (iz - CZ) ** 2
                if r2 <= (4 * SMEAR) ** 2:
                    w[idx(ix, iy, iz)] = (r2 / s2 - 3.0) * np.exp(-0.5 * r2 / s2)
    return w / np.linalg.norm(w)

BW = 1.6   # test-bump width (1.2: V1 14%; 1.4: 12%; lattice O(a^2) error demands 1.6)

def gauss_bump(cx, h):
    b = np.zeros(N)
    for ix in range(LX):
        for iy in range(LY):
            for iz in range(LZ):
                r2 = (ix - cx) ** 2 + (iy - CY) ** 2 + (iz - CZ) ** 2
                b[idx(ix, iy, iz)] = h * np.exp(-0.5 * r2 / BW ** 2)
    return b

def gauss_bump_dx(cx, h):
    b = np.zeros(N)
    for ix in range(LX):
        for iy in range(LY):
            for iz in range(LZ):
                r2 = (ix - cx) ** 2 + (iy - CY) ** 2 + (iz - CZ) ** 2
                b[idx(ix, iy, iz)] = h * ((ix - cx) / BW ** 2) \
                    * np.exp(-0.5 * r2 / BW ** 2)
    return b

# ------------------------------------------------- fast NESS (mode basis) ---

class Ness:
    def __init__(self, mu2_vec, noisy):
        K = K_matrix(mu2_vec)
        self.w2, self.U = np.linalg.eigh(K)
        del K
        self.w = np.sqrt(np.maximum(self.w2, 1e-14))
        ng = len(noisy)
        if ng:
            h = ETA / 2
            a = self.w2[:, None]
            b = self.w2[None, :]
            self.f11 = 1.0 / (2 * h * (a + b) + (b - a) ** 2 / (4 * h) + 4 * h ** 3)
            wts = [self.U.T @ wv for wv, _ in noisy]
            gam = [gv for _, gv in noisy]
            bvec = np.array([np.sum(wt ** 2 / (2 * self.w)) for wt in wts])
            M = np.zeros((ng, ng))
            for i in range(ng):
                for j in range(ng):
                    Xxx = self.f11 * np.outer(wts[j], wts[j])
                    M[i, j] = gam[j] * (wts[i] @ Xxx @ wts[i])
            ev = np.linalg.eigvals(M)
            if np.max(ev.real) >= 1.0:
                raise RuntimeError(f"parametric instability ({np.max(ev.real):.3f})")
            self.svals = np.linalg.solve(np.eye(ng) - M, bvec)
            self.Xxx = sum(gam[j] * self.svals[j] * self.f11 *
                           np.outer(wts[j], wts[j]) for j in range(ng))
            self.f11 = None            # free 3.3 GB — only needed to build Xxx
        else:
            self.f11 = None
            self.svals = np.zeros(0)
            self.Xxx = None

    def dx2(self):
        if self.Xxx is None:
            return np.zeros(N)
        return ((self.U @ self.Xxx) * self.U).sum(axis=1)

def xx_diag(ns):
    gs = ((ns.U / (2 * ns.w)) * ns.U).sum(axis=1)
    return gs + ns.dx2()

def hf_force(ns, dBdx):
    return -0.5 * float(dBdx @ xx_diag(ns))

def hartree(noisy_sites, g, label, base=None):
    noisy = tuple((weight(cx), gv) for cx, gv in noisy_sites)
    base_v = np.full(N, MU2) if base is None else base
    mu2 = base_v.copy()
    ns = None
    for it in range(MAXIT):
        ns = None                      # drop the old NESS before building the new
        ns = Ness(mu2, noisy)
        tgt = base_v + 3.0 * g * ns.dx2()
        new = (1 - ALPHA) * mu2 + ALPHA * tgt
        delta = np.max(np.abs(new - mu2))
        mu2 = new
        say(f"      [{label}] it {it:2d}  d(mu2) {delta:.2e}")
        if delta < TOL:
            break
    return Ness(mu2, noisy), mu2

# ------------------------------------------------------------------- main ---

if __name__ == "__main__":
    say(f"n21 — 3D Gauss rung. {LX}x{LY}x{LZ}, eta={ETA}, gamma={GAMMA}, g={GBIG}")

    ns0 = Ness(np.full(N, MU2), ())
    say(f"V0: vacuum dx2 max = {np.max(np.abs(ns0.dx2())):.2e}")

    say("V1: static-bump HF force vs exact dE0/dx_B")
    h_b = 0.5
    d0 = 6
    ax, bx = LX // 2 - d0 // 2, LX // 2 + d0 // 2
    def E0_of(bpos):
        mu2 = np.full(N, MU2) + gauss_bump(ax, h_b) + gauss_bump(bpos, h_b)
        return 0.5 * np.sum(np.sqrt(np.linalg.eigvalsh(K_matrix(mu2))))
    F_exact = -(E0_of(bx + 1) - E0_of(bx - 1)) / 2.0
    mu2 = np.full(N, MU2) + gauss_bump(ax, h_b) + gauss_bump(bx, h_b)
    F_hf = hf_force(Ness(mu2, ()), gauss_bump_dx(bx, h_b))
    say(f"V1: HF = {F_hf:+.5e}   exact = {F_exact:+.5e}   ratio = {F_hf/F_exact:.4f}")

    say("MAIN: force on passive test mass vs distance from monitored source")
    h_t = 0.2
    results = []
    for d in (6, 8, 10, 12):
        ax = LX // 2 - d // 2
        bx = ax + d
        bump = gauss_bump(bx, h_t)
        dBdx = gauss_bump_dx(bx, h_t)
        base = np.full(N, MU2) + bump
        ns_bg = Ness(base, ())
        F_bg = hf_force(ns_bg, dBdx)
        ns_bg = None                   # free before the Hartree sequence
        row = {}
        for g in (0.0, GBIG):
            ns, _ = hartree(((ax, GAMMA),), g, f"d{d} g{g}", base=base)
            row[g] = hf_force(ns, dBdx) - F_bg
            say(f"MAIN: d = {d:<3} g = {g:<4} F_test = {row[g]:+.5e}")
        dF = row[GBIG] - row[0.0]
        results.append((d, dF))
        say(f"MAIN: d = {d:<3} dF(g-induced) = {dF:+.5e}   [attraction = negative]")

    say("FIT: dF = C * d^-p * exp(-d/lambda)")
    ds = np.array([r[0] for r in results], float)
    dFs = -np.array([r[1] for r in results])
    if np.all(dFs > 0):
        best = None
        for lam in np.arange(8, 300, 2.0):
            y = np.log(dFs) + ds / lam
            A = np.vstack([np.ones_like(ds), -np.log(ds)]).T
            coef, res, *_ = np.linalg.lstsq(A, y, rcond=None)
            r = float(res[0]) if len(res) else 0.0
            if best is None or r < best[0]:
                best = (r, lam, coef[1])
        say(f"FIT: p = {best[2]:.3f}  lambda = {best[1]:.1f}  "
            f"(3D Gauss law: p = 2; gate: 1.6 <= p <= 2.4)")
        # local exponents for the record
        for i in range(len(ds) - 1):
            pl = np.log(dFs[i] / dFs[i + 1]) / np.log(ds[i + 1] / ds[i])
            say(f"FIT: local exponent {ds[i]:.0f}->{ds[i+1]:.0f}: {pl:.2f}")
    else:
        say("FIT: dF not uniformly attractive — gate FAILED as registered")
    say("done.")
