#!/usr/bin/env python3
"""
Numba-accelerated pure-gauge SU(3) Wilson-action lattice code for the
DFD-QCD Boundary Program.

Algorithm (all standard, textbook-correct):
  * Updates: multi-hit Metropolis on random SU(2) subgroups (ergodic) mixed
    with microcanonical overrelaxation reflections (Cabibbo-Marinari).  The
    Metropolis proposal is exactly symmetric (theta drawn symmetric about 0,
    axis uniform on S^2), so detailed balance holds for the Wilson measure
    exp[(beta/3) sum_plaq Re Tr U_plaq].
  * Observables: average plaquette; APE-smeared on-axis Wilson loops W(R,T)
    built from incremental line products; static potential V(R) from the
    large-T behaviour W(R,T) ~ exp[-V(R) T].

This is real lattice gauge theory.  It is NOT a continuum-extrapolated,
physical-point full-QCD result: it is pure gauge (quenched), single lattice
spacing per run, finite volume.  Those boundaries are reported honestly in
the summary JSON and in the companion analysis.

The decisive correctness check is the average plaquette versus the published
SU(3) Wilson-gauge values, e.g. <P>(beta=6.0) ~ 0.5937, <P>(beta=5.7) ~ 0.5491.
Run with --validate to print that comparison.
"""
from __future__ import annotations
import argparse, json, math, time
from pathlib import Path
import numpy as np
from numba import njit, prange

NC = 3
ND = 4

# ----------------------------------------------------------------------------
# Geometry: flat site index over an L^4 lattice, forward/backward neighbour
# tables.  site = ((t*L + z)*L + y)*L + x.
# ----------------------------------------------------------------------------

def build_neighbors(L: int):
    V = L ** ND
    fwd = np.empty((V, ND), dtype=np.int64)
    bwd = np.empty((V, ND), 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
        for mu in range(ND):
            cp = c.copy(); cp[mu] = (cp[mu] + 1) % L
            cm = c.copy(); cm[mu] = (cm[mu] - 1) % L
            fp = ((cp[3] * L + cp[2]) * L + cp[1]) * L + cp[0]
            fm = ((cm[3] * L + cm[2]) * L + cm[1]) * L + cm[0]
            fwd[s, mu] = fp
            bwd[s, mu] = fm
    return fwd, bwd


# ----------------------------------------------------------------------------
# 3x3 complex linear algebra helpers (unrolled-ish, numba nopython).
# ----------------------------------------------------------------------------

@njit(cache=True, inline='always')
def mat_mul(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(cache=True, inline='always')
def mat_mul_dag(a, b):
    # a @ b^dagger
    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(cache=True, inline='always')
def dag_mul(a, b):
    # a^dagger @ 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 += np.conj(a[k, i]) * b[k, j]
            c[i, j] = s
    return c


@njit(cache=True)
def project_su3(M, V, niter=7):
    # Modify V in place to maximize Re Tr(V^dag M) over SU(3), by iterated
    # SU(2)-subgroup (Cabibbo-Marinari) maximization.  V should be initialised
    # near the desired answer (e.g. the original link).  This is the correct
    # APE/HYP smearing projector (NOT Gram-Schmidt).
    for _ in range(niter):
        for isub in range(3):
            i = _SUB[isub][0]
            j = _SUB[isub][1]
            Y = dag_mul(V, M)  # V^dag @ M
            Yii = Y[i, i]; Yjj = Y[j, j]; Yij = Y[i, j]; Yji = Y[j, i]
            q0 = Yii.real + Yjj.real
            q1 = Yij.imag + Yji.imag
            q2 = Yij.real - Yji.real
            q3 = Yii.imag - Yjj.imag
            nrm = math.sqrt(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3)
            if nrm < 1e-14:
                continue
            a0 = q0 / nrm; a1 = q1 / nrm; a2 = q2 / nrm; a3 = q3 / nrm
            gii = a0 + 1j * a3; gij = a2 + 1j * a1
            gji = -a2 + 1j * a1; gjj = a0 - 1j * a3
            # V <- V @ g  (g embedded in cols/rows i,j)
            for row in range(3):
                vi = V[row, i]; vj = V[row, j]
                V[row, i] = vi * gii + vj * gji
                V[row, j] = vi * gij + vj * gjj
    return V


@njit(cache=True, inline='always')
def reunitarize(u):
    # Gram-Schmidt on rows, then fix third row = (row0 x row1)* for det=1.
    n0 = math.sqrt((u[0, 0].real**2 + u[0, 0].imag**2 +
                    u[0, 1].real**2 + u[0, 1].imag**2 +
                    u[0, 2].real**2 + u[0, 2].imag**2))
    for j in range(3):
        u[0, j] = u[0, j] / n0
    # row1 -= (row0^* . row1) row0
    dot = 0.0 + 0.0j
    for j in range(3):
        dot += np.conj(u[0, j]) * u[1, j]
    for j in range(3):
        u[1, j] = u[1, j] - dot * u[0, j]
    n1 = math.sqrt((u[1, 0].real**2 + u[1, 0].imag**2 +
                    u[1, 1].real**2 + u[1, 1].imag**2 +
                    u[1, 2].real**2 + u[1, 2].imag**2))
    for j in range(3):
        u[1, j] = u[1, j] / n1
    # row2 = conj(row0 x row1)
    u[2, 0] = np.conj(u[0, 1] * u[1, 2] - u[0, 2] * u[1, 1])
    u[2, 1] = np.conj(u[0, 2] * u[1, 0] - u[0, 0] * u[1, 2])
    u[2, 2] = np.conj(u[0, 0] * u[1, 1] - u[0, 1] * u[1, 0])
    return u


# ----------------------------------------------------------------------------
# Staple: A_mu(x) such that the link action is -(beta/3) Re Tr[ U_mu(x) A ].
# ----------------------------------------------------------------------------

@njit(cache=True)
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]
        # forward staple: U_nu(x+mu) U_mu(x+nu)^dag U_nu(x)^dag
        t1 = mat_mul_dag(U[spm, nu], U[spn, mu])
        t2 = mat_mul_dag(t1, U[s, nu])
        for i in range(3):
            for j in range(3):
                A[i, j] += t2[i, j]
        # backward staple: U_nu(x+mu-nu)^dag U_mu(x-nu)^dag U_nu(x-nu)
        # need (a^dag @ b^dag) for a=U_nu(x+mu-nu), b=U_mu(x-nu):
        b1 = np.zeros((3, 3), dtype=np.complex128)
        for i in range(3):
            for j in range(3):
                ssum = 0.0 + 0.0j
                for k in range(3):
                    ssum += np.conj(U[spmn, nu][k, i]) * np.conj(U[smn, mu][j, k])
                b1[i, j] = ssum
        b2 = mat_mul(b1, U[smn, nu])
        for i in range(3):
            for j in range(3):
                A[i, j] += b2[i, j]
    return A


# ----------------------------------------------------------------------------
# Average plaquette  <(1/3) Re Tr U_plaq>  over all sites and 6 planes.
# ----------------------------------------------------------------------------

@njit(cache=True, parallel=True)
def avg_plaquette(U, fwd):
    V = U.shape[0]
    acc = 0.0
    for s in prange(V):
        local = 0.0
        for mu in range(ND):
            for nu in range(mu + 1, ND):
                spm = fwd[s, mu]
                spn = fwd[s, nu]
                # U_mu(x) U_nu(x+mu) U_mu(x+nu)^dag U_nu(x)^dag
                a = mat_mul(U[s, mu], U[spm, nu])
                b = mat_mul_dag(a, U[spn, mu])
                c = mat_mul_dag(b, U[s, nu])
                local += (c[0, 0].real + c[1, 1].real + c[2, 2].real) / 3.0
        acc += local
    return acc / (V * 6.0)


# ----------------------------------------------------------------------------
# One combined update sweep over all links:
#   - nhit Metropolis hits per link on random SU(2) subgroups (ergodic)
#   - n_or overrelaxation reflections per link (fast decorrelation)
# Returns Metropolis acceptance fraction.
# ----------------------------------------------------------------------------

# subgroup index pairs
_SUB = ((0, 1), (0, 2), (1, 2))


@njit(cache=True)
def _apply_left_su2(U, s, mu, i, j, a0, a1, a2, a3):
    # U_mu(x) <- g . U_mu(x), g = embedded SU(2) in rows (i,j)
    gii = a0 + 1j * a3
    gij = a2 + 1j * a1
    gji = -a2 + 1j * a1
    gjj = a0 - 1j * a3
    for col in range(3):
        ui = U[s, mu][i, col]
        uj = U[s, mu][j, col]
        U[s, mu][i, col] = gii * ui + gij * uj
        U[s, mu][j, col] = gji * ui + gjj * uj


@njit(cache=True)
def update_sweep(U, fwd, bwd, beta, eps, nhit, n_or):
    V = U.shape[0]
    nacc = 0
    ntot = 0
    c = beta / 3.0
    for s in range(V):
        for mu in range(ND):
            A = staple(U, fwd, bwd, s, mu)
            # R = U_mu(x) . A  (3x3). Re Tr(g R) drives the update.
            R = mat_mul(U[s, mu], A)
            # ---- Metropolis hits ----
            for _ in range(nhit):
                isub = np.random.randint(0, 3)
                i = _SUB[isub][0]
                j = _SUB[isub][1]
                # propose near-identity SU(2): theta symmetric about 0, axis uniform
                theta = eps * np.random.normal()
                nx = np.random.normal(); ny = np.random.normal(); nz = np.random.normal()
                nn = math.sqrt(nx * nx + ny * ny + nz * nz)
                if nn < 1e-12:
                    continue
                st = math.sin(theta)
                a0 = math.cos(theta)
                a1 = st * nx / nn
                a2 = st * ny / nn
                a3 = st * nz / nn
                # dRe = Re Tr(g R_block) - Re Tr(R_block) for block (i,j)
                Rii = R[i, i]; Rjj = R[j, j]; Rij = R[i, j]; Rji = R[j, i]
                new = (a0 * (Rii.real + Rjj.real)
                       - a3 * (Rii.imag - Rjj.imag)
                       - a1 * (Rji.imag + Rij.imag)
                       + a2 * (Rji.real - Rij.real))
                old = Rii.real + Rjj.real
                dRe = new - old
                ntot += 1
                if dRe >= 0.0 or np.random.random() < math.exp(c * dRe):
                    nacc += 1
                    _apply_left_su2(U, s, mu, i, j, a0, a1, a2, a3)
                    # update R rows i,j: R <- g R
                    gii = a0 + 1j * a3; gij = a2 + 1j * a1
                    gji = -a2 + 1j * a1; gjj = a0 - 1j * a3
                    for col in range(3):
                        ri = R[i, col]; rj = R[j, col]
                        R[i, col] = gii * ri + gij * rj
                        R[j, col] = gji * ri + gjj * rj
            # ---- Overrelaxation reflections ----
            for _ in range(n_or):
                for isub in range(3):
                    i = _SUB[isub][0]
                    j = _SUB[isub][1]
                    Rii = R[i, i]; Rjj = R[j, j]; Rij = R[i, j]; Rji = R[j, i]
                    w0 = Rii.real + Rjj.real
                    w1 = -(Rji.imag + Rij.imag)
                    w2 = Rji.real - Rij.real
                    w3 = Rjj.imag - Rii.imag
                    k = math.sqrt(w0 * w0 + w1 * w1 + w2 * w2 + w3 * w3)
                    if k < 1e-12:
                        continue
                    w0 /= k; w1 /= k; w2 /= k; w3 /= k
                    # reflect identity (1,0,0,0) about unit w: a = 2 w0 w - (1,0,0,0)
                    a0 = 2.0 * w0 * w0 - 1.0
                    a1 = 2.0 * w0 * w1
                    a2 = 2.0 * w0 * w2
                    a3 = 2.0 * w0 * w3
                    _apply_left_su2(U, s, mu, i, j, a0, a1, a2, a3)
                    gii = a0 + 1j * a3; gij = a2 + 1j * a1
                    gji = -a2 + 1j * a1; gjj = a0 - 1j * a3
                    for col in range(3):
                        ri = R[i, col]; rj = R[j, col]
                        R[i, col] = gii * ri + gij * rj
                        R[j, col] = gji * ri + gjj * rj
            reunitarize(U[s, mu])
    if ntot == 0:
        return 1.0
    return nacc / ntot


@njit(cache=True)
def _update_link(U, fwd, bwd, s, mu, beta, eps, nhit, n_or):
    """Metropolis(nhit) + overrelax(n_or) update of the single link U[s,mu].
    Identical physics to update_sweep's inner body; safe to call in parallel
    over a fixed (site-parity, direction) set (no two such links share a staple)."""
    c = beta / 3.0
    nacc = 0
    ntot = 0
    A = staple(U, fwd, bwd, s, mu)
    R = mat_mul(U[s, mu], A)
    for _ in range(nhit):
        isub = np.random.randint(0, 3)
        i = _SUB[isub][0]; j = _SUB[isub][1]
        theta = eps * np.random.normal()
        nx = np.random.normal(); ny = np.random.normal(); nz = np.random.normal()
        nn = math.sqrt(nx * nx + ny * ny + nz * nz)
        if nn < 1e-12:
            continue
        st = math.sin(theta); a0 = math.cos(theta)
        a1 = st * nx / nn; a2 = st * ny / nn; a3 = st * nz / nn
        Rii = R[i, i]; Rjj = R[j, j]; Rij = R[i, j]; Rji = R[j, i]
        new = (a0 * (Rii.real + Rjj.real) - a3 * (Rii.imag - Rjj.imag)
               - a1 * (Rji.imag + Rij.imag) + a2 * (Rji.real - Rij.real))
        old = Rii.real + Rjj.real
        dRe = new - old
        ntot += 1
        if dRe >= 0.0 or np.random.random() < math.exp(c * dRe):
            nacc += 1
            _apply_left_su2(U, s, mu, i, j, a0, a1, a2, a3)
            gii = a0 + 1j * a3; gij = a2 + 1j * a1
            gji = -a2 + 1j * a1; gjj = a0 - 1j * a3
            for col in range(3):
                ri = R[i, col]; rj = R[j, col]
                R[i, col] = gii * ri + gij * rj
                R[j, col] = gji * ri + gjj * rj
    for _ in range(n_or):
        for isub in range(3):
            i = _SUB[isub][0]; j = _SUB[isub][1]
            Rii = R[i, i]; Rjj = R[j, j]; Rij = R[i, j]; Rji = R[j, i]
            w0 = Rii.real + Rjj.real
            w1 = -(Rji.imag + Rij.imag)
            w2 = Rji.real - Rij.real
            w3 = Rjj.imag - Rii.imag
            k = math.sqrt(w0 * w0 + w1 * w1 + w2 * w2 + w3 * w3)
            if k < 1e-12:
                continue
            w0 /= k; w1 /= k; w2 /= k; w3 /= k
            a0 = 2.0 * w0 * w0 - 1.0
            a1 = 2.0 * w0 * w1; a2 = 2.0 * w0 * w2; a3 = 2.0 * w0 * w3
            _apply_left_su2(U, s, mu, i, j, a0, a1, a2, a3)
            gii = a0 + 1j * a3; gij = a2 + 1j * a1
            gji = -a2 + 1j * a1; gjj = a0 - 1j * a3
            for col in range(3):
                ri = R[i, col]; rj = R[j, col]
                R[i, col] = gii * ri + gij * rj
                R[j, col] = gji * ri + gjj * rj
    reunitarize(U[s, mu])
    return nacc, ntot


@njit(cache=True, parallel=True)
def _sweep_color(U, fwd, bwd, sites, mu, beta, eps, nhit, n_or):
    nacc = 0
    ntot = 0
    for k in prange(sites.shape[0]):
        a, t = _update_link(U, fwd, bwd, sites[k], mu, beta, eps, nhit, n_or)
        nacc += a
        ntot += t
    return nacc, ntot


def build_parity(L):
    V = L ** ND
    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
    even = np.where(par == 0)[0].astype(np.int64)
    odd = np.where(par == 1)[0].astype(np.int64)
    return even, odd


def update_sweep_cb(U, fwd, bwd, even, odd, beta, eps, nhit, n_or):
    """Parallel red-black sweep: 2 parities x 4 directions, each prange-parallel."""
    na = 0
    nt = 0
    for sites in (even, odd):
        for mu in range(ND):
            a, t = _sweep_color(U, fwd, bwd, sites, mu, beta, eps, nhit, n_or)
            na += a; nt += t
    return na / nt if nt else 1.0


# ----------------------------------------------------------------------------
# APE smearing of spatial links only (mu in 1,2,3) for static-potential
# operators.  U_smeared = Proj_SU3[ (1-alpha) U + (alpha/4) sum_spatial_staples ].
# ----------------------------------------------------------------------------

@njit(cache=True, parallel=True)
def ape_smear_spatial(U, fwd, bwd, alpha, nstep):
    V = U.shape[0]
    Us = U.copy()
    tmp = U.copy()
    for _ in range(nstep):
        for s in prange(V):
            for mu in range(1, ND):  # spatial only
                acc = np.zeros((3, 3), dtype=np.complex128)
                for i in range(3):
                    for j in range(3):
                        acc[i, j] = (1.0 - alpha) * Us[s, mu][i, j]
                spm = fwd[s, mu]  # x+mu
                for nu in range(1, ND):
                    if nu == mu:
                        continue
                    spn = fwd[s, nu]        # x+nu
                    smn = bwd[s, nu]        # x-nu
                    smnpm = fwd[smn, mu]    # x-nu+mu
                    # forward fattening staple (oriented +mu, same as U_mu(x)):
                    #   C_f = U_nu(x) U_mu(x+nu) U_nu(x+mu)^dag
                    c1 = mat_mul(Us[s, nu], Us[spn, mu])
                    cf = mat_mul_dag(c1, Us[spm, nu])
                    # backward fattening staple:
                    #   C_b = U_nu(x-nu)^dag U_mu(x-nu) U_nu(x-nu+mu)
                    d1 = dag_mul(Us[smn, nu], Us[smn, mu])
                    cb = mat_mul(d1, Us[smnpm, nu])
                    for i in range(3):
                        for j in range(3):
                            acc[i, j] += (alpha / 4.0) * (cf[i, j] + cb[i, j])
                # init projection at the current (pre-smear) link, then project
                # the fattened sum 'acc' onto SU(3) by maximal-trace (Cabibbo-Marinari)
                for i in range(3):
                    for j in range(3):
                        tmp[s, mu][i, j] = Us[s, mu][i, j]
                project_su3(acc, tmp[s, mu], 7)
        for s in prange(V):
            for mu in range(1, ND):
                for i in range(3):
                    for j in range(3):
                        Us[s, mu][i, j] = tmp[s, mu][i, j]
    return Us


# ----------------------------------------------------------------------------
# On-axis Wilson loops W(R,T): R along each spatial direction (smeared links),
# T along time (mu=0, thin links).  Built from incremental line products.
# Returns array W[R, T], R=1..Rmax, T=1..Tmax (averaged over sites & 3 dirs).
# ----------------------------------------------------------------------------

@njit(cache=True)
def _shift_n(site, mu, n, fwd):
    s = site
    for _ in range(n):
        s = fwd[s, mu]
    return s


@njit(cache=True)
def wilson_loops(Uspat, Utime, fwd, Rmax, Tmax):
    V = Uspat.shape[0]
    W = np.zeros((Rmax + 1, Tmax + 1), dtype=np.float64)
    for d in range(1, ND):  # spatial direction for R
        # space line products SP[r,s] = U_d(s) U_d(s+d) ... (length r)
        SP = np.zeros((Rmax + 1, V, 3, 3), dtype=np.complex128)
        for s in range(V):
            for i in range(3):
                for j in range(3):
                    SP[1, s, i, j] = Uspat[s, d][i, j]
        for r in range(2, Rmax + 1):
            for s in range(V):
                sr = _shift_n(s, d, r - 1, fwd)
                prod = mat_mul(SP[r - 1, s], Uspat[sr, d])
                for i in range(3):
                    for j in range(3):
                        SP[r, s, i, j] = prod[i, j]
        # time line products TP[t,s] = U_0(s) ... (length t)
        TP = np.zeros((Tmax + 1, V, 3, 3), dtype=np.complex128)
        for s in range(V):
            for i in range(3):
                for j in range(3):
                    TP[1, s, i, j] = Utime[s, 0][i, j]
        for t in range(2, Tmax + 1):
            for s in range(V):
                st = _shift_n(s, 0, t - 1, fwd)
                prod = mat_mul(TP[t - 1, s], Utime[st, 0])
                for i in range(3):
                    for j in range(3):
                        TP[t, s, i, j] = prod[i, j]
        # assemble loops
        for r in range(1, Rmax + 1):
            for t in range(1, Tmax + 1):
                acc = 0.0
                for s in range(V):
                    s_r = _shift_n(s, d, r, fwd)   # corner after going R in space
                    s_t = _shift_n(s, 0, t, fwd)   # corner after going T in time
                    # W = SP[r,s] . TP[t, s_r] . SP[r, s_t]^dag . TP[t, s]^dag
                    m1 = mat_mul(SP[r, s], TP[t, s_r])
                    m2 = mat_mul_dag(m1, SP[r, s_t])
                    m3 = mat_mul_dag(m2, TP[t, s])
                    acc += (m3[0, 0].real + m3[1, 1].real + m3[2, 2].real) / 3.0
                W[r, t] += acc / V
    for r in range(1, Rmax + 1):
        for t in range(1, Tmax + 1):
            W[r, t] /= 3.0
    return W


# ----------------------------------------------------------------------------
# 0++ scalar glueball operator: zero-momentum, A1+ (cubic scalar) sum of
# spatial plaquettes on each timeslice (time axis = mu=0).  Built from smeared
# spatial links.  O(t) = sum_{x in slice t} sum_{(i,j) spatial} Re Tr U_plaq_ij / 3.
# The connected correlator C(t) = <O(t0+t)O(t0)> - <O>^2 (vacuum subtracted)
# decays as exp(-m_0++ t).  Time slices are indexed by the mu=0 coordinate s%L.
# ----------------------------------------------------------------------------

# spatial plane pairs (i<j), i,j in {1,2,3}
_SPLANES = ((1, 2), (1, 3), (2, 3))


@njit(cache=True)
def glueball_timeslices(Us, fwd, L):
    V = Us.shape[0]
    O = np.zeros(L, dtype=np.float64)
    for s in range(V):
        t0 = s % L  # mu=0 coordinate = "time"
        loc = 0.0
        for p in range(3):
            i = _SPLANES[p][0]
            j = _SPLANES[p][1]
            spi = fwd[s, i]
            spj = fwd[s, j]
            a = mat_mul(Us[s, i], Us[spi, j])
            b = mat_mul_dag(a, Us[spj, i])
            c = mat_mul_dag(b, Us[s, j])
            loc += (c[0, 0].real + c[1, 1].real + c[2, 2].real) / 3.0
        O[t0] += loc
    return O


# ----------------------------------------------------------------------------
# Driver
# ----------------------------------------------------------------------------

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 + 0.0j
    return U


def run(args):
    L = args.L
    fwd, bwd = build_neighbors(L)
    rng_seed = args.seed
    np.random.seed(rng_seed)  # numba uses numpy's RNG state under nopython
    U = cold_start(L)
    even, odd = build_parity(L) if args.parallel else (None, None)

    def do_sweep():
        if args.parallel:
            return update_sweep_cb(U, fwd, bwd, even, odd, args.beta, args.eps, args.nhit, args.n_or)
        return update_sweep(U, fwd, bwd, args.beta, args.eps, args.nhit, args.n_or)

    t0 = time.time()
    # thermalization
    acc_hist = []
    for it in range(args.thermal):
        a = do_sweep()
        acc_hist.append(a)
        if args.verbose and (it + 1) % max(1, args.thermal // 10) == 0:
            p = avg_plaquette(U, fwd)
            print(f"[therm {it+1}/{args.thermal}] plaq={p:.6f} acc={a:.3f} "
                  f"({time.time()-t0:.1f}s)", flush=True)

    # measurements
    rows = []
    glue_levels = [int(x) for x in args.glue_levels.split(",")] if args.glueball else []
    glue_data = []  # per-config arrays shape (nlevel, L)
    for m in range(args.measurements):
        for _ in range(args.skip):
            a = do_sweep()
            acc_hist.append(a)
        plaq = avg_plaquette(U, fwd)
        row = {"measurement": m, "plaquette": float(plaq)}
        if args.potential:
            Us = ape_smear_spatial(U, fwd, bwd, args.ape_alpha, args.ape_steps)
            W = wilson_loops(Us, U, fwd, args.Rmax, args.Tmax)
            for r in range(1, args.Rmax + 1):
                for t in range(1, args.Tmax + 1):
                    row[f"W_{r}x{t}"] = float(W[r, t])
            rows.append({**row})
        elif args.glueball:
            # operator at several APE-smearing levels (cumulative) -> variational basis
            Oset = np.zeros((len(glue_levels), L), dtype=np.float64)
            Us = U.copy()
            prev = 0
            for li, lvl in enumerate(glue_levels):
                if lvl > prev:
                    Us = ape_smear_spatial(Us, fwd, bwd, args.ape_alpha, lvl - prev)
                    prev = lvl
                src = U if lvl == 0 else Us
                Oset[li] = glueball_timeslices(src, fwd, L)
            glue_data.append(Oset)
            rows.append(row)
        else:
            rows.append(row)
        if args.verbose and (m + 1) % max(1, args.measurements // 10) == 0:
            print(f"[meas {m+1}/{args.measurements}] plaq={plaq:.6f} "
                  f"({time.time()-t0:.1f}s)", flush=True)

    plaqs = np.array([r["plaquette"] for r in rows])
    summary = {
        "code_status": "numba_pure_gauge_real_but_finite_volume_single_spacing",
        "gauge_group": "SU(3)",
        "action": "Wilson plaquette pure gauge (quenched)",
        "algorithm": f"metropolis(nhit={args.nhit}) + overrelax(n_or={args.n_or}) per sweep",
        "L": L,
        "dimensions": [L] * 4,
        "beta": args.beta,
        "thermal_sweeps": args.thermal,
        "measurements": args.measurements,
        "skip_sweeps_between_measurements": args.skip,
        "proposal_eps": args.eps,
        "seed": rng_seed,
        "mean_acceptance": float(np.mean(acc_hist)),
        "elapsed_sec": time.time() - t0,
        "plaquette_mean": float(np.mean(plaqs)),
        "plaquette_stderr": float(np.std(plaqs, ddof=1) / math.sqrt(len(plaqs))) if len(plaqs) > 1 else 0.0,
    }
    if args.potential:
        summary["ape_alpha"] = args.ape_alpha
        summary["ape_steps"] = args.ape_steps
        summary["Rmax"] = args.Rmax
        summary["Tmax"] = args.Tmax
        # mean Wilson loops + stderr
        wmean = {}
        for r in range(1, args.Rmax + 1):
            for t in range(1, args.Tmax + 1):
                vals = np.array([row[f"W_{r}x{t}"] for row in rows])
                wmean[f"W_{r}x{t}"] = float(np.mean(vals))
                wmean[f"W_{r}x{t}_stderr"] = float(np.std(vals, ddof=1) / math.sqrt(len(vals))) if len(vals) > 1 else 0.0
        summary["wilson_loops"] = wmean
    if args.glueball:
        summary["ape_alpha"] = args.ape_alpha
        summary["glue_levels"] = glue_levels

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    tag = f"su3fast_L{L}_b{args.beta}_seed{rng_seed}"
    if args.glueball and glue_data:
        arr = np.stack(glue_data, axis=0)  # (nconfig, nlevel, L)
        np.savez(outdir / f"{tag}_glueball.npz", O=arr, levels=np.array(glue_levels),
                 beta=args.beta, L=L)
    with (outdir / f"{tag}_summary.json").open("w") as f:
        json.dump(summary, f, indent=2)
    # CSV of per-measurement rows
    keys = list(rows[0].keys())
    with (outdir / f"{tag}_measurements.csv").open("w") as f:
        f.write(",".join(keys) + "\n")
        for r in rows:
            f.write(",".join(str(r.get(k, "")) for k in keys) + "\n")
    print(json.dumps({k: v for k, v in summary.items() if k != "wilson_loops"}, indent=2), flush=True)
    return summary


# reference plaquette values (infinite-volume, published) for validation
_REF_PLAQ = {5.5: 0.5276, 5.6: 0.5408, 5.7: 0.5491, 5.8: 0.5670, 6.0: 0.5937, 6.2: 0.6136}


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--L", type=int, default=8)
    p.add_argument("--beta", type=float, default=6.0)
    p.add_argument("--thermal", type=int, default=200)
    p.add_argument("--measurements", type=int, default=100)
    p.add_argument("--skip", type=int, default=5)
    p.add_argument("--eps", type=float, default=0.30)
    p.add_argument("--nhit", type=int, default=4)
    p.add_argument("--n_or", type=int, default=3)
    p.add_argument("--seed", type=int, default=1877)
    p.add_argument("--potential", action="store_true")
    p.add_argument("--parallel", action="store_true",
                   help="use the red-black checkerboard parallel gauge update (saturates cores on one lattice)")
    p.add_argument("--glueball", action="store_true")
    p.add_argument("--glue_levels", type=str, default="0,8,16,24",
                   help="comma list of cumulative APE smearing levels for the glueball operator basis")
    p.add_argument("--Rmax", type=int, default=6)
    p.add_argument("--Tmax", type=int, default=6)
    p.add_argument("--ape_alpha", type=float, default=0.5)
    p.add_argument("--ape_steps", type=int, default=15)
    p.add_argument("--outdir", type=str, default="results")
    p.add_argument("--verbose", action="store_true")
    p.add_argument("--validate", action="store_true",
                   help="print plaquette vs published reference value")
    args = p.parse_args()
    s = run(args)
    if args.validate:
        ref = _REF_PLAQ.get(round(args.beta, 2))
        meas = s["plaquette_mean"]
        err = s["plaquette_stderr"]
        print("\n=== PLAQUETTE VALIDATION ===", flush=True)
        print(f"beta={args.beta}  L={args.L}^4", flush=True)
        print(f"measured <P> = {meas:.5f} +/- {err:.5f}", flush=True)
        if ref is not None:
            dev = (meas - ref)
            print(f"published   <P> ~ {ref:.5f}", flush=True)
            print(f"deviation    = {dev:+.5f}  ({100*dev/ref:+.2f}%)", flush=True)
            print("NOTE: finite-volume; expect small positive deviation at small L.", flush=True)
        else:
            print("(no reference value tabulated for this beta)", flush=True)
