#!/usr/bin/env python3
"""
Continuum extrapolation + independent cross-checks for the pure-gauge string tension.

Addresses the review gates:
  * GATE B  per-beta disaggregation vs literature a*sqrt(sigma)
  * GATE C  r_min stability scan of sigma
  * GATE G  fix e = pi/12 (Luscher) to remove the sigma-e anticorrelation
  * GATE K  continuum trend: extrapolate r0*sqrt(sigma) linearly in (a/r0)^2 -> 1.193(10)

For each beta we build V(R) from the effective-mass plateau (jackknife over configs)
and fit the Cornell form three ways:
  (free26)   V0 + sigma R - e/R, free e, R in [2,6]
  (free36)   same, R in [3,6]
  (fixe26)   V0 + sigma R - (pi/12)/R, e FIXED, R in [2,6]   <- most robust
The fixe variant is used for the continuum extrapolation.
"""
from __future__ import annotations
import json, glob, os, math
import numpy as np

HBARC = 197.3269804
R0_FM = 0.5
E_STRING = math.pi / 12.0  # 0.2618, Luscher coefficient
# Necco-Sommer continuum value
R0SQ_CONT_REF = (1.193, 0.010)
# literature a*sqrt(sigma) (Wilson action) for the per-beta gate
LIT_ASQRTSIG = {5.8: 0.3133, 5.85: 0.2874, 6.0: 0.2189, 6.2: 0.1610}


def a_over_r0(beta):
    b = beta - 6.0
    return math.exp(-1.6804 - 1.7331 * b + 0.7849 * b * b - 0.4428 * b * b * b)


def load_csv(path):
    with open(path) as f:
        header = f.readline().strip().split(",")
        cols = {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:
                    cols[h].append(float(p))
                except ValueError:
                    cols[h].append(np.nan)
    return {h: np.array(v) for h, v in cols.items()}


def potential(Wmean, Rmax, Tmax, tmin=2, tmax=4):
    V = {}
    for R in range(1, Rmax + 1):
        vals = []
        for T in range(tmin, min(tmax, Tmax - 1) + 1):
            a = Wmean.get((R, T)); b = Wmean.get((R, T + 1))
            if a and 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(V, Rmin, Rmax, fix_e=None):
    Rs = [R for R in range(Rmin, Rmax + 1) if R in V]
    if len(Rs) < (2 if fix_e is not None else 3):
        return None
    Rs = np.array(Rs, float)
    y = np.array([V[int(R)] for R in Rs])
    if fix_e is None:
        A = np.vstack([np.ones_like(Rs), Rs, -1.0 / Rs]).T
        coef, *_ = np.linalg.lstsq(A, y, rcond=None)
        return {"sigma": coef[1], "e": coef[2], "V0": coef[0]}
    else:
        yy = y + fix_e / Rs  # V = V0 + sigma R - e/R  ->  V + e/R = V0 + sigma R
        A = np.vstack([np.ones_like(Rs), Rs]).T
        coef, *_ = np.linalg.lstsq(A, yy, rcond=None)
        return {"sigma": coef[1], "e": fix_e, "V0": coef[0]}


def analyze_beta(csv, beta, Rmax=8, Tmax=8, nbin=20):
    d = load_csv(csv)
    n = len(d["plaquette"])
    Wk = {(R, T): d[f"W_{R}x{T}"] for R in range(1, Rmax + 1) for T in range(1, Tmax + 1) if f"W_{R}x{T}" in d}

    def sigma_variant(idx, Rmin, Rfit, fix_e):
        Wm = {rt: float(np.mean(arr[idx])) for rt, arr in Wk.items()}
        V = potential(Wm, Rmax, Tmax)
        fit = cornell(V, Rmin, Rfit, fix_e)
        return fit["sigma"] if fit else np.nan

    variants = {
        "free26": dict(Rmin=2, Rfit=6, fix_e=None),
        "free36": dict(Rmin=3, Rfit=6, fix_e=None),
        "fixe26": dict(Rmin=2, Rfit=6, fix_e=E_STRING),
        "fixe36": dict(Rmin=3, Rfit=6, fix_e=E_STRING),
    }
    out = {"beta": beta, "n": n, "a_over_r0": a_over_r0(beta)}
    binsz = max(1, n // nbin)
    for name, kw in variants.items():
        full = sigma_variant(np.arange(n), **kw)
        jk = []
        for b in range(min(nbin, n)):
            mask = np.ones(n, bool); mask[b * binsz:(b + 1) * binsz] = False
            s = sigma_variant(np.where(mask)[0], **kw)
            if np.isfinite(s):
                jk.append(s)
        jk = np.array(jk)
        err = math.sqrt((len(jk) - 1) / len(jk) * np.sum((jk - jk.mean()) ** 2)) if len(jk) > 1 else float("nan")
        out[name] = {"sigma": float(full), "sigma_err": float(err)}
    # r_min scan (free e)
    out["rmin_scan"] = {}
    for rmin in [2, 3, 4]:
        s = sigma_variant(np.arange(n), rmin, 6, None)
        out["rmin_scan"][rmin] = float(s) if np.isfinite(s) else None
    return out


def main():
    files = sorted(glob.glob("results_prod/su3fast_L16_b*_measurements.csv"))
    rows = []
    for f in files:
        summ = f.replace("_measurements.csv", "_summary.json")
        beta = json.load(open(summ))["beta"]
        rows.append(analyze_beta(f, beta))
    rows.sort(key=lambda r: r["beta"])

    print("\n=== Per-beta string tension: fit-window cross-checks (a*sqrt(sigma)) ===")
    print(f"{'beta':>5} {'free[2,6]':>12} {'free[3,6]':>12} {'fixe[2,6]':>12} {'literature':>11}")
    for r in rows:
        b = r["beta"]
        def asq(v):
            s = r[v]["sigma"]; e = r[v]["sigma_err"]
            return f"{math.sqrt(s):.4f}({int(round(e/(2*math.sqrt(s))*1e4)):d})" if s > 0 else "  --  "
        lit = LIT_ASQRTSIG.get(b)
        lits = f"{lit:.4f}" if lit else "(interp)"
        print(f"{b:>5.2f} {asq('free26'):>12} {asq('free36'):>12} {asq('fixe26'):>12} {lits:>11}")

    print("\n=== r_min stability of a*sqrt(sigma) (free e), GATE C ===")
    for r in rows:
        sc = r["rmin_scan"]
        vals = "  ".join(f"rmin={k}:{math.sqrt(v):.4f}" if v and v > 0 else f"rmin={k}:--" for k, v in sc.items())
        print(f" beta={r['beta']:.2f}:  {vals}")

    # continuum extrapolation using fixe26 (most robust)
    print("\n=== Continuum extrapolation of r0*sqrt(sigma), linear in (a/r0)^2 ===")
    xs, ys, es = [], [], []
    for r in rows:
        ar0 = r["a_over_r0"]
        s = r["fixe26"]["sigma"]; serr = r["fixe26"]["sigma_err"]
        if s <= 0:
            continue
        rsq = math.sqrt(s) / ar0          # r0*sqrt(sigma)
        rsq_err = serr / (2 * math.sqrt(s)) / ar0
        x = ar0 ** 2                      # (a/r0)^2
        xs.append(x); ys.append(rsq); es.append(rsq_err)
        print(f" beta={r['beta']:.2f}  (a/r0)^2={x:.4f}  r0*sqrt(sigma)={rsq:.3f} +/- {rsq_err:.3f}")
    xs = np.array(xs); ys = np.array(ys); es = np.array(es)
    # weighted linear fit y = c0 + c1 x  ; intercept c0 = continuum
    W = 1.0 / es ** 2
    A = np.vstack([np.ones_like(xs), xs]).T
    AtW = A.T * W
    cov = np.linalg.inv(AtW @ A)
    coef = cov @ (AtW @ ys)
    c0, c1 = coef
    c0_err = math.sqrt(cov[0, 0])
    chi2 = np.sum(W * (ys - (c0 + c1 * xs)) ** 2)
    dof = len(xs) - 2
    print(f"\n CONTINUUM  r0*sqrt(sigma) = {c0:.3f} +/- {c0_err:.3f}   "
          f"(slope {c1:+.2f}, chi2/dof = {chi2/dof:.2f})")
    print(f" Necco-Sommer reference        = {R0SQ_CONT_REF[0]:.3f} +/- {R0SQ_CONT_REF[1]:.3f}")
    sqrt_sig_05 = c0 * HBARC / R0_FM
    sqrt_sig_055 = c0 * HBARC / 0.55
    print(f"\n sqrt(sigma) [r0=0.50 fm] = {sqrt_sig_05:.0f} +/- {c0_err*HBARC/R0_FM:.0f} MeV")
    print(f" sqrt(sigma) [r0=0.55 fm] = {sqrt_sig_055:.0f} +/- {c0_err*HBARC/0.55:.0f} MeV   (DFD-ledger scale)")
    print(" DFD ledger target          = 440 +/- 25 MeV")

    json.dump({"per_beta": rows,
               "continuum_r0_sqrt_sigma": [float(c0), float(c0_err)],
               "slope": float(c1), "chi2_dof": float(chi2 / dof),
               "sqrt_sigma_MeV_r0_0p50": float(sqrt_sig_05),
               "sqrt_sigma_MeV_r0_0p55": float(sqrt_sig_055)},
              open("results_prod/continuum_extrap.json", "w"), indent=2)
    print("\nwrote results_prod/continuum_extrap.json")


if __name__ == "__main__":
    main()
