"""Controlled isolation of the effective-mass effect on the s.p. spectrum.

Converge SIII for 208Pb, then FREEZE the central potential U, spin-orbit field W,
Coulomb, and density.  Scale only the effective-mass term,
    B_lambda(r) = hbar^2/2m + lambda * (B_SIII(r) - hbar^2/2m),
and re-diagonalize the identical potential.  lambda tunes m*/m with everything
else fixed.

Probe with PURE RADIAL ladders (same l,j, successive n): these isolate the
radial level density with no spin-orbit / l-splitting confound.  Extract the
power p in  spacing ~ (m*/m)^(-p)  by a log-log fit.
"""
import numpy as np
import matplotlib.pyplot as plt
import skyrme_hf as m
from skyrme_hf import Orbital as O, HBAR2_2M

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

# pure-radial neutron ladders in 208Pb (all bound): request n=1..3 of each (l,j)
RAD = {"s1/2": (0, 1), "p3/2": (1, 3), "p1/2": (1, 1), "d5/2": (2, 5)}
NMAX = 3

res = m.run_hf(m.SIII, m.closed_shell_208Pb(), grid=m.Grid(rmax=22, h=0.1),
               mix=0.22, maxiter=400, com_correction=True, verbose=False)
g = res["grid"]; r = g.r; fld = res["fields"]
B0 = fld["B_n"].copy(); U0, W0f = fld["U_n"], fld["W_n"]
interior = r < 4.0
# orbitals to request: n=1..NMAX for each (l,j)
wanted = [O(n, l, j2) for (l, j2) in RAD.values() for n in range(1, NMAX+1)]

def radial_spacings(lam):
    Blam = HBAR2_2M + lam * (B0 - HBAR2_2M)
    _, lev = m.solve_isospin(g, wanted, Blam, U0, W0f, np.zeros_like(r))
    e = {o.label(): en for o, en in lev}
    mstar = HBAR2_2M / np.mean(Blam[interior])
    sp = {}
    for name, (l, j2) in RAD.items():
        spec = "spdfghijklm"
        labs = [f"{n}{spec[l]}{j2}/2" for n in range(1, NMAX+1)]
        gaps = [e[labs[i+1]] - e[labs[i]] for i in range(NMAX-1)]
        sp[name] = np.mean(gaps)          # mean radial spacing for this (l,j)
    return mstar, sp

lams = np.linspace(0.4, 1.6, 13)
ms = []; series = {k: [] for k in RAD}
for lam in lams:
    mstar, sp = radial_spacings(lam)
    ms.append(mstar)
    for k in RAD:
        series[k].append(sp[k])
ms = np.array(ms)
series = {k: np.array(v) for k, v in series.items()}

# log-log fit: log(spacing) = -p*log(m*/m) + c  ->  slope = -p
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.6))
colors = {"s1/2": "#c0392b", "p3/2": "#2c7fb8", "p1/2": "#31a354", "d5/2": "#8e44ad"}
powers = {}
for k in RAD:
    p_fit = -np.polyfit(np.log(ms), np.log(series[k]), 1)[0]
    powers[k] = p_fit
    ax1.plot(ms, series[k], "o-", color=colors[k], ms=5,
             label=fr"{k}:  spacing $\propto (m^*/m)^{{-{p_fit:.2f}}}$")
ax1.axvline(HBAR2_2M/np.mean(B0[interior]), color="#888", lw=0.8, ls=":")
ax1.text(HBAR2_2M/np.mean(B0[interior]), ax1.get_ylim()[0], " SIII", fontsize=8,
         color="#555", va="bottom")
ax1.set_xlabel(r"$m^*/m$  (interior, frozen 208Pb potential)")
ax1.set_ylabel("mean radial level spacing (MeV)")
ax1.set_title("Pure-radial spacing vs effective mass")
ax1.legend(fontsize=8, loc="upper right")

# log-log panel to show power law cleanly
for k in RAD:
    ax2.loglog(ms, series[k], "o-", color=colors[k], ms=5)
mm = np.array([ms.min(), ms.max()])
# reference slopes -1 and -1/2 anchored at SIII s1/2 point
i0 = np.argmin(abs(ms - HBAR2_2M/np.mean(B0[interior])))
anch = series["s1/2"][i0]; m0 = ms[i0]
ax2.loglog(mm, anch*(mm/m0)**(-1.0), "--", color="#555", lw=0.9, label=r"slope $-1$ (box-like)")
ax2.loglog(mm, anch*(mm/m0)**(-0.5), ":", color="#555", lw=0.9, label=r"slope $-1/2$ (h.o.-like)")
ax2.set_xlabel(r"$m^*/m$"); ax2.set_ylabel("radial spacing (MeV)")
ax2.set_title(f"Power law:  mean exponent p = {np.mean(list(powers.values())):.2f}")
ax2.legend(fontsize=8, loc="lower left")

fig.suptitle(r"Controlled test: with potential frozen, radial spacing $\propto (m^*/m)^{-p}$,  $p\approx$"
             f"{np.mean(list(powers.values())):.2f}",
             y=1.03, fontsize=12, fontweight="bold")
fig.tight_layout()
fig.savefig("./fig8_mstar_controlled.png", bbox_inches="tight")
plt.close(fig)

print("m*/m and radial spacings:")
for i, lam in enumerate(lams):
    print(f"  m*/m={ms[i]:.3f}  " + "  ".join(f"{k}={series[k][i]:.2f}" for k in RAD))
print("\nfitted exponents p (spacing ~ (m*/m)^-p):")
for k in RAD:
    print(f"  {k}: p={powers[k]:.3f}")
print(f"  mean p = {np.mean(list(powers.values())):.3f}")
print("wrote fig8_mstar_controlled.png")