#!/usr/bin/env python3
"""
Small pure-gauge SU(3) Wilson-action Metropolis smoke test.

This is not a production lattice-QCD code. It is intentionally compact,
reproducible, and suitable for verifying the DFD-QCD boundary-program
workflow on a CPU-only laptop. It measures plaquettes and a few Wilson
loops on tiny lattices.
"""
from __future__ import annotations
import argparse, json, math, time
from pathlib import Path
from dataclasses import dataclass, asdict
import numpy as np

ND = 4
NC = 3


def su3_project(M: np.ndarray) -> np.ndarray:
    """Project a complex 3x3 matrix to SU(3) via QR and determinant correction."""
    Q, R = np.linalg.qr(M)
    # Make the QR decomposition deterministic with positive real diagonals.
    diag = np.diag(R)
    phases = diag / np.where(np.abs(diag) == 0, 1.0, np.abs(diag))
    Q = Q @ np.diag(np.conj(phases))
    det = np.linalg.det(Q)
    Q = Q / det ** (1.0 / 3.0)
    return Q.astype(np.complex128)


def random_su3_near_identity(rng: np.random.Generator, eps: float) -> np.ndarray:
    X = rng.normal(size=(NC, NC)) + 1j * rng.normal(size=(NC, NC))
    H = X - X.conj().T
    H -= np.trace(H) / NC * np.eye(NC, dtype=np.complex128)
    return su3_project(np.eye(NC, dtype=np.complex128) + eps * H)


def shift_site(x, mu, step, L):
    y = list(x)
    y[mu] = (y[mu] + step) % L
    return tuple(y)


def make_lattice(L: int) -> np.ndarray:
    U = np.zeros((L, L, L, L, ND, NC, NC), dtype=np.complex128)
    I = np.eye(NC, dtype=np.complex128)
    for idx in np.ndindex((L, L, L, L, ND)):
        U[idx] = I
    return U


def staple(U: np.ndarray, x, mu: int) -> np.ndarray:
    L = U.shape[0]
    S = np.zeros((NC, NC), dtype=np.complex128)
    for nu in range(ND):
        if nu == mu:
            continue
        x_plus_nu = shift_site(x, nu, +1, L)
        x_plus_mu = shift_site(x, mu, +1, L)
        x_minus_nu = shift_site(x, nu, -1, L)
        x_plus_mu_minus_nu = shift_site(x_plus_mu, nu, -1, L)
        # Forward staple V_f such that Re Tr[U_mu(x) V_f] is the forward plaquette.
        S += U[x_plus_mu + (nu,)] @ U[x_plus_nu + (mu,)].conj().T @ U[x + (nu,)].conj().T
        # Backward staple V_b for the plaquette in the negative-nu plane.
        S += U[x_plus_mu_minus_nu + (nu,)].conj().T @ U[x_minus_nu + (mu,)].conj().T @ U[x_minus_nu + (nu,)]
    return S


def sweep(U: np.ndarray, beta: float, eps: float, rng: np.random.Generator) -> tuple[int, int]:
    L = U.shape[0]
    accepted = 0
    total = 0
    for x in np.ndindex((L, L, L, L)):
        for mu in range(ND):
            S = staple(U, x, mu)
            old = U[x + (mu,)]
            proposal = random_su3_near_identity(rng, eps) @ old
            old_re = np.real(np.trace(old @ S))
            new_re = np.real(np.trace(proposal @ S))
            dS = -(beta / NC) * (new_re - old_re)
            if dS <= 0.0 or rng.random() < math.exp(-dS):
                U[x + (mu,)] = proposal
                accepted += 1
            total += 1
    return accepted, total


def wilson_loop(U: np.ndarray, R: int, T: int, mu: int = 0, nu: int = 1) -> float:
    L = U.shape[0]
    acc = 0.0
    count = 0
    for x in np.ndindex((L, L, L, L)):
        M = np.eye(NC, dtype=np.complex128)
        y = x
        for _ in range(R):
            M = M @ U[y + (mu,)]
            y = shift_site(y, mu, +1, L)
        for _ in range(T):
            M = M @ U[y + (nu,)]
            y = shift_site(y, nu, +1, L)
        for _ in range(R):
            y = shift_site(y, mu, -1, L)
            M = M @ U[y + (mu,)].conj().T
        for _ in range(T):
            y = shift_site(y, nu, -1, L)
            M = M @ U[y + (nu,)].conj().T
        acc += np.real(np.trace(M)) / NC
        count += 1
    return acc / count


def measure(U: np.ndarray, max_R: int, max_T: int) -> dict:
    loops = {}
    for R in range(1, max_R + 1):
        for T in range(1, max_T + 1):
            loops[f"W_{R}x{T}"] = float(wilson_loop(U, R, T))
    plaquette = loops["W_1x1"]
    return {"plaquette": plaquette, **loops}


def estimate_potential_and_sigma(rows: list[dict], max_R: int, max_T: int) -> dict:
    # Simple effective potential V(R) = -log(W(R,T+1)/W(R,T)), using averaged loops.
    mean = {}
    for key in rows[0].keys():
        if key.startswith("W_") or key == "plaquette":
            vals = [r[key] for r in rows if np.isfinite(r[key])]
            mean[key] = float(np.mean(vals))
            mean[key + "_stderr"] = float(np.std(vals, ddof=1) / math.sqrt(len(vals))) if len(vals) > 1 else 0.0
    potentials = {}
    for R in range(1, max_R + 1):
        key1 = f"W_{R}x1"
        key2 = f"W_{R}x2"
        if key1 in mean and key2 in mean and mean[key1] > 0 and mean[key2] > 0:
            potentials[str(R)] = float(-math.log(mean[key2] / mean[key1]))
    # crude string slope from V(2)-V(1), lattice units
    sigma_lat = None
    if "1" in potentials and "2" in potentials:
        sigma_lat = float(max(0.0, potentials["2"] - potentials["1"]))
    return {"means": mean, "effective_potential": potentials, "sigma_lattice_units_crude": sigma_lat}


def run(args):
    rng = np.random.default_rng(args.seed)
    U = make_lattice(args.L)
    log = []
    t0 = time.time()
    acc_hist = []
    for s in range(args.thermal):
        acc, tot = sweep(U, args.beta, args.eps, rng)
        acc_hist.append(acc / tot)
        if args.verbose and (s + 1) % max(1, args.thermal // 5) == 0:
            print(f"thermal {s+1}/{args.thermal}: acceptance={acc/tot:.3f}")
    rows = []
    for s in range(args.measurements):
        for _ in range(args.skip):
            acc, tot = sweep(U, args.beta, args.eps, rng)
            acc_hist.append(acc / tot)
        obs = measure(U, args.max_R, args.max_T)
        obs.update({"measurement": s, "beta": args.beta, "L": args.L, "acceptance_last": acc / tot})
        rows.append(obs)
        if args.verbose and (s + 1) % max(1, args.measurements // 5) == 0:
            print(f"measure {s+1}/{args.measurements}: plaq={obs['plaquette']:.5f}, acceptance={acc/tot:.3f}")
    analysis = estimate_potential_and_sigma(rows, args.max_R, args.max_T)
    summary = {
        "code_status": "smoke_test_not_production_lattice_qcd",
        "gauge_group": "SU(3)",
        "action": "Wilson plaquette pure gauge",
        "L": args.L,
        "dimensions": [args.L] * 4,
        "beta": args.beta,
        "thermal_sweeps": args.thermal,
        "measurements": args.measurements,
        "skip_sweeps_between_measurements": args.skip,
        "proposal_eps": args.eps,
        "seed": args.seed,
        "mean_acceptance": float(np.mean(acc_hist)),
        "elapsed_sec": time.time() - t0,
        **analysis,
    }
    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    with (outdir / f"su3_L{args.L}_b{args.beta}_seed{args.seed}_measurements.csv").open("w") as f:
        keys = list(rows[0].keys())
        f.write(",".join(keys) + "\n")
        for r in rows:
            f.write(",".join(str(r[k]) for k in keys) + "\n")
    with (outdir / f"su3_L{args.L}_b{args.beta}_seed{args.seed}_summary.json").open("w") as f:
        json.dump(summary, f, indent=2)
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--L", type=int, default=4)
    p.add_argument("--beta", type=float, default=5.6)
    p.add_argument("--thermal", type=int, default=20)
    p.add_argument("--measurements", type=int, default=20)
    p.add_argument("--skip", type=int, default=1)
    p.add_argument("--eps", type=float, default=0.22)
    p.add_argument("--seed", type=int, default=1877)
    p.add_argument("--max_R", type=int, default=2)
    p.add_argument("--max_T", type=int, default=2)
    p.add_argument("--outdir", type=str, default="results")
    p.add_argument("--verbose", action="store_true")
    run(p.parse_args())
