"""Test the solver across the full Vautherin-Brink doubly-magic chain."""
import numpy as np
import skyrme_hf as m

# experimental targets: (BE_total MeV, BE/A, rms_charge fm)
EXPT = {
    "16O":   (127.62, 7.976, 2.699),
    "40Ca":  (342.05, 8.551, 3.478),
    "48Ca":  (415.99, 8.666, 3.477),
    "90Zr":  (783.89, 8.710, 4.269),
    "208Pb": (1636.43, 7.867, 5.501),
}

NUCLEI = [m.closed_shell_16O(), m.closed_shell_40Ca(), m.closed_shell_48Ca(),
          m.closed_shell_90Zr(), m.closed_shell_208Pb()]

def run_all(force, h=0.1, rmax=22.0, mix=0.25):
    print(f"\n===== {force.name}  (m*/m = {m.mstar_over_m(force):.3f}) =====")
    print(f"{'nucleus':8s} {'E_HF':>9s} {'BE/A':>7s} {'rch':>6s}  |"
          f" {'E_exp':>9s} {'BE/A':>6s} {'rch':>6s}  | {'dBE/A':>6s} {'drch':>6s}")
    for nuc in NUCLEI:
        grid = m.Grid(rmax=rmax, h=h)
        res = m.run_hf(force, nuc, grid=grid, mix=mix, maxiter=300, verbose=False)
        be, bea, rch = EXPT[nuc.name]
        dbea = res['BE_per_A'] - bea
        drch = res['rms_ch'] - rch
        print(f"{nuc.name:8s} {res['E']:9.2f} {res['BE_per_A']:7.3f} "
              f"{res['rms_ch']:6.3f}  | {-be:9.2f} {bea:6.3f} {rch:6.3f}  |"
              f" {dbea:+6.3f} {drch:+6.3f}  [{res['iters']} it]")

def pb_spin_orbit(force, h=0.1, rmax=22.0, mix=0.25):
    """208Pb single-particle spectrum near the Fermi surface + ls splittings."""
    nuc = m.closed_shell_208Pb()
    grid = m.Grid(rmax=rmax, h=h)
    res = m.run_hf(force, nuc, grid=grid, mix=mix, maxiter=300, verbose=False)
    print(f"\n208Pb proton levels ({force.name}):")
    for o, e in sorted(res['levels_p'], key=lambda t: t[1]):
        print(f"   {o.label():8s} {e:8.3f}")
    print(f"208Pb neutron levels ({force.name}):")
    for o, e in sorted(res['levels_n'], key=lambda t: t[1]):
        print(f"   {o.label():8s} {e:8.3f}")

if __name__ == "__main__":
    for force in (m.SIII, m.SI, m.SII):
        run_all(force)
    pb_spin_orbit(m.SIII)