#!/usr/bin/env python3
"""
Nf=2 dynamical-fermion HMC for SU(3) with UNIMPROVED Wilson fermions.

Built from the locked-down method spec (Gattringer-Lang conventions):
  D = 1 - kappa * sum_mu [ (1-gamma_mu) U_mu(x) delta_{x+mu} + (1+gamma_mu) U_mu^dag(x-mu) delta_{x-mu} ]
  S_pf = phi^dag (M^dag M)^-1 phi,   M^dag = gamma5 M gamma5 (gamma5-Hermiticity)
  H = sum_{x,mu} (-tr P^2) + S_gauge + S_pf,  with P traceless anti-Hermitian momenta.

Validation gates (run with --test): gamma algebra, free-field Dirac, gamma5-Hermiticity,
CG residual, ANALYTIC vs NUMERICAL fermion force, reversibility, dH~dt^4, <exp(-dH)>=1.
Physics gate (run with --run): dynamical plaquette must exceed the quenched value at matched beta.

NOTE: numba caching is intentionally OFF here (no cache=True) to avoid the module-identity
cache poisoning seen with su3_fast; compile cost is paid once per process.
"""
from __future__ import annotations
import argparse, json, math, time
from pathlib import Path
import numpy as np
from numba import njit, prange

ND = 4
NC = 3
NS = 4  # Dirac components

# ---------------------------------------------------------------------------
# Gamma matrices: Euclidean Hermitian, chiral (Weyl) basis. gamma5=diag(1,1,-1,-1).
# ---------------------------------------------------------------------------
def _build_gammas():
    s1 = np.array([[0, 1], [1, 0]], dtype=np.complex128)
    s2 = np.array([[0, -1j], [1j, 0]], dtype=np.complex128)
    s3 = np.array([[1, 0], [0, -1]], dtype=np.complex128)
    I2 = np.eye(2, dtype=np.complex128)
    Z2 = np.zeros((2, 2), dtype=np.complex128)
    def block(A, B, C, D):
        return np.block([[A, B], [C, D]])
    # gamma_k = [[0, -i s_k],[ i s_k, 0]] ; gamma_4 = [[0, I],[I, 0]]
    g1 = block(Z2, -1j * s1, 1j * s1, Z2)
    g2 = block(Z2, -1j * s2, 1j * s2, Z2)
    g3 = block(Z2, -1j * s3, 1j * s3, Z2)
    g4 = block(Z2, I2, I2, Z2)
    g5 = np.diag([1, 1, -1, -1]).astype(np.complex128)
    return np.stack([g1, g2, g3, g4]), g5

GAMMA, GAMMA5 = _build_gammas()
# Dirac projectors used in the hopping term
PMINUS = np.stack([np.eye(4, dtype=np.complex128) - GAMMA[mu] for mu in range(4)])  # (1 - g_mu)
PPLUS = np.stack([np.eye(4, dtype=np.complex128) + GAMMA[mu] for mu in range(4)])   # (1 + g_mu)

# Gell-Mann matrices and generators T_a = lambda_a/2 (Hermitian, tr T_aT_b=delta/2)
def _gellmann():
    l = np.zeros((8, 3, 3), dtype=np.complex128)
    l[0] = [[0, 1, 0], [1, 0, 0], [0, 0, 0]]
    l[1] = [[0, -1j, 0], [1j, 0, 0], [0, 0, 0]]
    l[2] = [[1, 0, 0], [0, -1, 0], [0, 0, 0]]
    l[3] = [[0, 0, 1], [0, 0, 0], [1, 0, 0]]
    l[4] = [[0, 0, -1j], [0, 0, 0], [1j, 0, 0]]
    l[5] = [[0, 0, 0], [0, 0, 1], [0, 1, 0]]
    l[6] = [[0, 0, 0], [0, 0, -1j], [0, 1j, 0]]
    l[7] = (1 / math.sqrt(3)) * np.array([[1, 0, 0], [0, 1, 0], [0, 0, -2]], dtype=np.complex128)
    return l
LAMBDA = _gellmann()
TGEN = LAMBDA / 2.0  # Hermitian generators


# ---------------------------------------------------------------------------
# Geometry
# ---------------------------------------------------------------------------
def build_neighbors(L):
    V = L ** ND
    fwd = np.empty((V, ND), dtype=np.int64)
    bwd = np.empty((V, ND), dtype=np.int64)
    par = np.empty(V, dtype=np.int64)
    for s in range(V):
        c = [0, 0, 0, 0]; r = s
        for mu in range(ND):
            c[mu] = r % L; r //= L
        par[s] = (c[0] + c[1] + c[2] + c[3]) % 2
        for mu in range(ND):
            cp = c.copy(); cp[mu] = (cp[mu] + 1) % L
            cm = c.copy(); cm[mu] = (cm[mu] - 1) % L
            fwd[s, mu] = ((cp[3] * L + cp[2]) * L + cp[1]) * L + cp[0]
            bwd[s, mu] = ((cm[3] * L + cm[2]) * L + cm[1]) * L + cm[0]
    return fwd, bwd, par


def cold_start(L):
    V = L ** ND
    U = np.zeros((V, ND, 3, 3), dtype=np.complex128)
    for s in range(V):
        for mu in range(ND):
            for i in range(3):
                U[s, mu, i, i] = 1.0
    return U


# ---------------------------------------------------------------------------
# Wilson Dirac operator   out = in - kappa * hopping
# psi shape (V, 4, 3)  [site, dirac, colour]
# ---------------------------------------------------------------------------
@njit(parallel=True)
def apply_M(psi, U, fwd, bwd, kappa, Pm, Pp):
    V = psi.shape[0]
    out = psi.copy()
    for x in prange(V):
        for mu in range(ND):
            xf = fwd[x, mu]
            xb = bwd[x, mu]
            # forward: (1-g_mu) U_mu(x) psi(x+mu)
            cf = np.zeros((4, 3), dtype=np.complex128)  # U_mu(x) psi(x+mu), per dirac
            for d in range(4):
                for a in range(3):
                    acc = 0.0 + 0.0j
                    for b in range(3):
                        acc += U[x, mu, a, b] * psi[xf, d, b]
                    cf[d, a] = acc
            # backward: (1+g_mu) U_mu^dag(x-mu) psi(x-mu)
            cb = np.zeros((4, 3), dtype=np.complex128)
            for d in range(4):
                for a in range(3):
                    acc = 0.0 + 0.0j
                    for b in range(3):
                        acc += np.conj(U[xb, mu, b, a]) * psi[xb, d, b]
                    cb[d, a] = acc
            # apply Dirac projectors and subtract
            for d2 in range(4):
                for a in range(3):
                    accf = 0.0 + 0.0j
                    accb = 0.0 + 0.0j
                    for d in range(4):
                        accf += Pm[mu, d2, d] * cf[d, a]
                        accb += Pp[mu, d2, d] * cb[d, a]
                    out[x, d2, a] -= kappa * (accf + accb)
    return out


@njit(parallel=True)
def apply_g5(psi, g5):
    V = psi.shape[0]
    out = np.empty_like(psi)
    for x in prange(V):
        for d in range(4):
            for a in range(3):
                out[x, d, a] = g5[d, d] * psi[x, d, a]
    return out


def apply_Mdag(psi, U, fwd, bwd, kappa, Pm, Pp, g5):
    # M^dag = g5 M g5
    return apply_g5(apply_M(apply_g5(psi, g5), U, fwd, bwd, kappa, Pm, Pp), g5)


def apply_MdagM(psi, U, fwd, bwd, kappa, Pm, Pp, g5):
    return apply_Mdag(apply_M(psi, U, fwd, bwd, kappa, Pm, Pp), U, fwd, bwd, kappa, Pm, Pp, g5)


@njit(parallel=True)
def cdot(a, b):
    # sum conj(a) . b  over all components -> complex
    V = a.shape[0]
    re = 0.0; im = 0.0
    for x in prange(V):
        for d in range(4):
            for c in range(3):
                ar = a[x, d, c].real; ai = a[x, d, c].imag
                br = b[x, d, c].real; bi = b[x, d, c].imag
                re += ar * br + ai * bi
                im += ar * bi - ai * br
    return re + 1j * im


def cg_solve(phi, U, fwd, bwd, kappa, Pm, Pp, g5, tol=1e-9, maxit=2000):
    """Solve (M^dag M) x = phi.  Returns x and iteration count."""
    x = np.zeros_like(phi)
    r = phi.copy()
    p = r.copy()
    rs = cdot(r, r).real
    bnorm = math.sqrt(cdot(phi, phi).real)
    if bnorm == 0:
        return x, 0
    for it in range(maxit):
        Ap = apply_MdagM(p, U, fwd, bwd, kappa, Pm, Pp, g5)
        alpha = rs / cdot(p, Ap).real
        x += alpha * p
        r -= alpha * Ap
        rs_new = cdot(r, r).real
        if math.sqrt(rs_new) / bnorm < tol:
            return x, it + 1
        p = r + (rs_new / rs) * p
        rs = rs_new
    return x, maxit


# ---------------------------------------------------------------------------
# Gauge sector (staple, plaquette, TA, su3 exp, gauge force)
# ---------------------------------------------------------------------------
@njit
def _mm(a, b):
    c = np.zeros((3, 3), dtype=np.complex128)
    for i in range(3):
        for k in range(3):
            aik = a[i, k]
            for j in range(3):
                c[i, j] += aik * b[k, j]
    return c

@njit
def _mmd(a, b):
    c = np.zeros((3, 3), dtype=np.complex128)
    for i in range(3):
        for j in range(3):
            s = 0.0 + 0.0j
            for k in range(3):
                s += a[i, k] * np.conj(b[j, k])
            c[i, j] = s
    return c

@njit
def staple(U, fwd, bwd, s, mu):
    A = np.zeros((3, 3), dtype=np.complex128)
    spm = fwd[s, mu]
    for nu in range(ND):
        if nu == mu:
            continue
        spn = fwd[s, nu]; smn = bwd[s, nu]; spmn = bwd[spm, nu]
        t1 = _mmd(U[spm, nu], U[spn, mu])
        t2 = _mmd(t1, U[s, nu])
        b1 = np.zeros((3, 3), dtype=np.complex128)
        for i in range(3):
            for j in range(3):
                ss = 0.0 + 0.0j
                for k in range(3):
                    ss += np.conj(U[spmn, nu][k, i]) * np.conj(U[smn, mu][j, k])
                b1[i, j] = ss
        b2 = _mm(b1, U[smn, nu])
        for i in range(3):
            for j in range(3):
                A[i, j] += t2[i, j] + b2[i, j]
    return A

@njit(parallel=True)
def avg_plaquette(U, fwd):
    V = U.shape[0]; acc = 0.0
    for s in prange(V):
        for mu in range(ND):
            for nu in range(mu + 1, ND):
                spm = fwd[s, mu]; spn = fwd[s, nu]
                a = _mm(U[s, mu], U[spm, nu])
                b = _mmd(a, U[spn, mu])
                c = _mmd(b, U[s, nu])
                acc += (c[0, 0].real + c[1, 1].real + c[2, 2].real) / 3.0
    return acc / (V * 6.0)

@njit
def ta(M):
    # traceless anti-Hermitian projection
    out = np.zeros((3, 3), dtype=np.complex128)
    for i in range(3):
        for j in range(3):
            out[i, j] = 0.5 * (M[i, j] - np.conj(M[j, i]))
    tr = (out[0, 0] + out[1, 1] + out[2, 2]) / 3.0
    for i in range(3):
        out[i, i] -= tr
    return out

@njit
def su3_exp(A):
    # exp(A), A anti-Hermitian traceless; Taylor + Gram-Schmidt reunitarise
    R = np.eye(3, dtype=np.complex128)
    term = np.eye(3, dtype=np.complex128)
    for k in range(1, 20):
        term = _mm(term, A) / k
        for i in range(3):
            for j in range(3):
                R[i, j] += term[i, j]
    # reunitarise (Gram-Schmidt rows)
    n0 = math.sqrt(R[0, 0].real**2 + R[0, 0].imag**2 + R[0, 1].real**2 + R[0, 1].imag**2 + R[0, 2].real**2 + R[0, 2].imag**2)
    for j in range(3):
        R[0, j] /= n0
    dot = 0.0 + 0.0j
    for j in range(3):
        dot += np.conj(R[0, j]) * R[1, j]
    for j in range(3):
        R[1, j] -= dot * R[0, j]
    n1 = math.sqrt(R[1, 0].real**2 + R[1, 0].imag**2 + R[1, 1].real**2 + R[1, 1].imag**2 + R[1, 2].real**2 + R[1, 2].imag**2)
    for j in range(3):
        R[1, j] /= n1
    R[2, 0] = np.conj(R[0, 1] * R[1, 2] - R[0, 2] * R[1, 1])
    R[2, 1] = np.conj(R[0, 2] * R[1, 0] - R[0, 0] * R[1, 2])
    R[2, 2] = np.conj(R[0, 0] * R[1, 1] - R[0, 1] * R[1, 0])
    return R

@njit(parallel=True)
def exp_update_links(U, P, dt):
    # U_mu(x) <- exp(dt P_mu(x)) U_mu(x), parallel over sites
    V = U.shape[0]
    for s in prange(V):
        for mu in range(ND):
            E = su3_exp(dt * P[s, mu])
            M = _mm(E, U[s, mu])
            for i in range(3):
                for j in range(3):
                    U[s, mu, i, j] = M[i, j]


@njit(parallel=True)
def kinetic_njit(P):
    # T = sum_{x,mu} -Re tr(P^2) = (1/2) sum pi_a^2
    V = P.shape[0]
    acc = 0.0
    for s in prange(V):
        loc = 0.0
        for mu in range(ND):
            M = P[s, mu]
            tr = 0.0
            for i in range(3):
                for k in range(3):
                    tr += (M[i, k] * M[k, i]).real
            loc += -tr
        acc += loc
    return acc


@njit(parallel=True)
def gauge_force(U, fwd, bwd, beta):
    # S_g = const - (beta/3) Re tr[ U_mu(x) A_mu(x) ]  (A = staple, U A = plaquettes).
    # F^g_mu(x) = dS_g/dw = (beta/6) TA[ U_mu(x) A_mu(x) ].
    V = U.shape[0]
    F = np.zeros((V, ND, 3, 3), dtype=np.complex128)
    for s in prange(V):
        for mu in range(ND):
            A = staple(U, fwd, bwd, s, mu)
            UA = _mm(U[s, mu], A)  # U A  (= sum of plaquettes through the link)
            T = ta(UA)
            for i in range(3):
                for j in range(3):
                    F[s, mu, i, j] = (beta / 6.0) * T[i, j]
    return F

@njit(parallel=True)
def gauge_action(U, fwd, beta):
    # S_g = (beta/3) sum_plaq Re tr(1 - U_plaq) ; constant offset irrelevant
    V = U.shape[0]; acc = 0.0
    for s in prange(V):
        for mu in range(ND):
            for nu in range(mu + 1, ND):
                spm = fwd[s, mu]; spn = fwd[s, nu]
                a = _mm(U[s, mu], U[spm, nu])
                b = _mmd(a, U[spn, mu])
                c = _mmd(b, U[s, nu])
                acc += (3.0 - (c[0, 0].real + c[1, 1].real + c[2, 2].real))
    return (beta / 3.0) * acc


# ---------------------------------------------------------------------------
# Fermion force (analytic).  X = (M^dag M)^-1 phi ; Y = M X (full lattice).
# F^f_mu(x) = -kappa * TA{ U_mu(x) * [ (P-_mu Y(x+mu)) X(x)^dag + (P+_mu X(x+mu)) Y(x)^dag ] }
# where the spinor outer product over the Dirac index gives a 3x3 colour matrix.
# (Exact form is fixed by the numerical-force gate in --test.)
# ---------------------------------------------------------------------------
@njit(parallel=True)
def fermion_force(X, Y, U, fwd, kappa, Pm, Pp):
    # Derived: dS_pf = -2 Re[Y^dag dM X], Y = M X, X = (M^dag M)^-1 phi.
    # F^f_mu(x) = kappa * TA( B2 - B1 ), with
    #   B1_ab = sum_d (U_mu(x) X(x+mu))_{d,a} conj( ((1-g_mu) Y(x))_{d,b} )
    #   B2_ab = sum_d X(x)_{d,a} conj( (U_mu(x) (1+g_mu) Y(x+mu))_{d,b} )
    V = X.shape[0]
    F = np.zeros((V, ND, 3, 3), dtype=np.complex128)
    for x in prange(V):
        for mu in range(ND):
            xf = fwd[x, mu]
            # U_mu(x) X(x+mu)  (per dirac)
            UXf = np.zeros((4, 3), dtype=np.complex128)
            for d in range(4):
                for a in range(3):
                    acc = 0.0 + 0.0j
                    for c in range(3):
                        acc += U[x, mu, a, c] * X[xf, d, c]
                    UXf[d, a] = acc
            # (1-g_mu) Y(x)
            PmYx = np.zeros((4, 3), dtype=np.complex128)
            for d2 in range(4):
                for b in range(3):
                    acc = 0.0 + 0.0j
                    for d in range(4):
                        acc += Pm[mu, d2, d] * Y[x, d, b]
                    PmYx[d2, b] = acc
            # (1+g_mu) Y(x+mu), then U_mu(x) * that
            PpYf = np.zeros((4, 3), dtype=np.complex128)
            for d2 in range(4):
                for b in range(3):
                    acc = 0.0 + 0.0j
                    for d in range(4):
                        acc += Pp[mu, d2, d] * Y[xf, d, b]
                    PpYf[d2, b] = acc
            UPpYf = np.zeros((4, 3), dtype=np.complex128)
            for d in range(4):
                for a in range(3):
                    acc = 0.0 + 0.0j
                    for c in range(3):
                        acc += U[x, mu, a, c] * PpYf[d, c]
                    UPpYf[d, a] = acc
            C = np.zeros((3, 3), dtype=np.complex128)
            for a in range(3):
                for b in range(3):
                    acc = 0.0 + 0.0j
                    for d in range(4):
                        acc += X[x, d, a] * np.conj(UPpYf[d, b])      # B2
                        acc -= UXf[d, a] * np.conj(PmYx[d, b])        # - B1
                    C[a, b] = acc
            T = ta(C)
            for i in range(3):
                for j in range(3):
                    F[x, mu, i, j] = kappa * T[i, j]
    return F


def pf_action(phi, U, fwd, bwd, kappa, Pm, Pp, g5, tol):
    X, nit = cg_solve(phi, U, fwd, bwd, kappa, Pm, Pp, g5, tol=tol)
    return cdot(phi, X).real, X, nit


# ---------------------------------------------------------------------------
# Numerical fermion force (ground-truth validator): dS_pf/dw_a via finite diff.
# Returns force matrix at link (x0,mu0) as sum_a (dS/dw_a) T_a-style anti-herm.
# ---------------------------------------------------------------------------
def numeric_fermion_force_link(phi, U, fwd, bwd, kappa, x0, mu0, h=1e-4, tol=1e-11):
    Fmat = np.zeros((3, 3), dtype=np.complex128)
    U0 = U[x0, mu0].copy()
    for a in range(8):
        Ta = TGEN[a]
        # U -> exp(+/- h * i Ta) U0   (i Ta is anti-Hermitian)
        Ep = np.eye(3, dtype=np.complex128)
        # matrix exp of +/- i h Ta (small) via expm-ish: use scipy-free series
        Aplus = 1j * h * Ta
        Eplus = _expm_small(Aplus)
        Eminus = _expm_small(-Aplus)
        U[x0, mu0] = Eplus @ U0
        Sp, _, _ = pf_action(phi, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5, tol)
        U[x0, mu0] = Eminus @ U0
        Sm, _, _ = pf_action(phi, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5, tol)
        dS = (Sp - Sm) / (2 * h)  # dS/dw_a
        Fmat += dS * (1j * Ta)    # assemble anti-Herm force matrix: F = sum_a (dS/dw_a)(i T_a)
    U[x0, mu0] = U0
    return Fmat


def _expm_small(A):
    R = np.eye(3, dtype=np.complex128)
    term = np.eye(3, dtype=np.complex128)
    for k in range(1, 24):
        term = term @ A / k
        R = R + term
    return R


# ---------------------------------------------------------------------------
# HMC
# ---------------------------------------------------------------------------
def draw_momenta(V, rng):
    # P = sum_a pi_a (i T_a),  pi_a ~ N(0,1).  Anti-Hermitian traceless.
    pis = rng.normal(size=(V, ND, 8))
    P = np.zeros((V, ND, 3, 3), dtype=np.complex128)
    for a in range(8):
        P += pis[..., a][..., None, None] * (1j * TGEN[a])[None, None]
    return P, pis

def kinetic(P):
    # T = sum_{x,mu} (-Re tr P^2) = sum (1/2) pi_a^2  (parallel kernel)
    return kinetic_njit(P)

def total_force(U, fwd, bwd, beta, phi, kappa, tol):
    Fg = gauge_force(U, fwd, bwd, beta)
    X, nit = cg_solve(phi, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5, tol=tol)
    Y = apply_M(X, U, fwd, bwd, kappa, PMINUS, PPLUS)
    Ff = fermion_force(X, Y, U, fwd, kappa, PMINUS, PPLUS)
    return Fg + Ff, nit

def leapfrog(U, P, fwd, bwd, beta, phi, kappa, n_md, dt, force_tol):
    # force uses a LOOSE CG tolerance (deterministic -> reversibility preserved);
    # the exact action with tight CG is used only in accept/reject.
    F, _ = total_force(U, fwd, bwd, beta, phi, kappa, force_tol)
    P = P - 0.5 * dt * F
    for step in range(n_md):
        exp_update_links(U, P, dt)  # parallel U drift
        F, _ = total_force(U, fwd, bwd, beta, phi, kappa, force_tol)
        if step == n_md - 1:
            P = P - 0.5 * dt * F
        else:
            P = P - dt * F
    return U, P

def hmc_trajectory(U, fwd, bwd, par, beta, kappa, n_md, dt, rng, tol=1e-8, force_tol=1e-5):
    V = U.shape[0]
    # refresh momenta
    P, _ = draw_momenta(V, rng)
    # heatbath pseudofermion: phi = M^dag R, R complex gaussian var 1/2 per real comp
    R = (rng.normal(size=(V, 4, 3)) + 1j * rng.normal(size=(V, 4, 3))) / math.sqrt(2.0)
    phi = apply_Mdag(R, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5)
    # initial H
    Spf0, X0, _ = pf_action(phi, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5, tol)
    H0 = kinetic(P) + gauge_action(U, fwd, beta) + Spf0
    U0 = U.copy()
    U, P = leapfrog(U, P, fwd, bwd, beta, phi, kappa, n_md, dt, force_tol)
    Spf1, _, _ = pf_action(phi, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5, tol)
    H1 = kinetic(P) + gauge_action(U, fwd, beta) + Spf1
    dH = H1 - H0
    if dH < 0 or rng.random() < math.exp(-dH):
        return U, dH, 1
    else:
        return U0, dH, 0


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
def run_tests(L=4, kappa=0.155, beta=5.6, seed=1):
    rng = np.random.default_rng(seed)
    fwd, bwd, par = build_neighbors(L)
    V = L ** ND
    print(f"=== Wilson HMC validation  L={L}^4  kappa={kappa}  beta={beta} ===")

    # GATE 1: gamma algebra
    ok = True
    for mu in range(4):
        for nu in range(4):
            anti = GAMMA[mu] @ GAMMA[nu] + GAMMA[nu] @ GAMMA[mu]
            expect = 2 * (mu == nu) * np.eye(4)
            if np.max(np.abs(anti - expect)) > 1e-13:
                ok = False
    g5c = GAMMA[0] @ GAMMA[1] @ GAMMA[2] @ GAMMA[3]
    g5_ok = np.max(np.abs(g5c - GAMMA5)) < 1e-13 or np.max(np.abs(g5c + GAMMA5)) < 1e-13
    print(f"[GATE gamma]  {{g_mu,g_nu}}=2d: {'PASS' if ok else 'FAIL'};  "
          f"g1g2g3g4=+/-g5: {'PASS' if g5_ok else 'FAIL'} (g5c diag={np.round(np.real(np.diag(g5c)),3)})")

    # random gauge config
    U = cold_start(L)
    for s in range(V):
        for mu in range(ND):
            A = np.zeros((3, 3), dtype=np.complex128)
            for a in range(8):
                A += rng.normal() * 0.3 * (1j * TGEN[a])
            U[s, mu] = su3_exp(A) @ U[s, mu]

    # GATE 2: gamma5-Hermiticity  ||g5 M g5 - M^dag|| / ||M||
    psi = (rng.normal(size=(V, 4, 3)) + 1j * rng.normal(size=(V, 4, 3)))
    chi = (rng.normal(size=(V, 4, 3)) + 1j * rng.normal(size=(V, 4, 3)))
    Mpsi = apply_M(psi, U, fwd, bwd, kappa, PMINUS, PPLUS)
    Mdchi = apply_Mdag(chi, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5)
    # <chi, M psi> must equal <M^dag chi, psi>
    lhs = cdot(chi, Mpsi); rhs = cdot(Mdchi, psi)
    print(f"[GATE g5-herm]  <chi,M psi> - <M^dag chi,psi> = {abs(lhs-rhs):.2e}  "
          f"{'PASS' if abs(lhs-rhs) < 1e-10 else 'FAIL'}")

    # GATE 3: free-field Dirac eigenvalue (U=1) at zero momentum -> D = 1 - 8 kappa on the
    # constant mode's Wilson part; check apply on constant spinor.
    Uf = cold_start(L)
    const = np.zeros((V, 4, 3), dtype=np.complex128); const[:, 0, 0] = 1.0
    out = apply_M(const, Uf, fwd, bwd, kappa, PMINUS, PPLUS)
    # zero-momentum: hopping sums (1-g)+(1+g)=2 per direction * kappa -> out = (1 - 8 kappa) on the mode
    val = out[0, 0, 0]
    print(f"[GATE free-field]  D(p=0) on const = {val.real:.6f}  expected {1-8*kappa:.6f}  "
          f"{'PASS' if abs(val-(1-8*kappa)) < 1e-9 else 'FAIL'}")

    # GATE 4: CG residual
    phi = apply_Mdag(chi, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5)
    Xcg, nit = cg_solve(phi, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5, tol=1e-10)
    resid = apply_MdagM(Xcg, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5) - phi
    rr = math.sqrt(cdot(resid, resid).real / cdot(phi, phi).real)
    print(f"[GATE CG]  iters={nit}  rel.resid={rr:.2e}  {'PASS' if rr < 1e-9 else 'FAIL'}")

    # GATE 5: analytic vs numerical fermion force on a few links
    X = Xcg
    Y = apply_M(X, U, fwd, bwd, kappa, PMINUS, PPLUS)
    Fan = fermion_force(X, Y, U, fwd, kappa, PMINUS, PPLUS)
    maxrel = 0.0
    for (x0, mu0) in [(0, 0), (1, 2), (5, 3), (7, 1)]:
        Fnum = numeric_fermion_force_link(phi, U, fwd, bwd, kappa, x0, mu0)
        # analytic returns force already with -kappa & TA; numeric returns dS/dw assembled.
        # The MD uses P -= dt*F with F = dS/dU (TA). Compare Fan vs TA(Fnum-style). Numeric gives
        # F_a=dS/dw_a assembled as sum_a (dS/dw_a)(i T_a) which is exactly the anti-herm gradient.
        diff = np.max(np.abs(Fan[x0, mu0] - ta(Fnum)))
        scale = np.max(np.abs(Fan[x0, mu0])) + 1e-12
        maxrel = max(maxrel, diff / scale)
    print(f"[GATE ferm-force]  max rel(analytic - numeric) = {maxrel:.2e}  "
          f"{'PASS' if maxrel < 1e-3 else 'FAIL  <-- fix fermion_force sign/gamma/X-Y'}")

    # GATE 6: reversibility
    P, _ = draw_momenta(V, rng)
    R = (rng.normal(size=(V, 4, 3)) + 1j * rng.normal(size=(V, 4, 3))) / math.sqrt(2.0)
    phi = apply_Mdag(R, U, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5)
    Ucopy = U.copy(); Pcopy = P.copy()
    U2, P2 = leapfrog(U.copy(), P.copy(), fwd, bwd, beta, phi, kappa, 8, 0.05, 1e-10)
    U3, P3 = leapfrog(U2, -P2, fwd, bwd, beta, phi, kappa, 8, 0.05, 1e-10)
    dU = max(np.max(np.abs(U3[s, mu] - Ucopy[s, mu])) for s in range(V) for mu in range(ND))
    print(f"[GATE reversibility]  max|dU| = {dU:.2e}  {'PASS' if dU < 1e-8 else 'FAIL'}")

    # GATE 7: dH ~ dt^? scaling (energy conservation) -> confirms force correctness
    print("[GATE dH-scaling]  (tau=1 fixed; expect <dH> shrink steeply as dt->0)")
    for n_md in [5, 10, 20, 40]:
        dt = 1.0 / n_md
        Ut = U.copy()
        P, _ = draw_momenta(V, rng)
        R = (rng.normal(size=(V, 4, 3)) + 1j * rng.normal(size=(V, 4, 3))) / math.sqrt(2.0)
        phi = apply_Mdag(R, Ut, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5)
        S0, _, _ = pf_action(phi, Ut, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5, 1e-10)
        H0 = kinetic(P) + gauge_action(Ut, fwd, beta) + S0
        Ut, P = leapfrog(Ut, P, fwd, bwd, beta, phi, kappa, n_md, dt, 1e-10)
        S1, _, _ = pf_action(phi, Ut, fwd, bwd, kappa, PMINUS, PPLUS, GAMMA5, 1e-10)
        H1 = kinetic(P) + gauge_action(Ut, fwd, beta) + S1
        print(f"    n_md={n_md:>3} dt={dt:.3f}  dH={H1-H0:+.4e}")


def run_hmc(args):
    rng = np.random.default_rng(args.seed)
    L = args.L
    fwd, bwd, par = build_neighbors(L)
    U = cold_start(L)
    t0 = time.time()
    dHs = []; accs = []; plaqs = []
    for it in range(args.ntraj):
        U, dH, acc = hmc_trajectory(U, fwd, bwd, par, args.beta, args.kappa,
                                    args.n_md, args.dt, rng, tol=args.tol, force_tol=args.force_tol)
        dHs.append(dH); accs.append(acc)
        if it >= args.therm:
            plaqs.append(avg_plaquette(U, fwd))
        if args.verbose and (it + 1) % max(1, args.ntraj // 20) == 0:
            p = avg_plaquette(U, fwd)
            print(f"[traj {it+1}/{args.ntraj}] plaq={p:.5f} dH={dH:+.3e} "
                  f"acc={np.mean(accs):.3f} <e^-dH>={np.mean(np.exp(-np.array(dHs))):.3f} "
                  f"({time.time()-t0:.1f}s)", flush=True)
    summary = {
        "code": "Nf=2 unimproved Wilson HMC (un-preconditioned)",
        "L": L, "beta": args.beta, "kappa": args.kappa,
        "n_md": args.n_md, "dt": args.dt, "tau": args.n_md * args.dt,
        "ntraj": args.ntraj, "therm": args.therm, "seed": args.seed,
        "acceptance": float(np.mean(accs)),
        "exp_minus_dH": float(np.mean(np.exp(-np.array(dHs)))),
        "mean_abs_dH": float(np.mean(np.abs(dHs))),
        "plaquette_dynamical_mean": float(np.mean(plaqs)) if plaqs else None,
        "plaquette_dynamical_stderr": float(np.std(plaqs, ddof=1) / math.sqrt(len(plaqs))) if len(plaqs) > 1 else None,
        "elapsed_sec": time.time() - t0,
    }
    Path(args.outdir).mkdir(parents=True, exist_ok=True)
    tag = f"wilsonhmc_L{L}_b{args.beta}_k{args.kappa}_seed{args.seed}"
    json.dump(summary, open(Path(args.outdir) / f"{tag}_summary.json", "w"), indent=2)
    print(json.dumps(summary, indent=2), flush=True)
    return summary


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--test", action="store_true")
    p.add_argument("--L", type=int, default=4)
    p.add_argument("--beta", type=float, default=5.6)
    p.add_argument("--kappa", type=float, default=0.155)
    p.add_argument("--n_md", type=int, default=10)
    p.add_argument("--dt", type=float, default=0.1)
    p.add_argument("--ntraj", type=int, default=200)
    p.add_argument("--therm", type=int, default=40)
    p.add_argument("--tol", type=float, default=1e-8, help="tight CG tol for the Hamiltonian (accept/reject)")
    p.add_argument("--force_tol", type=float, default=1e-5, help="loose CG tol for the MD force")
    p.add_argument("--seed", type=int, default=1)
    p.add_argument("--outdir", type=str, default="results_hmc")
    p.add_argument("--verbose", action="store_true")
    a = p.parse_args()
    if a.test:
        run_tests(L=a.L, kappa=a.kappa, beta=a.beta, seed=a.seed)
    else:
        run_hmc(a)
