"""n41 — the dimension gate. Exact n40 mechanism on 2D (L=48) and 3D (L=12)
bipartite tori, no new ingredient. Per pre-registration n41_preregistration.md
(committed before this file existed). Gates: lattice-Coulomb form (<0.7%
frac RMS resid, 3x margin over screened/linear competitors) and attractive
sign, in BOTH dimensions, else headstone.
"""
import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spl

MU2 = 1e-6; F = 0.3; AL = 1.0; BE = 1.0

V2 = [(1,0),(2,1),(3,0),(3,2),(5,0),(4,3),(6,1),(5,4),(7,0),(7,2),
      (9,0),(9,4),(11,0),(11,4),(13,0),(13,2),(15,0),(15,4),(17,0),(17,2),
      (19,0),(19,4),(21,0),(21,2),(23,0)]
V3 = [(1,0,0),(1,1,1),(2,1,0),(3,0,0),(2,2,1),(3,1,1),(3,2,0),(4,1,0),
      (3,3,1),(4,2,1),(5,0,0),(4,3,0),(5,2,0),(4,4,1),(5,3,1),(6,1,0),
      (5,4,2),(6,3,0)]

def lattice(L, dim):
    """sites, edges (site pairs), unsigned incidence A (site x edge),
    oriented edge-difference D (edge x edge graph Laplacian factor)."""
    ns = L**dim
    def idx(c): return int(np.ravel_multi_index(np.mod(c, L), (L,)*dim))
    edges = []
    for s in range(ns):
        c = np.array(np.unravel_index(s, (L,)*dim))
        for ax in range(dim):
            c2 = c.copy(); c2[ax] += 1
            edges.append((s, idx(c2), ax))
    ne = len(edges)
    rows = []; cols = []
    for e, (i, j, ax) in enumerate(edges):
        rows += [i, j]; cols += [e, e]
    A = sp.csr_matrix((np.ones(2*ne), (rows, cols)), shape=(ns, ne))
    # D: difference of dk between an edge and its translate along each axis
    # (the |D dk|^2 smoothness term, direct generalization of the ring's D)
    epos = {}
    for e, (i, j, ax) in enumerate(edges): epos[(i, ax)] = e
    r2 = []; c2_ = []; v2 = []
    row = 0
    for e, (i, j, ax) in enumerate(edges):
        ci = np.array(np.unravel_index(i, (L,)*dim))
        for ax2 in range(dim):
            cn = ci.copy(); cn[ax2] += 1
            e2 = epos[(idx(cn), ax)]
            r2 += [row, row]; c2_ += [e, e2]; v2 += [1.0, -1.0]
            row += 1
    D = sp.csr_matrix((v2, (r2, c2_)), shape=(row, ne))
    return ns, ne, edges, A, D

def drainage_solver(L, dim):
    ns, ne, edges, A, D = lattice(L, dim)
    H = AL * sp.eye(ne) + BE * (D.T @ D)
    lu = spl.splu(H.tocsc())
    X = lu.solve(A.T.toarray())          # H^-1 A^T   (ne x ns)
    W = A @ X                            # A H^-1 A^T (ns x ns)
    def solve(cons):
        lam, *_ = np.linalg.lstsq(W, -cons, rcond=None)
        return X @ lam
    return ns, edges, A, solve

def Egs(L, dim, edges, dk):
    ns = L**dim
    K = np.zeros((ns, ns))
    k = 1.0 + dk
    for e, (i, j, ax) in enumerate(edges):
        K[i, i] += k[e]; K[j, j] += k[e]; K[i, j] -= k[e]; K[j, i] -= k[e]
    K[np.diag_indices(ns)] += MU2
    w2 = np.linalg.eigvalsh(K)
    return 0.5 * np.sum(np.sqrt(np.maximum(w2, 1e-15)))

def green_torus(L, dim):
    k = 2*np.pi*np.fft.fftfreq(L)
    ks = np.meshgrid(*([k]*dim), indexing='ij')
    lam = 2*dim - 2*sum(np.cos(x) for x in ks)
    lam_inv = np.where(lam > 1e-12, 1/np.where(lam > 1e-12, lam, 1), 0.0)
    return np.fft.ifftn(lam_inv).real

def affine_fit(y, x):
    Ad = np.vstack([np.ones_like(x), x]).T
    cf, *_ = np.linalg.lstsq(Ad, y, rcond=None)
    resid = y - Ad @ cf
    return cf[1], np.sqrt(np.mean(resid**2)) / (y.max() - y.min())

def run_dim(L, dim, vecs, fscan):
    print(f"===== {dim}D, L={L} =====")
    ns, edges, A, solve = drainage_solver(L, dim)
    def pair_cons(vec, f):
        c = np.zeros(ns)
        c[0] = f
        c[int(np.ravel_multi_index(np.mod(vec, L), (L,)*dim))] = f
        return c
    G = green_torus(L, dim)
    g = np.array([G[(0,)*dim] - G[tuple(np.mod(v, L))] for v in vecs])
    d = np.array([np.linalg.norm(v) for v in vecs])
    out = {}
    for f in fscan:
        V = []
        for vec in vecs:
            dk = solve(pair_cons(vec, f))
            V.append(Egs(L, dim, edges, dk))
        V = np.array(V)
        b, fr = affine_fit(V, g)
        # competitors
        fr_scr = min(affine_fit(V, np.exp(-d/lam))[1] for lam in np.linspace(0.5, 30, 120))
        fr_lin = affine_fit(V, d)[1]
        print(f"  F={f}: Coulomb fit b={b:+.4e} fracRMS={fr:.4%} | "
              f"best-screened={fr_scr:.4%} linear-d={fr_lin:.4%}")
        g1 = fr < 0.007 and fr_scr > 3*fr and fr_lin > 3*fr
        g2 = b > 0
        print(f"  G1 (form) {'PASS' if g1 else 'FAIL'}   G2 (sign) "
              f"{'PASS (attractive)' if g2 else 'FAIL (repulsive)'}")
        out[f] = dict(V=V.tolist(), b=b, fr=fr, fr_scr=fr_scr, fr_lin=fr_lin,
                      g1=bool(g1), g2=bool(g2))
    # C1: non-conserving control on first 8 vecs
    Vc = []
    for vec in vecs[:8]:
        c = pair_cons(vec, F)
        dk = np.zeros(len(edges))
        inc = {s: [] for s in np.nonzero(c)[0]}
        for e, (i, j, ax) in enumerate(edges):
            if i in inc: inc[i].append(e)
            if j in inc: inc[j].append(e)
        for s, es in inc.items():
            for e in es: dk[e] += -c[s]/(2*dim)
        Vc.append(Egs(L, dim, edges, dk))
    Vc = np.array(Vc)
    bc, frc = affine_fit(Vc, g[:8])
    frc_scr = min(affine_fit(Vc, np.exp(-d[:8]/lam))[1] for lam in np.linspace(0.5, 30, 120))
    c1_fails_coulomb = not (frc < 0.007 and frc_scr > 3*frc)
    print(f"  C1 control: Coulomb fracRMS={frc:.3%} screened={frc_scr:.3%} "
          f"-> control {'FAILS Coulomb (good: conservation is the cause)' if c1_fails_coulomb else 'PASSES Coulomb (EXPERIMENT INCONCLUSIVE)'}")
    # C2: carrier spectrum of a representative drainage field
    dk = solve(pair_cons(vecs[6], F))
    frac_stag = carrier_fraction(dk, edges, L, dim)
    print(f"  C2 carrier: fraction of |dk|^2 within pi/4 of the staggered point = {frac_stag:.1%}")
    # C3: same-parity infeasibility
    c = np.zeros(ns); c[0] = F
    c[int(np.ravel_multi_index((2,)+(0,)*(dim-1), (L,)*dim))] = F
    stag = np.array([(-1)**np.sum(np.unravel_index(s, (L,)*dim)) for s in range(ns)])
    print(f"  C3 neutrality: same-parity staggered charge = {np.dot(stag, c):+.2f} (nonzero = infeasible sector, as in 1D)")
    out['controls'] = dict(frc=frc, frc_scr=frc_scr, carrier=frac_stag)
    return out

def carrier_fraction(dk, edges, L, dim):
    tot = 0.0; stag = 0.0
    q = 2*np.pi*np.fft.fftfreq(L)
    qs = np.meshgrid(*([q]*dim), indexing='ij')
    dist = np.sqrt(sum((np.abs(np.abs(x) - np.pi))**2 for x in qs))
    mask = dist < np.pi/4
    for ax in range(dim):
        f = np.zeros((L,)*dim)
        for e, (i, j, a) in enumerate(edges):
            if a == ax: f[np.unravel_index(i, (L,)*dim)] = dk[e]
        P = np.abs(np.fft.fftn(f))**2
        tot += P.sum(); stag += P[mask].sum()
    return stag/tot

if __name__ == "__main__":
    r2 = run_dim(48, 2, V2, (0.3, 0.15))
    r3 = run_dim(12, 3, V3, (0.3, 0.15))
    verdict2 = r2[0.3]['g1'] and r2[0.3]['g2']
    verdict3 = r3[0.3]['g1'] and r3[0.3]['g2']
    print("=====")
    print(f"2D: {'ALIVE' if verdict2 else 'gate fired'}   3D: {'ALIVE' if verdict3 else 'gate fired'}")
    np.save("n41_results.npy", np.array([r2, r3], dtype=object), allow_pickle=True)