"""Does the single-particle spectral stretch scale as 1/(m*/m)?

Level density near the Fermi surface g(e) ~ m*, so within-shell level spacing
~ 1/m*. If the empirical (surface) effective mass is ~1 (particle-vibration
coupling), a force with m*/m < 1 should over-stretch the spectrum by ~ m/m*.
Test: compare calc vs experimental WITHIN-shell energy spans in 208Pb across
five forces spanning m*/m = 0.56 -> 0.91.
"""
import numpy as np
import matplotlib.pyplot as plt
import skyrme_hf as m
from skyrme_hf import Orbital as O

plt.rcParams.update({"figure.dpi": 140, "savefig.dpi": 140, "font.size": 10,
                     "axes.spines.top": False, "axes.spines.right": False,
                     "font.family": "DejaVu Sans"})

def parse(label):
    spec = "spdfghijklm"
    return O(int(label[0]), spec.index(label[1]), int(label.split("/")[0][2:]))

# 208Pb experimental single-particle energies (SWV, arXiv:0709.3525), split into
# within-shell sub-ladders (same side of the magic gap):
LADDERS = {
    "n-hole": {"1h9/2": -11.40, "2f7/2": -9.81, "1i13/2": -9.24,
               "3p3/2": -8.26, "2f5/2": -7.94, "3p1/2": -7.37},
    "n-part": {"2g9/2": -3.94, "1i11/2": -3.16, "1j15/2": -2.51, "3d5/2": -2.37,
               "4s1/2": -1.90, "2g7/2": -1.44, "3d3/2": -1.40},
    "p-hole": {"1g7/2": -11.49, "2d5/2": -9.70, "1h11/2": -9.36,
               "2d3/2": -8.36, "3s1/2": -8.01},
    "p-part": {"1h9/2": -3.80, "2f7/2": -2.90, "1i13/2": -2.10,
               "2f5/2": -0.97, "3p3/2": -0.68, "3p1/2": -0.16},
}
ISO = {"n-hole": "n", "n-part": "n", "p-hole": "p", "p-part": "p"}

FORCES = [m.SII, m.SLy4, m.SIII, m.SkMs, m.SI]

def analyze(force):
    res = m.run_hf(force, m.closed_shell_208Pb(), grid=m.Grid(rmax=22, h=0.1),
                   mix=0.22, maxiter=400, com_correction=True, verbose=False)
    ratios, rms_all = [], []
    for lad, exp in LADDERS.items():
        wanted = [parse(lb) for lb in exp]
        lp, ln = m.sp_levels(res, wanted, wanted)
        levs = lp if ISO[lad] == "p" else ln
        calc = {o.label(): e for o, e in levs}
        exp_span = max(exp.values()) - min(exp.values())
        calc_span = max(calc[k] for k in exp) - min(calc[k] for k in exp)
        ratios.append(calc_span / exp_span)
        rms_all += [calc[k] - exp[k] for k in exp]
    # de-mean each ladder handled by span; report mean span ratio + overall rms
    return np.mean(ratios), np.sqrt(np.mean(np.array(rms_all)**2))

rows = []
for f in FORCES:
    stretch, rms = analyze(f)
    rows.append((f.name, m.mstar_over_m(f), stretch, rms))
    print(f"{f.name:5s}  m*/m={rows[-1][1]:.3f}  span_ratio={stretch:.3f}  rms={rms:.2f}")

names = [r[0] for r in rows]
mstar = np.array([r[1] for r in rows])
stretch = np.array([r[2] for r in rows])
inv = 1.0 / mstar

# ---- plot: spectral stretch vs 1/(m*/m) --------------------------------------
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.6))

ax1.plot(inv, stretch, "o", ms=9, color="#c0392b", zorder=5)
for n, x, y in zip(names, inv, stretch):
    ax1.annotate(n, (x, y), textcoords="offset points", xytext=(8, -3), fontsize=9)
# reference line: stretch = m/m* (i.e. y = x) would mean empirical m*_eff = 1
xx = np.linspace(inv.min()*0.98, inv.max()*1.02, 50)
ax1.plot(xx, xx, "--", color="#888", lw=1, label=r"$y=x$  (empirical $m^*_{\rm surf}=m$)")
# best-fit line through the data
a, b = np.polyfit(inv, stretch, 1)
ax1.plot(xx, a*xx + b, "-", color="#2c7fb8", lw=1.3,
         label=f"fit: {a:.2f}·(m/m*) {b:+.2f}")
r = np.corrcoef(inv, stretch)[0, 1]
ax1.set_xlabel(r"$m/m^*$  (inverse effective mass)")
ax1.set_ylabel(r"within-shell spectral stretch  $\Delta_{\rm calc}/\Delta_{\rm exp}$")
ax1.set_title(f"208Pb: level spacing tracks $1/m^*$   (r = {r:.3f})")
ax1.legend(fontsize=8.5, loc="upper left")

ax2.plot(mstar, [r[3] for r in rows], "s", ms=9, color="#31a354", zorder=5)
for n, x, y in zip(names, mstar, [r[3] for r in rows]):
    ax2.annotate(n, (x, y), textcoords="offset points", xytext=(8, -3), fontsize=9)
ax2.set_xlabel(r"$m^*/m$")
ax2.set_ylabel("208Pb s.p. energy RMS deviation (MeV)")
ax2.set_title("Overall spectrum error vs effective mass")

fig.suptitle("The single-particle spectrum error is an effective-mass effect",
             y=1.02, fontsize=12, fontweight="bold")
fig.tight_layout()
fig.savefig("./fig7_mstar_stretch.png", bbox_inches="tight")
plt.close(fig)
print(f"\ncorrelation r(stretch, 1/m*) = {r:.3f};  fit slope {a:.2f}")
print("wrote fig7_mstar_stretch.png")