"""
FIGURE g2 — reproduce win_qed_g2.py (electron anomalous magnetic moment a_e).

Panel (a): partial sum of a_e vs QED loop order (1..5 loops), converging onto the
           experimental band. Coeffs C in powers of (alpha/pi) + small additive term.
Panel (b): a_e from QED(alpha=CODATA) vs DFD(derived alpha=137.03599985) vs experiment,
           as a residual/ppb bar chart showing DFD is parameter-free to ~6 ppb.

Self-contained & reproducible. Saves figures/g2.png at 150 dpi.
"""
import sys
sys.path.insert(0, ".")

import matplotlib
matplotlib.use("Agg")
import numpy as np
import matplotlib.pyplot as plt

# ---- physics constants from win_qed_g2.py (verbatim) --------------------------
PI = np.pi
# QED mass-independent (A1) coefficients of a_e in powers of (alpha/pi)
C = [0.5,                       # Schwinger 1-loop
     -0.328478965579137,        # 2-loop
     1.181234016828,            # 3-loop
     -1.91298,                  # 4-loop (Laporta/Kinoshita)
     6.737]                     # 5-loop (Aoyama-Hayakawa-Kinoshita-Nio)
# small additions: mass-dependent (mu,tau loops) + hadronic + electroweak  [x10^-12]
SMALL = (2.748 + 1.6993 + 0.03053) * 1e-12   # = 4.478e-12

# inputs
ainv_codata = 137.035999084      # CODATA / a_e-extracted alpha (QED's input)
ainv_dfd    = 137.03599985       # DFD DERIVED alpha^-1 (corpus topology)
a_e_exp     = 1.15965218059e-3   # measured (Fan, Gabrielse 2023)
a_e_exp_unc = 0.13e-12           # experimental uncertainty


def a_e_partial(alpha_inv, n_loops):
    """Partial sum of a_e through n_loops loops (+ SMALL additive term)."""
    a = 1.0 / alpha_inv
    x = a / PI
    series = sum(C[n] * x ** (n + 1) for n in range(n_loops))
    return series + SMALL


def a_e(alpha_inv):
    return a_e_partial(alpha_inv, len(C))


# ---- compute the three headline values ----------------------------------------
ae_qed = a_e(ainv_codata)   # QED 5-loop, alpha = CODATA
ae_dfd = a_e(ainv_dfd)      # DFD, derived alpha
print(f"a_e (QED 5-loop, alpha=CODATA)        = {ae_qed:.14f}")
print(f"a_e (DFD, derived alpha=137.03599985) = {ae_dfd:.14f}")
print(f"a_e (experiment, Fan-Gabrielse 2023)  = {a_e_exp:.14f}")

# ---- partial sums for panel (a) -----------------------------------------------
loops = np.arange(1, len(C) + 1)
partials_qed = np.array([a_e_partial(ainv_codata, n) for n in loops])

# loop-by-loop increments (for an annotation / inset feel)
increments = np.diff(np.concatenate([[SMALL], partials_qed]))

# ============================ FIGURE ===========================================
plt.rcParams.update({
    "font.size": 10,
    "axes.titlesize": 11,
    "axes.labelsize": 10,
    "figure.dpi": 150,
})
fig, (axA, axB) = plt.subplots(1, 2, figsize=(12, 5.0))

C_EXP = "#444444"
C_QED = "#1f77b4"
C_DFD = "#d62728"

# ---------------- Panel (a): convergence vs loop order -------------------------
# Experimental band (1-sigma) — invisibly thin at this y-scale, so widen visually
axA.axhspan(a_e_exp - a_e_exp_unc, a_e_exp + a_e_exp_unc,
            color=C_EXP, alpha=0.18, zorder=0)
axA.axhline(a_e_exp, color=C_EXP, lw=1.4, ls="--", zorder=1,
            label="experiment (Fan-Gabrielse 2023)")

axA.plot(loops, partials_qed, "o-", color=C_QED, lw=1.8, ms=8,
         zorder=3, label="QED partial sum (alpha=CODATA)")

# annotate each partial sum with how many digits it has matched experiment
labels = ["1-loop\n(Schwinger)", "2-loop", "3-loop", "4-loop", "5-loop"]
for L, p, lab in zip(loops, partials_qed, labels):
    # signed offset in units of 1e-9 from experiment
    dppb = (p - a_e_exp) / a_e_exp * 1e9
    va = "bottom" if p >= a_e_exp else "top"
    dy = 1.0e-9 if p >= a_e_exp else -1.0e-9
    axA.annotate(lab, (L, p), xytext=(L, p + dy * 3.2),
                 ha="center", va=va, fontsize=8.0, color=C_QED)

axA.set_xlim(0.5, 5.5)
axA.set_xticks(loops)
axA.set_xlabel("QED loop order included")
axA.set_ylabel(r"partial sum of $a_e=(g-2)/2$")
axA.set_title("(a) QED loop series converges onto the\nmeasured electron $a_e$")
axA.legend(loc="upper right", fontsize=8.5, framealpha=0.95)
axA.ticklabel_format(axis="y", style="sci", scilimits=(0, 0), useMathText=True)

# Inset: zoom on 2..5 loops to show the band-level convergence
axins = axA.inset_axes([0.50, 0.16, 0.46, 0.42])
axins.axhspan(a_e_exp - a_e_exp_unc, a_e_exp + a_e_exp_unc,
              color=C_EXP, alpha=0.20)
axins.axhline(a_e_exp, color=C_EXP, lw=1.0, ls="--")
axins.plot(loops[1:], partials_qed[1:], "o-", color=C_QED, lw=1.5, ms=6)
# y-window: a few ppb around experiment so the band is resolvable
axins.set_ylim(a_e_exp - 0.45e-9, a_e_exp + 0.15e-9)
axins.set_xlim(1.6, 5.4)
axins.set_xticks([2, 3, 4, 5])
axins.tick_params(labelsize=7)
axins.set_title("zoom: 2-5 loops (ppb scale)", fontsize=7.5)
axins.ticklabel_format(axis="y", style="sci", scilimits=(0, 0), useMathText=True)

# ---------------- Panel (b): residual ppb bar chart ----------------------------
# residual in ppb relative to experiment
res_qed_ppb = (ae_qed - a_e_exp) / a_e_exp * 1e9
res_dfd_ppb = (ae_dfd - a_e_exp) / a_e_exp * 1e9
exp_unc_ppb = a_e_exp_unc / a_e_exp * 1e9   # ~0.11 ppb
# DFD prediction error relative to QED (its alpha-limited budget)
dfd_vs_qed_ppb = (ae_dfd - ae_qed) / ae_qed * 1e9

names = ["QED\n(alpha = CODATA,\nfitted input)", "DFD\n(alpha derived,\nparameter-free)"]
vals  = [res_qed_ppb, res_dfd_ppb]
cols  = [C_QED, C_DFD]
xpos  = np.arange(len(names))

bars = axB.bar(xpos, vals, width=0.55, color=cols, edgecolor="black", lw=0.8, zorder=3)

# experimental uncertainty band around zero (the target = experiment)
axB.axhspan(-exp_unc_ppb, exp_unc_ppb, color=C_EXP, alpha=0.22, zorder=0,
            label=f"experiment $\\pm1\\sigma$ ({exp_unc_ppb:.2f} ppb)")
axB.axhline(0, color=C_EXP, lw=1.2, ls="--", zorder=1)

# value labels on bars
for b, v in zip(bars, vals):
    off = 0.25 if v >= 0 else -0.25
    va = "bottom" if v >= 0 else "top"
    axB.text(b.get_x() + b.get_width() / 2, v + off,
             f"{v:+.2f} ppb", ha="center", va=va, fontsize=9.5, fontweight="bold")

axB.set_xticks(xpos)
axB.set_xticklabels(names, fontsize=8.8)
axB.set_ylabel(r"residual vs experiment  $(a_e-a_e^{\rm exp})/a_e^{\rm exp}$  [ppb]")
axB.set_title("(b) DFD predicts $a_e$ parameter-free to ~6 ppb\n"
              "(QED uses $a_e$ itself to fix $\\alpha$)")
axB.set_ylim(min(vals) - 2.2, max(abs(min(vals)), abs(max(vals)), exp_unc_ppb) + 2.4)
axB.legend(loc="upper left", fontsize=8.5, framealpha=0.95)

# annotate the parameter-free message + the 3 raw values
txt = (f"$\\alpha^{{-1}}$ (CODATA) = {ainv_codata}\n"
       f"$\\alpha^{{-1}}$ (DFD)    = {ainv_dfd}\n"
       f"$\\Delta\\alpha$ = {1e9*(ainv_dfd-ainv_codata)/ainv_codata:+.1f} ppb\n"
       f"$\\Rightarrow$ DFD $a_e$ vs QED = {dfd_vs_qed_ppb:+.2f} ppb")
axB.text(0.035, 0.035, txt, transform=axB.transAxes,
         ha="left", va="bottom", fontsize=8.0,
         bbox=dict(boxstyle="round,pad=0.45", fc="#fbf7e8", ec="#b8a25a", lw=0.8))

fig.suptitle("Electron anomalous moment $a_e=(g-2)/2$: QED's 12-digit machinery, "
             "made parameter-free by DFD's derived $\\alpha$",
             fontsize=12, y=0.995)

# small footer with the three exact numbers
foot = (f"$a_e$:  exp = {a_e_exp:.14f}   |   "
        f"QED 5-loop = {ae_qed:.14f}   |   "
        f"DFD = {ae_dfd:.14f}")
fig.text(0.5, 0.005, foot, ha="center", va="bottom", fontsize=8.0, family="monospace")

fig.tight_layout(rect=[0, 0.035, 1, 0.96])
out = "./figures/g2.png"
fig.savefig(out, dpi=150, bbox_inches="tight")
print(f"saved: {out}")
