#!/usr/bin/env python3
"""
0++ scalar glueball analysis from per-config timeslice operators (*_glueball.npz).

For each smearing level (and via a GEVP over the level basis):
  - connected correlator  C(t) = <O(t0+t) O(t0)>_{t0} - <O>^2,  averaged over
    source slices t0 and configs, symmetrised under t -> Nt - t (periodic);
  - effective mass m_eff(t) = arccosh[(C(t-1)+C(t+1)) / (2 C(t))]  (periodic form),
    falling back to log[C(t)/C(t+1)] when needed;
  - GEVP: solve C(t) v = lambda(t,t0) C(t0) v; ground-state mass from the
    leading eigenvalue lambda0(t) ~ exp(-m (t - t0));
  - jackknife errors over configs.

Scale: a from Necco-Sommer r0/a (r0 = 0.5 fm) at the run's beta, so m_0++ in MeV
and the dimensionless m_0++/sqrt(sigma) and r0*m_0++ are reported.  Pure gauge,
finite volume: a heavy 0++ on a short time extent is statistics-limited; treat
the number as an exploratory plateau estimate, not a precision spectroscopy result.
"""
from __future__ import annotations
import argparse, json, math, glob, os
import numpy as np

HBARC = 197.3269804
R0_FM = 0.5


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 connected_corr(O):
    """O shape (ncfg, Nt). Return C(t), t=0..Nt//2, vacuum-subtracted, source-averaged."""
    ncfg, Nt = O.shape
    Omean = O.mean()
    Oc = O - Omean
    # correlator per config averaged over source slice t0
    C = np.zeros(Nt)
    for dt in range(Nt):
        # <O(t0+dt) O(t0)> averaged over t0 and configs
        prod = (Oc * np.roll(Oc, -dt, axis=1)).mean()
        C[dt] = prod
    # symmetrise periodic
    Cs = np.array([0.5 * (C[t] + C[(Nt - t) % Nt]) for t in range(Nt)])
    return Cs[: Nt // 2 + 1]


def jackknife_corr(O):
    ncfg, Nt = O.shape
    full = connected_corr(O)
    samples = []
    for i in range(ncfg):
        idx = np.ones(ncfg, bool); idx[i] = False
        samples.append(connected_corr(O[idx]))
    samples = np.array(samples)
    mean = samples.mean(0)
    err = np.sqrt((ncfg - 1) * np.mean((samples - mean) ** 2, axis=0))
    return full, err


def eff_mass(C):
    Nt2 = len(C)
    m = []
    for t in range(1, Nt2 - 1):
        if C[t] > 0 and C[t - 1] + C[t + 1] > 0:
            arg = (C[t - 1] + C[t + 1]) / (2 * C[t])
            if arg >= 1:
                m.append((t, math.acosh(arg)))
            elif C[t + 1] > 0:
                m.append((t, math.log(C[t] / C[t + 1])))
    # also simple log at t=0->1
    out = {}
    for t, val in m:
        out[t] = val
    if C[0] > 0 and C[1] > 0 and 0 not in out:
        out[0] = math.log(C[0] / C[1])
    return out


def _corr_matrix(Ofull):
    """Symmetric correlation matrix Csym[t][a,b], vacuum-subtracted, time-symmetrised."""
    ncfg, nlev, Nt = Ofull.shape
    Oc = Ofull - Ofull.mean(axis=0, keepdims=True)  # remove vacuum per (level,t)
    Cmat = np.zeros((Nt, nlev, nlev))
    for a in range(nlev):
        for b in range(nlev):
            for dt in range(Nt):
                Cmat[dt, a, b] = (Oc[:, a, :] * np.roll(Oc[:, b, :], -dt, axis=1)).mean()
    # enforce a,b symmetry and time-reflection symmetry
    Cmat = 0.5 * (Cmat + np.transpose(Cmat, (0, 2, 1)))
    Csym = np.array([0.5 * (Cmat[t] + Cmat[(Nt - t) % Nt]) for t in range(Nt)])
    return Csym


def gevp_lambda0(Csym, t0=2):
    """Leading GEVP eigenvalue lambda0(t) from the symmetric whitening
       C(t0)^{-1/2} C(t) C(t0)^{-1/2}. Returns dict t-> lambda0(t)."""
    Nt = Csym.shape[0]
    out = {}
    C0 = Csym[t0]
    # symmetric positive-definite whitening via eigendecomposition of C(t0)
    w, Vv = np.linalg.eigh(C0)
    if np.any(w <= 0):
        # regularise: keep positive subspace
        keep = w > max(w) * 1e-6
        if keep.sum() == 0:
            return out
        w = w[keep]; Vv = Vv[:, keep]
    Winv = Vv @ np.diag(1.0 / np.sqrt(w)) @ Vv.T.conj()
    for t in range(t0, Nt):
        M = Winv @ Csym[t] @ Winv
        M = 0.5 * (M + M.T)
        ev = np.linalg.eigvalsh(M)
        ev = ev[ev > 0]
        if len(ev):
            out[t] = float(ev.max())  # ground state = largest eigenvalue
    return out


def gevp_masses(Ofull, t0=2):
    """Ground-state effective mass from the GEVP leading eigenvalue:
       m_eff(t) = ln[lambda0(t)/lambda0(t+1)]  (ratio form; stable)."""
    Csym = _corr_matrix(Ofull)
    lam = gevp_lambda0(Csym, t0=t0)
    masses = {}
    ts = sorted(lam)
    for t in ts:
        if (t + 1) in lam and lam[t] > 0 and lam[t + 1] > 0:
            masses[t] = math.log(lam[t] / lam[t + 1])
    return masses


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--glob", default="results_glue/*_glueball.npz")
    p.add_argument("--out", default="results_glue/glueball_analysis.json")
    args = p.parse_args()
    files = sorted(glob.glob(args.glob))
    if not files:
        print("no glueball npz files match", args.glob)
        return
    results = []
    for f in files:
        d = np.load(f)
        O = d["O"]  # (ncfg, nlevel, Nt)
        levels = list(d["levels"])
        beta = float(d["beta"]); L = int(d["L"])
        ncfg = O.shape[0]
        aR0 = a_over_r0(beta); a_fm = R0_FM * aR0
        print(f"\n=== {os.path.basename(f)}  beta={beta} L={L}^4  ncfg={ncfg} ===")
        per_level = {}
        for li, lvl in enumerate(levels):
            C, Cerr = jackknife_corr(O[:, li, :])
            em = eff_mass(C)
            per_level[str(lvl)] = {
                "C": [float(x) for x in C],
                "C_err": [float(x) for x in Cerr],
                "eff_mass": {str(t): float(v) for t, v in em.items()},
            }
            emstr = "  ".join(f"m({t})={v:.3f}" for t, v in sorted(em.items()))
            print(f"  APE={lvl:>2}: C(0..)= " + " ".join(f"{c:.3e}" for c in C[:5]))
            print(f"          eff mass: {emstr}")
        gm = gevp_masses(O, t0=1)
        gmstr = "  ".join(f"m({t})={v:.3f}" for t, v in sorted(gm.items()))
        print(f"  GEVP ground state eff mass: {gmstr}  (diagnostic; noise-prone at short Nt)")
        # PRIMARY estimate: arccosh effective-mass plateau (t=1,2) of the most-smeared operator,
        # which best overlaps the 0++ ground state. (GEVP is reported only as a cross-check;
        # on a short time extent it is noise-dominated beyond t~2.)
        top_em = per_level[str(levels[-1])]["eff_mass"]
        plat = [top_em[str(t)] for t in (1, 2) if str(t) in top_em]
        m_lat = float(np.mean(plat)) if plat else (list(top_em.values())[0] if top_em else None)
        # jackknife error on the plateau mass (heaviest smearing level)
        m_lat_err = None
        if m_lat is not None:
            top = O[:, len(levels) - 1, :]
            jkm = []
            for i in range(ncfg):
                idx = np.ones(ncfg, bool); idx[i] = False
                C = connected_corr(top[idx])
                vs = []
                for t in (1, 2):
                    arg = (C[t - 1] + C[t + 1]) / (2 * C[t]) if C[t] != 0 else 0
                    if arg > 1:
                        vs.append(math.acosh(arg))
                if vs:
                    jkm.append(np.mean(vs))
            jkm = np.array(jkm)
            if len(jkm) > 1:
                m_lat_err = float(math.sqrt((len(jkm) - 1) / len(jkm) * np.sum((jkm - jkm.mean()) ** 2)))
        rec = {"file": os.path.basename(f), "beta": beta, "L": L, "ncfg": ncfg,
               "a_fm": a_fm, "levels": [int(x) for x in levels],
               "per_level": per_level, "gevp_eff_mass": {str(t): float(v) for t, v in gm.items()}}
        if m_lat is not None and m_lat > 0:
            rec["m0pp_lattice"] = float(m_lat)
            rec["m0pp_lattice_err"] = m_lat_err
            rec["m0pp_MeV"] = float(m_lat * HBARC / a_fm)
            rec["m0pp_MeV_err"] = float(m_lat_err * HBARC / a_fm) if m_lat_err else None
            rec["r0_m0pp"] = float(m_lat / aR0)
            es = f" +/- {m_lat_err:.3f}" if m_lat_err else ""
            print(f"  -> a*m(0++) = {m_lat:.3f}{es}  =>  {rec['m0pp_MeV']:.0f} MeV   (r0*m={rec['r0_m0pp']:.2f})")
        results.append(rec)
    json.dump(results, open(args.out, "w"), indent=2)
    print(f"\nwrote {args.out}")
    print("DFD ledger target: m(0++) = 1.69 +/- 0.10 GeV;  continuum m/sqrt(sigma) ~ 3.6-3.7, r0*m ~ 4.2")


if __name__ == "__main__":
    main()
