#!/usr/bin/env python3
"""
Static-potential / string-tension analysis for the DFD-QCD pure-gauge runs.

Pipeline per beta:
  1. Read per-measurement Wilson loops W(R,T) from the *_measurements.csv.
  2. Static potential V(R): linear fit of ln W(R,T) vs T over a stable window
     (W(R,T) ~ c(R) exp[-V(R) T]); slope = -V(R).
  3. Cornell fit V(R) = V0 + sigma*R - e/R  (linear LS in basis [1, R, 1/R]);
     'sigma' is the string tension in lattice units (sigma a^2).
  4. Errors by jackknife over the measurement ensemble (binned).
  5. Scale setting: Necco-Sommer (hep-lat/0108008) r0/a(beta) parametrisation
     for the Wilson gauge action with r0 = 0.5 fm -> a in fm -> sqrt(sigma) in
     MeV, plus the scale-independent product r0*sqrt(sigma).

Honest boundary: pure gauge (quenched), single lattice spacing per beta, finite
volume.  The Necco-Sommer r0 scale is an EXTERNAL experimental calibration, not
a DFD prediction; it is the standard way to attach MeV to a lattice number.
"""
from __future__ import annotations
import argparse, json, math, glob, os
import numpy as np

HBARC = 197.3269804  # MeV*fm
R0_FM = 0.5          # Sommer scale convention


def a_over_r0(beta):
    """Necco-Sommer parametrisation, valid 5.7 <= beta <= 6.92 (Wilson gauge)."""
    b = beta - 6.0
    ln = -1.6804 - 1.7331 * b + 0.7849 * b * b - 0.4428 * b * b * b
    return math.exp(ln)


def load_csv(path):
    with open(path) as f:
        header = f.readline().strip().split(",")
        data = {h: [] for h in header}
        for line in f:
            parts = line.strip().split(",")
            if len(parts) != len(header):
                continue
            for h, p in zip(header, parts):
                try:
                    data[h].append(float(p))
                except ValueError:
                    data[h].append(np.nan)
    return {h: np.array(v) for h, v in data.items()}


def potential_from_W(Wmean, Rmax, Tmax, tmin=2, tmax=4):
    """V(R) from the effective-mass plateau:
       V_eff(R,T) = ln[W(R,T)/W(R,T+1)], averaged over the plateau window
       T in [tmin, tmax].  APE smearing produces an early, flat plateau, so a
       fixed intermediate window is robust and avoids large-T noise."""
    V = {}
    for R in range(1, Rmax + 1):
        vals = []
        for T in range(tmin, min(tmax, Tmax - 1) + 1):
            a = Wmean.get((R, T), np.nan)
            b = Wmean.get((R, T + 1), np.nan)
            if np.isfinite(a) and np.isfinite(b) and a > 1e-9 and b > 1e-9:
                vals.append(math.log(a / b))
        if vals:
            V[R] = float(np.mean(vals))
    return V


def cornell_fit(V, Rmin, Rmax):
    """Fit V(R) = c0 + sigma*R - e/R. Returns (sigma, e, c0)."""
    Rs = [R for R in range(Rmin, Rmax + 1) if R in V]
    if len(Rs) < 3:
        return None
    Rs = np.array(Rs, float)
    y = np.array([V[int(R)] for R in Rs])
    A = np.vstack([np.ones_like(Rs), Rs, -1.0 / Rs]).T
    coef, *_ = np.linalg.lstsq(A, y, rcond=None)
    c0, sigma, e = coef
    return sigma, e, c0


def analyze_one(csv_path, beta, Rmax, Tmax, tmin, Rmin, Rfit, nbin):
    d = load_csv(csv_path)
    nmeas = len(d["plaquette"])
    # collect per-measurement W arrays
    Wkeys = {}
    for R in range(1, Rmax + 1):
        for T in range(1, Tmax + 1):
            k = f"W_{R}x{T}"
            if k in d:
                Wkeys[(R, T)] = d[k]

    def sigma_from_indices(idx):
        Wmean = {rt: float(np.mean(arr[idx])) for rt, arr in Wkeys.items()}
        V = potential_from_W(Wmean, Rmax, Tmax, tmin, tmin + 2)
        fit = cornell_fit(V, Rmin, Rfit)
        return (fit[0] if fit else np.nan), V

    full_sigma, Vfull = sigma_from_indices(np.arange(nmeas))

    # jackknife over bins
    nbin = min(nbin, nmeas)
    binsz = nmeas // nbin
    jk = []
    for b in range(nbin):
        mask = np.ones(nmeas, bool)
        mask[b * binsz:(b + 1) * binsz] = False
        s, _ = sigma_from_indices(np.where(mask)[0])
        if np.isfinite(s):
            jk.append(s)
    jk = np.array(jk)
    if len(jk) > 1:
        sigma_mean = float(np.mean(jk))
        sigma_err = float(math.sqrt((len(jk) - 1) / len(jk) * np.sum((jk - np.mean(jk))**2)))
    else:
        sigma_mean, sigma_err = float(full_sigma), float("nan")

    sigma = full_sigma if np.isfinite(full_sigma) else sigma_mean
    sqrt_sigma_lat = math.sqrt(sigma) if sigma > 0 else float("nan")
    # propagate error: d(sqrt s)=ds/(2 sqrt s)
    sqrt_sigma_lat_err = sigma_err / (2 * sqrt_sigma_lat) if sigma > 0 else float("nan")

    # scale setting
    aR0 = a_over_r0(beta)
    a_fm = R0_FM * aR0
    sqrt_sigma_MeV = sqrt_sigma_lat * HBARC / a_fm
    sqrt_sigma_MeV_err = sqrt_sigma_lat_err * HBARC / a_fm
    r0_sqrt_sigma = sqrt_sigma_lat / aR0  # = sqrt(sigma a^2) * (r0/a)

    return {
        "beta": beta,
        "n_measurements": nmeas,
        "csv": os.path.basename(csv_path),
        "V_of_R": {str(R): float(Vfull[R]) for R in sorted(Vfull)},
        "sigma_lattice": float(sigma),
        "sigma_lattice_jk_err": sigma_err,
        "sqrt_sigma_lattice": sqrt_sigma_lat,
        "sqrt_sigma_lattice_err": sqrt_sigma_lat_err,
        "a_over_r0_NeccoSommer": aR0,
        "a_fm": a_fm,
        "sqrt_sigma_MeV": sqrt_sigma_MeV,
        "sqrt_sigma_MeV_err": sqrt_sigma_MeV_err,
        "r0_sqrt_sigma_dimensionless": r0_sqrt_sigma,
        "fit_window": {"tmin": tmin, "Rmin": Rmin, "Rfit": Rfit},
    }


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--glob", default="results_prod/su3fast_L*_measurements.csv")
    p.add_argument("--Rmax", type=int, default=8)
    p.add_argument("--Tmax", type=int, default=8)
    p.add_argument("--tmin", type=int, default=2)
    p.add_argument("--Rmin", type=int, default=2)
    p.add_argument("--Rfit", type=int, default=6)
    p.add_argument("--nbin", type=int, default=20)
    p.add_argument("--out", default="results_prod/string_tension_analysis.json")
    args = p.parse_args()

    files = sorted(glob.glob(args.glob))
    if not files:
        print("no files match", args.glob)
        return
    results = []
    for f in files:
        # recover beta from companion summary
        summ = f.replace("_measurements.csv", "_summary.json")
        beta = None
        if os.path.exists(summ):
            beta = json.load(open(summ)).get("beta")
        if beta is None:
            continue
        try:
            r = analyze_one(f, beta, args.Rmax, args.Tmax, args.tmin, args.Rmin, args.Rfit, args.nbin)
            results.append(r)
        except Exception as e:
            print(f"  [skip {os.path.basename(f)}: {e}]")
    results.sort(key=lambda r: r["beta"])

    print("\n=== DFD-QCD pure-gauge string tension (quenched, finite volume) ===")
    print("DFD ledger target:  sqrt(sigma) = 440 +/- 25 MeV   [r0=0.5 fm Sommer scale]")
    print(f"{'beta':>5} {'a[fm]':>7} {'sqrt(sig)a':>11} {'r0*sqrt(sig)':>13} {'sqrt(sig)[MeV]':>16}")
    for r in results:
        print(f"{r['beta']:>5.2f} {r['a_fm']:>7.4f} "
              f"{r['sqrt_sigma_lattice']:>7.4f}    "
              f"{r['r0_sqrt_sigma_dimensionless']:>10.3f}    "
              f"{r['sqrt_sigma_MeV']:>7.0f} +/- {r['sqrt_sigma_MeV_err']:>4.0f}")
    print("\nReference (continuum, world pure-gauge):  r0*sqrt(sigma) ~ 1.16")
    print("                                          sqrt(sigma)    ~ 440-465 MeV")

    json.dump(results, open(args.out, "w"), indent=2)
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()
