"""n46b — the infrared gate. All-Fourier pipeline for the information-metric
drainage: G block-diagonal in q, W's symbol, V(d) by FFT. Per
n46b_preregistration.md (committed before this file).
"""
import numpy as np
import itertools

MU2 = 1e-6

def wdisp(qs):
    """w_k = sqrt(sum_a 2-2cos k_a + mu^2); qs: array (...,dim)"""
    return np.sqrt(np.sum(2 - 2*np.cos(qs), axis=-1) + MU2)

def Ghat_W(L, dim):
    """Return Whate(q) on the full q-grid, plus the q-grid."""
    k1 = 2*np.pi*np.fft.fftfreq(L)
    grids = np.meshgrid(*([k1]*dim), indexing='ij')
    kvec = np.stack(grids, axis=-1)             # (L,..,L,dim)
    w = wdisp(kvec)                              # (L,..,L)
    shape = w.shape
    NS = L**dim
    kflat = kvec.reshape(NS, dim)
    wflat = w.reshape(NS)
    # Ghat_ab(q) = (1/(8*NS)) sum_k kern(w_k, w_{q-k}) u_a(k)u_a(q-k) conj(u_b(k)u_b(q-k))
    # u_a(k) = 1 - exp(-i k_a)   (site-difference along +a)
    # kern(w1,w2) = 1/((w1+w2)^2 w1 w2)
    Wq = np.zeros(NS)
    idx = np.arange(NS)
    coords = np.array(np.unravel_index(idx, (L,)*dim)).T   # integer momenta
    for iq in range(NS):
        qc = coords[iq]
        # momentum transfer q = k' - k  ->  k' = k + q (mod L)
        kp = np.mod(coords + qc[None, :], L)
        ikp = np.ravel_multi_index(kp.T, (L,)*dim)
        w1 = wflat; w2 = wflat[ikp]
        kern = 1.0/((w1+w2)**2 * w1 * w2)
        u1 = 1 - np.exp(1j*kflat)                 # u_a(k)
        u2 = 1 - np.exp(1j*kflat[ikp])            # u_a(k+q)
        X = np.conj(u1) * u2                      # (NS,dim)
        M = np.einsum('n,na,nb->ab', kern, X.conj(), X) / (8*NS)
        # Ahat_a(q): (A dk)(r) = sum_a dk_a(r-e_a) + dk_a(r) -> 1 + exp(-i q_a)
        qphys = 2*np.pi*qc/L
        Av = 1 + np.exp(-1j*qphys)                # (dim,)
        # W(q) = Av^dagger Ghat^-1 Av  -- guard singular Ghat (q=0: u=0)
        try:
            x = np.linalg.solve(M, Av)
            Wq[iq] = np.real(np.conj(Av) @ x)
        except np.linalg.LinAlgError:
            Wq[iq] = np.inf
    return Wq.reshape(shape)

def V_of_pairs(L, dim, vecs, f):
    Wq = Ghat_W(L, dim)
    invW = np.where(np.isfinite(Wq) & (np.abs(Wq) > 1e-9), 1/np.where(np.abs(Wq)>1e-9, Wq, 1), 0.0)
    ker = np.fft.ifftn(invW).real               # W^+ kernel (zero modes dropped)
    out = []
    for v in vecs:
        # c = f (delta_0 + delta_v): c^T W^+ c = f^2 (2 ker(0) + 2 ker(v))
        out.append(0.5 * f*f * (2*ker[(0,)*dim] + 2*ker[tuple(np.mod(v, L))]))
    return np.array(out), Wq

# ---- V0: reproduce n46 2D direct build -----------------------------------
V2 = [(1,0),(2,1),(3,2),(5,0),(4,3),(6,1),(7,2),(7,4),(9,0),(9,4),(11,0)]
Vf, Wq24 = V_of_pairs(24, 2, V2, 0.3)
# reference values from the direct build (n46_infometric_2d.py, this session)
exec(open('n46_infometric_2d.py').read().split('for f in')[0])   # rebuild direct G, W (prints)
Vd = []
for v in V2:
    c = np.zeros(NS); c[site((0,0))] = 0.3; c[site(v)] = 0.3
    lam, *_ = np.linalg.lstsq(W, -c, rcond=None)
    dk = Ginv_At @ lam
    Vd.append(0.5*dk@G@dk)
Vd = np.array(Vd)
v0err = np.abs(Vf - Vd).max()/np.abs(Vd).max()
print(f"V0 cross-check (Fourier vs direct, 2D L=24): max rel = {v0err:.2e} -> {'PASS' if v0err < 1e-8 else 'FAIL'}")

# ---- IR gates -------------------------------------------------------------
def ir_report(L, dim, Wq=None):
    if Wq is None: Wq = Ghat_W(L, dim)
    stagg = (L//2,)*dim
    z = abs(Wq[stagg])/np.abs(Wq[np.isfinite(Wq)]).max()
    S = []
    for dq in (1, 2):
        i = list(stagg); i[0] = (L//2 + dq) % L
        qp = dq*2*np.pi/L
        S.append(Wq[tuple(i)]/qp**2)
    print(f"{dim}D L={L}: |W(stagg)|/max = {z:.2e}; S(q1)={S[0]:.4f} S(q2)={S[1]:.4f} ratio-1 = {abs(S[0]/S[1]-1):.4f}")
    return z, S

print("IR1/IR2:")
z24, S24 = ir_report(24, 2, Wq24)
z48, S48 = ir_report(48, 2)
z96, S96 = ir_report(96, 2)
print(f"2D stiffness convergence: S_min(96)/S_min(48) - 1 = {abs(S96[0]/S48[0]-1):.4f} (gate < 0.02)")
z3, S3 = ir_report(24, 3) if False else ir_report(12, 3)
z3b, S3b = ir_report(24, 3)

# ---- 3D sign --------------------------------------------------------------
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)]
Vv, _ = V_of_pairs(12, 3, V3, 0.3)
def green(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)
    li = np.where(lam > 1e-12, 1/np.where(lam > 1e-12, lam, 1), 0.0)
    return np.fft.ifftn(li).real
Gr = green(12, 3)
g = np.array([Gr[0,0,0] - Gr[tuple(np.mod(v,12))] for v in V3])
Ad = np.vstack([np.ones_like(g), g]).T
cf, *_ = np.linalg.lstsq(Ad, Vv, rcond=None)
fr = np.sqrt(np.mean((Vv-Ad@cf)**2))/np.ptp(Vv)
print(f"3D SIGN: slope b = {cf[1]:+.6e} ({'ATTRACTIVE' if cf[1]>0 else 'REPULSIVE'}); bare-Green fracRMS = {fr:.3%} (ruler only)")
print(f"dressed stiffness (toy 1/G_N): 2D S = {S96[0]:.3f}, 3D S = {S3b[0]:.3f}")