"""
n20 — The 2D rung of the Gauss ladder. Does the record-shadow force follow
the dimensional law of gravity?

The mechanism measured in n18 (1D) predicts F ~ P/(sphere area): the force
on a monitored object tracks the partner's record-flux intensity, which
dilutes geometrically and is screened only at the medium's transparency
length. The dimensional ladder:
    1D: constant force            -- MEASURED (n18, +-5% over d=24-64)
    2D: F ~ 1/d                   -- THIS SCRIPT
    3D: F ~ 1/d^2 (Newton)        -- follows by the same argument if 2D holds

PRE-REGISTERED GATES
 V0  Undriven NESS == exact ground state (built into the solver: matched
     uniform bath; verified numerically anyway).
 V1  Force observable calibration: for STATIC coherent bumps, the
     momentum-balance force must match the exact ground-state energy
     gradient -dE0/dd (central difference) to ~10%.
 G0  g = 0 baseline: pump repulsion, small (the 2D analog of n13).
 MAIN dF(d) = F(g) - F(0) fitted to C * d^-p * exp(-d/lambda), lambda free:
     SUCCESS if p in [0.7, 1.3]  (2D Gauss law, allowing lattice + box bias)
     FAILURE if p < 0.4 (no dilution: artifact) or p > 1.7 or attraction absent.

Method: exact Gaussian NESS in the eigenbasis of K. Uniform damping eta on
every site with the bath matched to the CURRENT K's ground state makes the
drift A = A_H - (eta/2) I block-diagonal per mode pair; each Lyapunov solve
is then a closed-form 2x2 block:
    with h = eta/2, a = w_k^2, b = w_l^2, source s only in (p,p):
    x11 = s / (2h(a+b) + (b-a)^2/(4h) + 4h^3)
    x12 = x11 (h - (b-a)/(4h));  x21 = x11 (h + (b-a)/(4h))
    x22 = x11 ((a+b)/2 + 2h^2)
Grain noise: n13's self-consistent record noise (measurement of (w.x),
diffusion on (w.p)), grains = sum-zero Laplacian-of-Gaussian profiles
(isotropic, IR-safe). Hartree dressing as in n18: mu2 = MU2 + 3 g dx2,
dx2 = noise-induced <x^2> excess (>= 0). Force: momentum balance on a box
around each grain, ABBA-subtracted; calibrated by gate V1.

Runtime: ~1-2 h for the full scan (numpy/BLAS, N = 2048 sites).
"""

import time
import numpy as np

# macOS Accelerate raises spurious FP flags on healthy matmuls (verified n18)
np.seterr(divide='ignore', over='ignore', invalid='ignore')

LX, LY = 64, 32
N      = LX * LY
MU2    = 4e-4
ETA    = 0.05
GAMMA  = 0.05
GBIG   = 1.0
SMEAR  = 1.5
BOX    = 3          # half-width of the momentum-balance box
ALPHA  = 0.6        # Hartree damping
TOL    = 1e-8
MAXIT  = 40
CY     = LY // 2
T0     = time.time()

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

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

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

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

def weight(cx, cy):
    """Sum-zero isotropic grain: Laplacian of a Gaussian."""
    w = np.zeros(N)
    for ix in range(LX):
        for iy in range(LY):
            r2 = (ix - cx) ** 2 + (iy - cy) ** 2
            if r2 <= (4 * SMEAR) ** 2:
                s2 = SMEAR ** 2
                w[idx(ix, iy)] = (r2 / s2 - 2.0) * np.exp(-0.5 * r2 / s2)
    return w / np.linalg.norm(w)

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

def gauss_bump_dx(cx, cy, h):
    """d(bump)/d(x0) analytically: force kernel for Hellmann-Feynman."""
    b = np.zeros(N)
    for ix in range(LX):
        for iy in range(LY):
            r2 = (ix - cx) ** 2 + (iy - cy) ** 2
            b[idx(ix, iy)] = h * ((ix - cx) / 4.0) * np.exp(-0.5 * r2 / 4.0)
    return b

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

class Ness:
    """Exact Gaussian NESS with uniform matched bath + n13 record noise."""
    def __init__(self, mu2_vec, noisy):
        # noisy = ((weight_vector, gamma), ...)
        self.K = K_matrix(mu2_vec)
        self.w2, self.U = np.linalg.eigh(self.K)
        self.w = np.sqrt(np.maximum(self.w2, 1e-14))
        h = ETA / 2
        a = self.w2[:, None]
        b = self.w2[None, :]
        den = 2 * h * (a + b) + (b - a) ** 2 / (4 * h) + 4 * h ** 3
        self.f11 = 1.0 / den
        self.f12 = self.f11 * (h - (b - a) / (4 * h))
        self.f21 = self.f11 * (h + (b - a) / (4 * h))
        self.f22 = self.f11 * ((a + b) / 2 + 2 * h * h)
        self.noisy = noisy
        # self-consistent record noise (n13): s_i = <(w_i . x)^2>
        ng = len(noisy)
        if ng:
            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.Xpp = sum(gam[j] * self.svals[j] * self.f22 *
                           np.outer(wts[j], wts[j]) for j in range(ng))
            self.Xxp = sum(gam[j] * self.svals[j] * self.f12 *
                           np.outer(wts[j], wts[j]) for j in range(ng))
        else:
            self.svals = np.zeros(0)
            self.Xxx = self.Xpp = self.Xxp = None

    def site_xx(self):
        """Full site-basis <x x> matrix (ground + noise excess)."""
        gs = (self.U / (2 * self.w)) @ self.U.T
        if self.Xxx is None:
            return gs
        return gs + self.U @ self.Xxx @ self.U.T

    def site_pp(self):
        gs = (self.U * (self.w / 2)) @ self.U.T
        if self.Xpp is None:
            return gs
        return gs + self.U @ self.Xpp @ self.U.T

    def site_xp(self):
        if self.Xxp is None:
            return np.zeros((N, N))
        return self.U @ self.Xxp @ self.U.T

    def dx2(self):
        """Noise-induced <x^2> excess per site (>= 0)."""
        if self.Xxx is None:
            return np.zeros(N)
        return ((self.U @ self.Xxx) * self.U).sum(axis=1)

# --------------------------------------------------------------- Hartree ----

def hartree(noisy_sites, g, label, base=None):
    """noisy_sites = ((cx, cy, gamma), ...); base = frozen stiffness profile
    (e.g. the passive test bump) -> converged Ness + mu2."""
    noisy = tuple((weight(cx, cy), gv) for cx, cy, gv in noisy_sites)
    base_v = np.full(N, MU2) if base is None else base
    mu2 = base_v.copy()
    for it in range(MAXIT):
        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
        if it % 5 == 0 or delta < TOL:
            say(f"      [{label}] it {it:2d}  d(mu2) {delta:.2e}")
        if delta < TOL:
            break
    return Ness(mu2, noisy), mu2

# ------------------------------------------------------- force observable ---

def xx_diag(ns):
    """Site-basis <x_i^2>: ground diagonal + noise excess. No GEMM needed."""
    gs = ((ns.U / (2 * ns.w)) * ns.U).sum(axis=1)
    return gs + ns.dx2()

def hf_force(ns, dBdx):
    """Hellmann-Feynman x-force on the test bump: F = -<dH/dx0>
    = -(1/2) sum_i (dB/dx0)_i <x_i^2>.  (H contains B_i x_i^2 / 2;
    dB/dx0 = -dB/dx at fixed site, handled by the analytic kernel.)
    Exact in any state, noisy or not. Sign calibrated by gate V1."""
    return -0.5 * float(dBdx @ xx_diag(ns))

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

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

    # -- V0: vacuum == ground state ----------------------------------------
    ns0 = Ness(np.full(N, MU2), ())
    say(f"V0: vacuum dx2 max = {np.max(np.abs(ns0.dx2())):.2e} (exact 0 by construction)")

    # -- V1: HF force calibration on static bumps ---------------------------
    # Two static bumps: HF force on B must match the exact ground-state
    # energy gradient dE0/dd (both computed independently).
    say("V1: static-bump HF force vs exact dE0/dd")
    h_b = 0.5
    d0 = 12
    ax, bx = LX // 2 - d0 // 2, LX // 2 + d0 // 2
    def E0_of(bpos):
        mu2 = np.full(N, MU2) + gauss_bump(ax, CY, h_b) + gauss_bump(bpos, CY, 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   # -dE0/dx_B
    mu2 = np.full(N, MU2) + gauss_bump(ax, CY, h_b) + gauss_bump(bx, CY, h_b)
    ns = Ness(mu2, ())
    F_hf = hf_force(ns, gauss_bump_dx(bx, CY, h_b))
    say(f"V1: HF force = {F_hf:+.5e}   exact -dE0/dx_B = {F_exact:+.5e}   "
        f"ratio = {F_hf / F_exact:.4f}   (attraction = negative x-force on B)")

    # -- G0 + MAIN: source -> test-mass distance scan ------------------------
    # Newton's own geometry: a monitored grain (the source) at ax, a PASSIVE
    # stiffness bump (the test mass) at bx = ax + d. Force on the test mass
    # by HF, background-subtracted with the bump-alone config.
    say("MAIN: force on a passive test mass vs distance from a monitored source")
    h_t = 0.2
    results = []
    for d in (8, 12, 16, 20, 26):
        ax = LX // 2 - d // 2
        bx = ax + d
        bump = gauss_bump(bx, CY, h_t)
        dBdx = gauss_bump_dx(bx, CY, h_t)
        base = np.full(N, MU2) + bump
        # bump alone (self-force background; ~0 by symmetry, boundary bias)
        F_bg = hf_force(Ness(base, ()), dBdx)
        row = {}
        for g in (0.0, GBIG):
            ns, _ = hartree(((ax, CY, 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, row[0.0], row[GBIG], dF))
        say(f"MAIN: d = {d:<3} dF(g-induced) = {dF:+.5e}   [attraction = negative]")

    # -- fit ------------------------------------------------------------------
    say("FIT: dF = C * d^-p * exp(-d/lambda)")
    ds = np.array([r[0] for r in results], float)
    dFs = -np.array([r[3] for r in results])   # attraction = negative x-force on B
    if np.all(dFs > 0):
        # crude 2-parameter grid then refine p by least squares on log
        best = None
        for lam in np.arange(10, 200, 2.5):
            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"(2D Gauss law: p = 1; gate: 0.7 <= p <= 1.3)")
    else:
        say("FIT: dF not uniformly attractive — gate FAILED as registered")
    say("done.")
