#!/usr/bin/env python3
"""
Plots for the DFD-QCD pure-gauge string-tension program:
  1. Static potential V(R) vs R with Cornell fit, per beta.
  2. sqrt(sigma) in MeV vs beta, with the DFD ledger target band 440+/-25 MeV.
  3. Plaquette validation: measured vs published <P>(beta).
Reads results_prod/string_tension_analysis.json and the *_summary.json files.
"""
from __future__ import annotations
import json, glob, os, math
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

PROD = "results_prod"
FIG = "figures"
os.makedirs(FIG, exist_ok=True)


def cornell(R, c0, sigma, e):
    return c0 + sigma * R - e / R


def plot_potentials(results):
    plt.figure(figsize=(7, 5))
    colors = plt.cm.viridis(np.linspace(0, 0.85, len(results)))
    for r, col in zip(results, colors):
        V = r["V_of_R"]
        Rs = np.array(sorted(int(k) for k in V))
        Vs = np.array([V[str(R)] for R in Rs])
        plt.plot(Rs, Vs, "o", color=col, label=f"$\\beta={r['beta']}$ (a={r['a_fm']:.3f} fm)")
        # overlay Cornell fit using fitted sigma; refit c0,e quickly for the curve
        Rfit = r["fit_window"]["Rfit"]; Rmin = r["fit_window"]["Rmin"]
        mask = (Rs >= Rmin) & (Rs <= Rfit)
        if mask.sum() >= 3:
            A = np.vstack([np.ones(mask.sum()), Rs[mask], -1.0 / Rs[mask]]).T
            coef, *_ = np.linalg.lstsq(A, Vs[mask], rcond=None)
            xx = np.linspace(Rs[mask].min(), Rs[mask].max(), 50)
            plt.plot(xx, coef[0] + coef[1] * xx - coef[2] / xx, "-", color=col, lw=1.2)
    plt.xlabel("R  (lattice units)")
    plt.ylabel("V(R)  (lattice units)")
    plt.title("Static quark potential — DFD-QCD pure gauge SU(3), $16^4$\n"
              "points: lattice; lines: Cornell fit $V_0+\\sigma R-e/R$")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(f"{FIG}/static_potential.png", dpi=140)
    plt.close()


def plot_sqrt_sigma(results):
    betas = [r["beta"] for r in results]
    ss = [r["sqrt_sigma_MeV"] for r in results]
    er = [r["sqrt_sigma_MeV_err"] for r in results]
    plt.figure(figsize=(7, 5))
    plt.axhspan(440 - 25, 440 + 25, color="orange", alpha=0.25,
                label="DFD ledger target $440\\pm25$ MeV")
    plt.axhline(440, color="orange", ls="--", lw=1)
    plt.errorbar(betas, ss, yerr=er, fmt="s", color="navy", capsize=4,
                 label="lattice (Sommer $r_0=0.5$ fm)")
    plt.xlabel("$\\beta = 6/g^2$")
    plt.ylabel("$\\sqrt{\\sigma}$  (MeV)")
    plt.title("String tension vs lattice coupling\n"
              "(scale set by Necco-Sommer $r_0/a$; quenched, finite volume)")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(f"{FIG}/sqrt_sigma_vs_beta.png", dpi=140)
    plt.close()


def plot_plaquette_validation():
    ref = {5.5: 0.5276, 5.6: 0.5408, 5.7: 0.5491, 5.8: 0.5670, 6.0: 0.5937, 6.2: 0.6136}
    meas = []
    for f in glob.glob(f"{PROD}/su3fast_L*_summary.json") + glob.glob("results_validate/su3fast_L*_summary.json"):
        d = json.load(open(f))
        meas.append((d["beta"], d["plaquette_mean"], d.get("plaquette_stderr", 0), d["L"]))
    if not meas:
        return
    bx = np.linspace(5.45, 6.25, 100)
    plt.figure(figsize=(7, 5))
    rb = sorted(ref)
    plt.plot(rb, [ref[b] for b in rb], "k-", label="published $\\langle P\\rangle$")
    for b, m, e, L in sorted(meas):
        plt.errorbar([b], [m], yerr=[max(e, 1e-5)], fmt="o", capsize=4,
                     label=f"this code $L={L}$")
    plt.xlabel("$\\beta$")
    plt.ylabel("$\\langle P\\rangle$  average plaquette")
    plt.title("Plaquette validation vs published SU(3) Wilson values")
    plt.legend(fontsize=8)
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(f"{FIG}/plaquette_validation.png", dpi=140)
    plt.close()


def plot_continuum():
    cpath = f"{PROD}/continuum_extrap.json"
    if not os.path.exists(cpath):
        return
    d = json.load(open(cpath))
    rows = sorted(d["per_beta"], key=lambda r: r["beta"])
    xs, ys, es = [], [], []
    for r in rows:
        ar0 = r["a_over_r0"]; s = r["fixe26"]["sigma"]; se = r["fixe26"]["sigma_err"]
        if s <= 0:
            continue
        xs.append(ar0 ** 2); ys.append(math.sqrt(s) / ar0); es.append(se / (2 * math.sqrt(s)) / ar0)
    xs = np.array(xs); ys = np.array(ys); es = np.array(es)
    c0, c0e = d["continuum_r0_sqrt_sigma"]; c1 = d["slope"]
    plt.figure(figsize=(7, 5))
    plt.errorbar(xs, ys, yerr=es, fmt="o", color="navy", capsize=4, label="lattice (fixed-e Cornell)")
    xx = np.linspace(0, max(xs) * 1.05, 50)
    plt.plot(xx, c0 + c1 * xx, "b--", lw=1, label=f"linear $(a/r_0)^2$ fit")
    plt.errorbar([0], [c0], yerr=[c0e], fmt="s", color="crimson", capsize=5,
                 label=f"continuum {c0:.3f}({int(c0e*1000)})")
    plt.axhspan(1.193 - 0.010, 1.193 + 0.010, color="green", alpha=0.2,
                label="Necco-Sommer 1.193(10)")
    plt.xlabel("$(a/r_0)^2$"); plt.ylabel("$r_0\\sqrt{\\sigma}$")
    plt.title("Continuum trend of $r_0\\sqrt{\\sigma}$ (pure gauge SU(3), $16^4$)")
    plt.legend(fontsize=8); plt.grid(alpha=0.3); plt.tight_layout()
    plt.savefig(f"{FIG}/continuum_r0sqrtsigma.png", dpi=140); plt.close()
    print("wrote continuum_r0sqrtsigma.png")


def plot_hmc_plaquette():
    import glob as _g
    dyn = {}
    for f in _g.glob("results_hmc/wilsonhmc_*_summary.json"):
        d = json.load(open(f))
        if d.get("plaquette_dynamical_mean"):
            dyn[d["kappa"]] = (d["plaquette_dynamical_mean"], d.get("plaquette_dynamical_stderr") or 0)
    qf = "results_hmc/su3fast_L4_b5.6_seed561_summary.json"
    if not os.path.exists(qf) or not dyn:
        return
    q = json.load(open(qf))["plaquette_mean"]
    plt.figure(figsize=(6.5, 5))
    plt.axhline(q, color="gray", ls="--", label=f"quenched <P>={q:.4f}")
    ks = sorted(dyn)
    plt.errorbar(ks, [dyn[k][0] for k in ks], yerr=[dyn[k][1] for k in ks],
                 fmt="o", color="darkred", capsize=4, label="Nf=2 dynamical")
    plt.xlabel("$\\kappa$ (sea-quark hopping parameter)"); plt.ylabel("$\\langle P\\rangle$")
    plt.title("Sea-quark plaquette rise — dynamical Nf=2 Wilson HMC\n($4^4$, $\\beta=5.6$)")
    plt.legend(); plt.grid(alpha=0.3); plt.tight_layout()
    plt.savefig(f"{FIG}/hmc_plaquette_rise.png", dpi=140); plt.close()
    print("wrote hmc_plaquette_rise.png")


def main():
    apath = f"{PROD}/string_tension_analysis.json"
    if os.path.exists(apath):
        results = json.load(open(apath))
        results.sort(key=lambda r: r["beta"])
        plot_potentials(results)
        plot_sqrt_sigma(results)
        print("wrote static_potential.png, sqrt_sigma_vs_beta.png")
    plot_plaquette_validation()
    print("wrote plaquette_validation.png")
    plot_continuum()
    plot_hmc_plaquette()


if __name__ == "__main__":
    main()
