"""Self-consistent Kohn-Sham (LDA) screening of a point charge Z in jellium.

A nucleus of charge Z sits at the centre of a uniform electron gas of density
n0 = 3 / (4 pi r_s^3) with a rigid neutralising background. The Kohn-Sham equations are
solved in spherical symmetry, spin-unpolarised, with Dirac exchange and the Perdew-Zunger
(1981) parametrisation of the Ceperley-Alder correlation energy. The output is the set of
phase shifts delta_l at the Fermi momentum and the scattering strength

    h = (4 hbar^2 k_F / m) * sum_l (l + 1) sin^2(delta_l - delta_{l+1}).

Atomic units (hartree, bohr) inside; h = 4 k_F * sum in units of e^2 = 14.3996 eV angstrom.

Method. Radial functions u(r) are integrated outward with Numerov's method on the mesh
r = x^2 (uniform in x, which removes the Coulomb singularity) from a power series at the
origin. Scattering states for k on a Gauss-Legendre grid in [0, k_F] are normalised to
sin(kr - l pi/2 + delta_l) by matching to Riccati-Bessel functions at two radii outside
the potential. Bound states are found by bisection on the node count. The induced density is

    dn(r) = sum_bound 2(2l+1) u_nl^2 / (4 pi r^2)
          + (1 / pi^2 r^2) sum_k w_k sum_l (2l+1) [u_l(k,r)^2 - u_l^free(k,r)^2],

with the free reference integrated by the same code. The effective potential is
V = -Z/r + V_H[dn] + v_xc(n0 + dn) - v_xc(n0) inside the radius R and zero outside. The
Poisson equation is solved with V_H(R) set to the value that the Friedel oscillation of
the density implies there, (pi / k_F^2) dn(R) to leading order; setting it to zero instead
is an option and is one of the checks. Neutrality is not imposed: the charge inside R and
the Friedel sum are computed from the converged solution and compared with Z.
Self-consistency is reached by Anderson mixing of the potential residual after it has
been filtered by q^2 / (q^2 + k_TF^2) (the Thomas-Fermi screened Poisson equation, solved
with its radial Green function).

Run:  python jellium.py           all tables and checks, writes jellium.json (3 to 4 minutes)
      python jellium.py --quick   the Z = 1 and Z = 5 tables only, nothing written (10 seconds)
"""

import argparse
import json
import math
import time
from pathlib import Path

import numpy as np
from scipy.integrate import cumulative_simpson, quad
from scipy.optimize import brentq
from scipy.special import gammaln, kve, spherical_jn, spherical_yn

HERE = Path(__file__).parent
E2 = 14.399645  # e^2 in eV angstrom = 1 hartree bohr
BOHR = 0.529177211
HARTREE_EV = 27.211386
PI = math.pi

RS_PROTON = (1.0, 1.3, 1.5, 1.8, 2.0, 2.07, 2.5, 3.0, 4.0, 5.0, 6.0)
RS_BORON = (1.5, 1.8, 2.07)

PZ_JOIN = 0.01  # half-width in r_s of the join between the two branches of the PZ81 fit

# Default numerical parameters. R is the radius beyond which the potential is set to zero.
DEFAULTS = {"R_over_rs": 16.0, "R_min": 40.0, "dx": 0.004, "nk_min": 64, "lmax": 12, "tol": 1e-9}


def default_nk(kf, R):
    """Number of Gauss-Legendre points in k. The integrand contains cos(2 k r), which in the
    variable t = 2 k / k_F - 1 is cos(k_F r (t + 1)); an n-point rule integrates it to
    rounding error once 2n exceeds k_F R by several times (k_F R)^(1/3), so n has to grow
    with the box. Below that the quadrature aliases and the charge at large r is wrong."""
    a = kf * R
    return max(DEFAULTS["nk_min"], int(math.ceil(0.5 * a + 5.0 * a ** (1.0 / 3.0))) + 8)


# ----------------------------------------------------------------------------- LDA

def v_xc(n, kind="pz81"):
    """Exchange-correlation potential in hartree for electron density n (bohr^-3).

    pz81: Dirac exchange + Perdew-Zunger 1981 fit to Ceperley-Alder (unpolarised).
    pw92: Dirac exchange + Perdew-Wang 1992 fit to the same data (for comparison).
    x: exchange only.  hartree: zero.
    """
    n = np.maximum(np.asarray(n, dtype=float), 1e-30)
    if kind == "hartree":
        return np.zeros_like(n)
    rs = (3.0 / (4.0 * PI * n)) ** (1.0 / 3.0)
    vx = -((3.0 * n / PI) ** (1.0 / 3.0))
    if kind == "x":
        return vx
    if kind == "hl":  # Hedin and Lundqvist 1971
        return vx * (1.0 + 0.7734 * (rs / 21.0) * np.log1p(21.0 / rs))
    if kind == "gl":  # Gunnarsson and Lundqvist 1976
        return vx * (1.0 + 0.0545 * rs * np.log1p(11.4 / rs))
    if kind == "pz81":
        gamma, b1, b2 = -0.1423, 1.0529, 0.3334
        a, b, c, d = 0.0311, -0.048, 0.0020, -0.0116
        sq = np.sqrt(rs)
        den = 1.0 + b1 * sq + b2 * rs
        low = gamma * (1.0 + (7.0 / 6.0) * b1 * sq + (4.0 / 3.0) * b2 * rs) / den**2
        ln = np.log(rs)
        high = a * ln + (b - a / 3.0) + (2.0 / 3.0) * c * rs * ln + (2.0 * d - c) / 3.0 * rs
        # The published fit has a jump of 2.8e-5 hartree in v_c at r_s = 1. A background at
        # exactly r_s = 1 crosses it at every half period of the Friedel oscillation and the
        # iteration cannot converge below that size, so the two branches are joined linearly
        # across 1 - PZ_JOIN < r_s < 1 + PZ_JOIN. Outside that window this is the published fit.
        t = np.clip((rs - (1.0 - PZ_JOIN)) / (2.0 * PZ_JOIN), 0.0, 1.0)
        return vx + (1.0 - t) * high + t * low
    if kind == "pw92":
        a, a1, b1, b2, b3, b4 = 0.0310907, 0.21370, 7.5957, 3.5876, 1.6382, 0.49294
        sq = np.sqrt(rs)
        q0 = -2.0 * a * (1.0 + a1 * rs)
        q1 = 2.0 * a * (b1 * sq + b2 * rs + b3 * rs * sq + b4 * rs * rs)
        q1p = a * (b1 / sq + 2.0 * b2 + 3.0 * b3 * sq + 4.0 * b4 * rs)
        log = np.log1p(1.0 / q1)
        ec = q0 * log
        dec = -2.0 * a * a1 * log - q0 * q1p / (q1 * q1 + q1)
        return vx + ec - rs / 3.0 * dec
    raise ValueError(kind)


# ----------------------------------------------------------------------------- mesh

class Mesh:
    """r = x^2 on a uniform x grid. The potential lives on [0, Rc]; the grid runs a little
    further so that scattering states can be matched to free solutions at two radii."""

    def __init__(self, R, dx, pad=1.5):
        self.dx = dx
        self.nc = int(round(math.sqrt(R) / dx))
        self.Rc = (self.nc * dx) ** 2
        self.N = int(math.ceil(math.sqrt(self.Rc + pad) / dx))
        self.x = dx * np.arange(self.N + 1)
        self.r = self.x**2
        self.na = self.nc + 2
        self.nb = self.N

    def cumint(self, f_of_r, upto=None):
        """Running integral over r from 0 of f(r), for f given on the grid (or its first part)."""
        m = len(f_of_r) if upto is None else upto
        return cumulative_simpson(f_of_r[:m] * 2.0 * self.x[:m], dx=self.dx, initial=0.0)


def riccati(l, z):
    """Riccati-Bessel functions S = z j_l(z), C = -z y_l(z); u = cos(d) S + sin(d) C."""
    return z * spherical_jn(l, z), -z * spherical_yn(l, z)


# ----------------------------------------------------------------------------- radial solver

class Waves:
    """Regular solutions of u'' = [l(l+1)/r^2 + 2 V(r) - e2] u for a list of (l, e2).

    e2 = k^2 for scattering states and 2E < 0 for bound states. With r = x^2 and
    u = sqrt(2x) y the equation is y'' = g y,
        g = [4 l(l+1) + 3/4] / x^2 + 8 x^2 V - 4 x^2 e2,
    and 8 x^2 V = -8 Z + 8 x^2 V_s with V_s = V + Z/r finite at the origin.
    """

    def __init__(self, mesh, l, e2):
        self.mesh = mesh
        self.l = np.asarray(l, dtype=int)
        self.e2 = np.asarray(e2, dtype=float)
        x = mesh.x[1:, None]
        base = np.zeros((mesh.N + 1, self.l.size))
        base[1:] = (4.0 * self.l * (self.l + 1) + 0.75) / x**2 - 4.0 * x**2 * self.e2
        self.tbase = 1.0 - (mesh.dx**2 / 12.0) * base
        # Series is used up to index n1(l); beyond it Numerov is stable for that l.
        self.n1 = np.minimum(6 + 3 * self.l, mesh.nc // 2)

    def solve(self, v2, z, v0, v1, logscale):
        """Return (W, T) with y = W / T on the grid. v2 = 8 x^2 V on the grid."""
        mesh = self.mesh
        t = self.tbase - (mesh.dx**2 / 12.0) * v2[:, None]
        n1, n1max, n1min = self.n1, int(self.n1.max()), int(self.n1.min())
        r = mesh.r[1 : n1max + 1, None]
        lf = self.l.astype(float)
        e = 0.5 * self.e2
        a1 = -z / (lf + 1.0)
        a2 = (-2.0 * z * a1 + 2.0 * (v0 - e)) / (2.0 * (2.0 * lf + 3.0))
        a3 = (-2.0 * z * a2 + 2.0 * (v0 - e) * a1 + 2.0 * v1) / (3.0 * (2.0 * lf + 4.0))
        with np.errstate(under="ignore"):
            u = np.exp((lf + 1.0) * np.log(r) + logscale) * (1.0 + r * (a1 + r * (a2 + r * a3)))
        wser = t[1 : n1max + 1] * u / np.sqrt(2.0 * mesh.x[1 : n1max + 1, None])
        c = 12.0 / t - 10.0
        w = np.empty_like(t)
        w[0] = 0.0
        w[1 : n1max + 1] = wser
        with np.errstate(under="ignore"):
            for n in range(n1min, n1max):
                rec = c[n] * w[n] - w[n - 1]
                w[n + 1] = np.where(n + 1 <= n1, wser[n], rec)
            for n in range(n1max, mesh.N):
                np.multiply(c[n], w[n], out=w[n + 1])
                w[n + 1] -= w[n - 1]
        return w, t

    def nodes(self, w):
        """Number of sign changes of each solution between the series region and r_b."""
        s = np.signbit(w[int(self.n1.min()) : self.mesh.nb + 1])
        return np.count_nonzero(s[1:] != s[:-1], axis=0)

    def u_at(self, w, t, n):
        return math.sqrt(2.0 * self.mesh.x[n]) * w[n] / t[n]


def free_logscale(l, k):
    """log of k^(l+1) / (2l+1)!!, so that the series start equals the free S_l(kr)."""
    l = np.asarray(l, dtype=float)
    return (l + 1.0) * np.log(k) - (gammaln(2.0 * l + 2.0) - l * math.log(2.0) - gammaln(l + 1.0))


class Continuum:
    """Scattering states on a Gauss-Legendre grid in k and the density they induce."""

    def __init__(self, mesh, kf, nk, lmax, thresh=1e-7):
        self.mesh = mesh
        t, w = np.polynomial.legendre.leggauss(nk)
        k = 0.5 * kf * (t + 1.0)
        wk = 0.5 * kf * w
        ll, kk = np.meshgrid(np.arange(lmax + 1), np.arange(nk), indexing="ij")
        ll, kk = ll.ravel(), kk.ravel()
        # Drop (l, k) pairs whose free solution is negligible everywhere inside the grid.
        with np.errstate(under="ignore"):
            s_edge = k[kk] * mesh.r[-1] * spherical_jn(ll, k[kk] * mesh.r[-1])
        keep = np.abs(s_edge) > thresh
        self.l, self.k = ll[keep], k[kk[keep]]
        self.wt = wk[kk[keep]] * (2 * self.l + 1)
        self.waves = Waves(mesh, self.l, self.k**2)
        self.logscale = free_logscale(self.l, self.k)
        self.sa, self.ca = riccati(self.l, self.k * mesh.r[mesh.na])
        self.sb, self.cb = riccati(self.l, self.k * mesh.r[mesh.nb])
        self.det = self.sa * self.cb - self.sb * self.ca
        w0, t0 = self.waves.solve(np.zeros(mesh.N + 1), 0.0, 0.0, 0.0, self.logscale)
        self.free = self._raw(w0, t0)

    def amplitudes(self, w, t):
        ua, ub = self.waves.u_at(w, t, self.mesh.na), self.waves.u_at(w, t, self.mesh.nb)
        alpha = (ua * self.cb - ub * self.ca) / self.det
        beta = (self.sa * ub - self.sb * ua) / self.det
        return alpha, beta

    def _raw(self, w, t):
        alpha, beta = self.amplitudes(w, t)
        y = w / t
        with np.errstate(under="ignore"):
            np.square(y, out=y)
            d = y @ (self.wt / (alpha**2 + beta**2))
        return d * 2.0 * self.mesh.x

    def induced(self, v2, z, v0, v1):
        """Induced density of the continuum, (1/pi^2 r^2) sum w (2l+1) [u^2 - u_free^2]."""
        w, t = self.waves.solve(v2, z, v0, v1, self.logscale)
        dn = np.zeros(self.mesh.N + 1)
        dn[1:] = (self._raw(w, t)[1:] - self.free[1:]) / (PI**2 * self.mesh.r[1:] ** 2)
        return dn


# ----------------------------------------------------------------------------- bound states

def _decay_ratio(l, kappa, ra, rb):
    """d(rb) / d(ra) for the decaying free solution d(r) = kappa r k_l(kappa r)."""
    za, zb = kappa * ra, kappa * rb
    return np.sqrt(zb / za) * kve(l + 0.5, zb) / kve(l + 0.5, za) * np.exp(-(zb - za))


def _count_below(mesh, l, energies, v2, z, v0, v1):
    """Number of bound states of angular momentum l below each (negative) energy.

    Equal to the number of nodes of the outward solution inside r_b plus one if its
    continuation by free solutions beyond r_b has a further node.
    """
    energies = np.atleast_1d(np.asarray(energies, dtype=float))
    wv = Waves(mesh, np.full(energies.size, l), 2.0 * energies)
    scale = -(l + 1.0) * math.log(mesh.r[int(wv.n1[0])])
    w, t = wv.solve(v2, z, v0, v1, scale)
    ua, ub = wv.u_at(w, t, mesh.na), wv.u_at(w, t, mesh.nb)
    kappa = np.sqrt(-2.0 * energies)
    grow = ub - ua * _decay_ratio(l, kappa, mesh.r[mesh.na], mesh.r[mesh.nb])
    extra = np.signbit(grow) != np.signbit(ub)
    return wv.nodes(w) + extra, wv, w, t


def bound_states(mesh, v2, veff, z, v0, v1, lmax_b, e_floor=1e-9):
    """All bound states with l <= lmax_b and E < -e_floor. Returns (list, density)."""
    r, x = mesh.r, mesh.x
    dens = np.zeros(mesh.N + 1)
    found = []
    e_low = -(0.5 * z * z + 1.0)
    for l in range(lmax_b + 1):
        scan = -np.geomspace(-e_low, e_floor, 72)
        counts = _count_below(mesh, l, scan, v2, z, v0, v1)[0]
        if counts[0] != 0:
            raise RuntimeError("bound-state scan does not start below the spectrum")
        for j in range(int(counts[-1])):
            hi_i = int(np.argmax(counts > j))
            lo, hi = scan[hi_i - 1], scan[hi_i]
            for _ in range(60):
                trial = np.linspace(lo, hi, 26)[1:-1]
                above = _count_below(mesh, l, trial, v2, z, v0, v1)[0] > j
                if above[0]:
                    hi = trial[0]
                elif not above[-1]:
                    lo = trial[-1]
                else:
                    i = int(np.argmax(above))
                    lo, hi = trial[i - 1], trial[i]
                if hi - lo <= 4e-16 * abs(lo):
                    break
            e = float(0.5 * (lo + hi))
            _, wv, w, t = _count_below(mesh, l, [e], v2, z, v0, v1)
            u = np.sqrt(2.0 * x) * w[:, 0] / t[:, 0]
            # Outward integration picks up the growing solution far inside the barrier:
            # cut the function at the first minimum of |u| beyond the outer turning point.
            f = np.zeros(mesh.N + 1)
            f[1:] = l * (l + 1) / r[1:] ** 2 + 2.0 * veff[1:] - 2.0 * e
            inside = np.nonzero(f[1:] < 0.0)[0]
            nt = int(inside[-1]) + 1 if inside.size else 1
            au = np.abs(u)
            rising = np.nonzero(au[nt + 1 :] > au[nt:-1])[0]
            cut = nt + int(rising[0]) if rising.size else mesh.N
            u[cut + 1 :] = 0.0
            norm = mesh.cumint(u * u)[-1]
            if cut == mesh.N:
                kappa = math.sqrt(-2.0 * e)
                rn = r[-1]
                tail = quad(lambda s: float(_decay_ratio(l, kappa, rn, rn + s)) ** 2, 0.0, np.inf)[0]
                norm += u[-1] ** 2 * tail
            inside_box = mesh.cumint(u * u, mesh.nc + 1)[-1] / norm
            occ = 2 * (2 * l + 1)
            dens[1:] += occ * u[1:] ** 2 / (4.0 * PI * r[1:] ** 2 * norm)
            found.append({"l": l, "n": j + l + 1, "E_Ha": e, "E_eV": e * HARTREE_EV,
                          "electrons": occ, "fraction_inside_R": float(inside_box)})
    return found, dens


# ----------------------------------------------------------------------------- potential

def hartree_smooth(mesh, dn, z, v_edge=0.0):
    """V_H + Z/r on [0, Rc] for induced density dn, with V_H(Rc) = v_edge. Returns (V, Q(r))."""
    nc = mesh.nc
    r = mesh.r[: nc + 1]
    q = mesh.cumint(4.0 * PI * r**2 * dn[: nc + 1])
    p = mesh.cumint(4.0 * PI * r * dn[: nc + 1])
    v = np.empty(nc + 1)
    v[1:] = q[1:] / r[1:] + (p[-1] - p[1:])
    v[0] = p[-1]
    return v + (z - q[-1]) / mesh.Rc + v_edge, q


def screened_filter(mesh, res, kappa):
    """Apply q^2 / (q^2 + kappa^2) to a radial function on [0, Rc] that vanishes at Rc.

    res - kappa^2 y with (-lap + kappa^2) y = res, solved with the Green function
    sinh(kappa r<) sinh(kappa (Rc - r>)) / (kappa sinh(kappa Rc)) for r y.
    """
    nc = mesh.nc
    r = mesh.r[: nc + 1]
    a = np.sinh(kappa * r)
    b = np.exp(-kappa * r) * (-np.expm1(-2.0 * kappa * (mesh.Rc - r))) / (-np.expm1(-2.0 * kappa * mesh.Rc))
    i1 = mesh.cumint(a * r * res)
    i2 = mesh.cumint(b * r * res)
    i2 = i2[-1] - i2
    out = np.empty(nc + 1)
    out[1:] = res[1:] - kappa * (b[1:] * i1[1:] + a[1:] * i2[1:]) / r[1:]
    out[0] = res[0] - kappa**2 * i2[0]
    return out


class Anderson:
    def __init__(self, beta, depth, weight):
        self.beta, self.depth, self.weight = beta, depth, weight
        self.xs, self.fs = [], []

    def step(self, x, f):
        self.xs.append(x.copy())
        self.fs.append(f.copy())
        self.xs, self.fs = self.xs[-(self.depth + 1) :], self.fs[-(self.depth + 1) :]
        m = len(self.xs) - 1
        if m == 0:
            return x + self.beta * f
        df = np.array([f - fi for fi in self.fs[:-1]])
        dx = np.array([x - xi for xi in self.xs[:-1]])
        a = (df * self.weight) @ df.T
        rhs = (df * self.weight) @ f
        gam = np.linalg.lstsq(a + 1e-12 * np.trace(a) / m * np.eye(m), rhs, rcond=None)[0]
        return (x - gam @ dx) + self.beta * (f - gam @ df)


# ----------------------------------------------------------------------------- phase shifts

def phase_shifts(mesh, kf, v2, z, v0, v1, lmax):
    """Absolute phase shifts at k_F (zero at infinite energy, n_bound * pi at zero energy).

    The value modulo pi comes from the two-point match. The multiple of pi comes from
    counting the nodes of the solution and of the free solution inside r_b.
    """
    l = np.arange(lmax + 1)
    wv = Waves(mesh, l, np.full(l.size, kf * kf))
    scale = free_logscale(l, kf)
    w, t = wv.solve(v2, z, v0, v1, scale)
    wf, tf = wv.solve(np.zeros(mesh.N + 1), 0.0, 0.0, 0.0, scale)
    sa, ca = riccati(l, kf * mesh.r[mesh.na])
    sb, cb = riccati(l, kf * mesh.r[mesh.nb])

    def match(w, t):
        ua, ub = wv.u_at(w, t, mesh.na), wv.u_at(w, t, mesh.nb)
        num, den = sa * ub - sb * ua, ua * cb - ub * ca
        # Solutions that have underflowed to zero (l far above k r) have no phase shift.
        return np.arctan(np.divide(num, den, out=np.zeros_like(num), where=den != 0.0))

    # The free solution integrated by the same code has a phase error of order dx^4
    # against the analytic Riccati-Bessel functions; subtracting it removes that error.
    principal = match(w, t) - match(wf, tf)
    frac = np.mod(np.arctan2(sb, cb), PI)
    return PI * (wv.nodes(w) - wv.nodes(wf)) + np.mod(frac + principal, PI) - frac


def strength_sum(delta):
    l = np.arange(delta.size - 1)
    return float(np.sum((l + 1) * np.sin(delta[:-1] - delta[1:]) ** 2))


# ----------------------------------------------------------------------------- SCF

def solve(z, rs, R=None, dx=None, nk=None, lmax=None, tol=None, xc="pz81", maxit=400, verbose=False,
          bc="friedel", fields=None):
    """Self-consistent solution for charge z at density parameter rs. Returns a record of
    plain numbers; if a dict is passed as fields it receives r, dn and V + Z/r on [0, R]."""
    p = dict(DEFAULTS)
    kf = (9.0 * PI / 4.0) ** (1.0 / 3.0) / rs
    R = max(p["R_min"], p["R_over_rs"] * rs) if R is None else R
    dx = p["dx"] if dx is None else dx
    nk = default_nk(kf, R) if nk is None else nk
    lmax = p["lmax"] if lmax is None else lmax
    tol = p["tol"] if tol is None else tol
    t_start = time.time()

    n0 = 3.0 / (4.0 * PI * rs**3)
    k_tf = math.sqrt(4.0 * kf / PI)
    mesh = Mesh(R, dx)
    nc, r = mesh.nc, mesh.r
    cont = Continuum(mesh, kf, nk, lmax)
    lmax_b = 0 if z <= 2 else 2
    vxc0 = float(v_xc(n0, xc))

    def v2_of(vs):
        v2 = np.zeros(mesh.N + 1)
        v2[: nc + 1] = -8.0 * z + 8.0 * r[: nc + 1] * vs
        veff = np.zeros(mesh.N + 1)
        veff[1 : nc + 1] = -z / r[1 : nc + 1] + vs[1:]
        return v2, veff, float(vs[0]), float((vs[8] - vs[0]) / r[8])

    def density(vs):
        v2, veff, v0, v1 = v2_of(vs)
        dn = cont.induced(v2, z, v0, v1)
        states, nb = bound_states(mesh, v2, veff, z, v0, v1, lmax_b)
        dn += nb
        # Value at the origin from the cusp condition n(r) = n(0) exp(-2 Z r).
        dn[0] = (n0 + dn[1]) * math.exp(2.0 * z * r[1]) - n0
        return dn, states

    def potential(dn):
        edge = 0.0
        if bc == "friedel":
            # Far from the nucleus dn = A cos(2 k_F r + phi) / r^3, and the solution of the
            # Poisson equation that decays is V_H = (pi / k_F^2) [dn - (r^3 dn)' / (k_F^2 r^4)]
            # up to relative order 1 / (k_F r)^2.
            x3 = r**3 * dn
            slope = (x3[nc + 1] - x3[nc - 1]) / (r[nc + 1] - r[nc - 1])
            edge = PI / kf**2 * (dn[nc] - slope / (kf**2 * mesh.Rc**4))
        vh, q = hartree_smooth(mesh, dn, z, edge)
        return vh + v_xc(n0 + dn[: nc + 1], xc) - vxc0, q

    rr = r[: nc + 1]
    vs = np.empty(nc + 1)
    vs[1:] = z * (-np.expm1(-k_tf * rr[1:])) / rr[1:]
    vs[0] = z * k_tf
    mixer = Anderson(beta=0.7 if z <= 2 else 0.4, depth=6, weight=rr**2 * 2.0 * mesh.x[: nc + 1] * dx)
    err, it = np.inf, 0
    for it in range(1, maxit + 1):
        dn, states = density(vs)
        vout, q = potential(dn)
        res = vout - vs
        err = float(np.max(np.abs(res)))
        if verbose:
            eb = ", ".join(f"{s['n']}{'spdf'[s['l']]} {s['E_Ha']:.6f}" for s in states)
            print(f"   it {it:3d}  max|dV| {err:9.2e}  Q(R)-Z {q[-1] - z:+.5f}  [{eb}]", flush=True)
        if err < tol:
            break
        vs = mixer.step(vs, screened_filter(mesh, res, k_tf))
    converged = err < tol

    if fields is not None:
        fields.update(r=rr.copy(), dn=dn[: nc + 1].copy(), v_smooth=vs.copy())
    v2, veff, v0, v1 = v2_of(vs)
    l_all = int(kf * mesh.Rc) + 12
    delta_all = phase_shifts(mesh, kf, v2, z, v0, v1, l_all)
    delta = delta_all[: lmax + 1]
    weights = 2 * np.arange(l_all + 1) + 1
    friedel = float(2.0 / PI * np.sum(weights[: lmax + 1] * delta))
    friedel_all = float(2.0 / PI * np.sum(weights * delta_all))
    orbitals = sum(2 * st["l"] + 1 for st in states)
    s = strength_sum(delta)
    # Levinson's theorem as a check on the absolute phase: delta_0(k -> 0) / pi = bound s states.
    levinson = float(phase_shifts(mesh, 1e-5, v2, z, v0, v1, 0)[0] / PI)
    # Bound states with l above those searched during the iteration (there should be none).
    extra = [st for st in bound_states(mesh, v2, veff, z, v0, v1, lmax_b + 2)[0] if st["l"] > lmax_b]
    # Charge inside the box, and its mean over the last Friedel period.
    last = rr >= mesh.Rc - PI / kf
    q_mean = float(np.trapezoid(q[last], rr[last]) / (rr[last][-1] - rr[last][0]))
    return {
        "Z": z,
        "rs": rs,
        "xc": xc,
        "kF_per_bohr": kf,
        "kF_per_A": kf / BOHR,
        "n0_per_bohr3": n0,
        "delta": [float(d) for d in delta],
        "delta0": float(delta[0]),
        "delta1": float(delta[1]),
        "delta2": float(delta[2]),
        "delta3": float(delta[3]),
        "friedel_sum": friedel,
        "friedel_sum_all_l": friedel_all,
        "l_all": l_all,
        "bound_states": states,
        "bound_orbitals": orbitals,
        "friedel_sum_continuum": friedel - 2 * orbitals,
        "levinson_delta0_over_pi_at_k0": levinson,
        "unexpected_bound_states": extra,
        "sum": s,
        "sum_terms": [float((l + 1) * math.sin(delta[l] - delta[l + 1]) ** 2) for l in range(4)],
        "h_over_e2": 4.0 * kf * s,
        "h_eV_A": 4.0 * kf * s * E2,
        "h_unitarity_eV_A": 4.0 * kf * E2,
        "friction_Q_au": 4.0 * kf * kf * s / (3.0 * PI),
        "density_at_nucleus_per_bohr3": float(n0 + dn[0]),
        "charge_in_box": float(q[-1]),
        "charge_in_box_period_mean": q_mean,
        "numerics": {"R_bohr": mesh.Rc, "dx": dx, "radial_points": mesh.N, "nk": nk, "lmax": lmax,
                     "tol_Ha": tol, "bc": bc, "iterations": it, "residual_Ha": err, "converged": converged,
                     "pairs": int(cont.l.size), "seconds": round(time.time() - t_start, 2)},
    }


# ----------------------------------------------------------------------------- comparisons

def yukawa_model(rs, lmax=8):
    """The model of packing.py solved with this file's radial code: V = -exp(-kappa r) / r
    with kappa fixed by the Friedel sum rule over l <= 8."""
    kf = (9.0 * PI / 4.0) ** (1.0 / 3.0) / rs
    mesh = Mesh(40.0, DEFAULTS["dx"])
    nc = mesh.nc
    r = mesh.r[: nc + 1]
    w = 2 * np.arange(lmax + 1) + 1

    def shifts(kappa):
        vs = np.empty(nc + 1)
        vs[1:] = -np.expm1(-kappa * r[1:]) / r[1:]
        vs[0] = kappa
        v2 = np.zeros(mesh.N + 1)
        v2[: nc + 1] = -8.0 + 8.0 * r * vs
        return phase_shifts(mesh, kf, v2, 1.0, float(vs[0]), float((vs[8] - vs[0]) / r[8]), lmax)

    kappa = brentq(lambda q: 2.0 / PI * np.sum(w * shifts(q)) - 1.0, 0.3, 5.0, xtol=1e-10)
    d = shifts(kappa)
    return {"kappa_per_bohr": kappa, "kappa_per_A": kappa / BOHR, "delta0": float(d[0]), "delta1": float(d[1]),
            "delta2": float(d[2]), "h_eV_A": 4.0 * kf * strength_sum(d) * E2}


def linear_response_check(rs, z=1e-3, xc="pz81"):
    """Weak-charge limit against the analytic result. For Z -> 0 the induced density is
    dn(q) / Z = k_TF^2 F / [q^2 (1 + f_xc N0 F) + k_TF^2 F], with F the static Lindhard
    function of q / 2 k_F, N0 = k_F / pi^2 and f_xc = d v_xc / d n at n0. The part of the
    self-consistent density that is odd in Z, [dn(Z) - dn(-Z)] / 2Z, is compared with the
    Fourier transform of that expression at a set of radii."""
    kf = (9.0 * PI / 4.0) ** (1.0 / 3.0) / rs
    n0 = 3.0 / (4.0 * PI * rs**3)
    fxc = float((v_xc(n0 * 1.0001, xc) - v_xc(n0 * 0.9999, xc)) / (2e-4 * n0))
    q = np.arange(0.0, 400.0 * kf, 5e-4 * kf)
    x = q[1:] / (2.0 * kf)
    f = np.ones_like(q)
    with np.errstate(divide="ignore", invalid="ignore"):
        f[1:] = 0.5 + (1.0 - x * x) / (4.0 * x) * np.log(np.abs((1.0 + x) / (1.0 - x)))
    f[1:][~np.isfinite(f[1:])] = 0.5
    k_tf2 = 4.0 * kf / PI
    g = k_tf2 * f / (q * q * (1.0 + fxc * kf / PI**2 * f) + k_tf2 * f)
    plus, minus = {}, {}
    solve(z, rs, xc=xc, fields=plus)
    solve(-z, rs, xc=xc, fields=minus)
    odd = (plus["dn"] - minus["dn"]) / (2.0 * z)
    rows = []
    for target in (0.25, 0.5, 1.0, 2.0, 4.0, 8.0):
        i = int(np.argmin(np.abs(plus["r"] - target)))
        r = float(plus["r"][i])
        exact = float(np.trapezoid(q * np.sin(q * r) * g, q) / (2.0 * PI**2 * r))
        rows.append({"r_bohr": r, "scf": float(odd[i]), "linear_response": exact})
    scale = max(abs(row["linear_response"]) for row in rows)
    return {"rs": rs, "Z": z, "xc": xc, "f_xc": fxc, "points": rows,
            "max_abs_deviation_over_peak": max(abs(row["scf"] - row["linear_response"]) for row in rows) / scale}


def pz_join_check():
    """What the join of the two PZ81 branches does at r_s = 1, the one density it affects."""
    global PZ_JOIN
    keep = PZ_JOIN
    out = {"joined": brief(solve(1, 1.0)), "pw92": brief(solve(1, 1.0, xc="pw92")),
           "rs_0.999": brief(solve(1, 0.999)), "rs_1.001": brief(solve(1, 1.001))}
    PZ_JOIN = 1e-9
    try:
        raw = solve(1, 1.0, maxit=150)
    finally:
        PZ_JOIN = keep
    out["as_published"] = dict(brief(raw), residual_Ha=raw["numerics"]["residual_Ha"])
    return out


def brief(rec):
    return {"delta0": rec["delta0"], "delta1": rec["delta1"], "delta2": rec["delta2"], "delta3": rec["delta3"],
            "friedel_sum": rec["friedel_sum"], "h_eV_A": rec["h_eV_A"], "h_over_e2": rec["h_over_e2"],
            "friction_Q_au": rec["friction_Q_au"],
            "density_at_nucleus_per_bohr3": rec["density_at_nucleus_per_bohr3"],
            "bound_states": [{"l": st["l"], "n": st["n"], "E_Ha": st["E_Ha"]} for st in rec["bound_states"]],
            "iterations": rec["numerics"]["iterations"], "converged": rec["numerics"]["converged"]}


def convergence(z, rs, base):
    """Repeat a run with each numerical parameter doubled (or the step and tolerance reduced)."""
    nm = base["numerics"]
    R = max(DEFAULTS["R_min"], DEFAULTS["R_over_rs"] * rs)
    # Doubling R also enlarges the k grid, which default_nk ties to k_F R.
    variants = {"R_x2": {"R": 2.0 * R}, "nk_x2": {"nk": 2 * nm["nk"]}, "lmax_x2": {"lmax": 2 * nm["lmax"]},
                "dx_half": {"dx": 0.5 * nm["dx"]}, "tol_tenth": {"tol": 0.1 * nm["tol_Ha"]},
                "VH_zero_at_R": {"bc": "zero"}}
    out = {}
    for name, kw in variants.items():
        r = solve(z, rs, **kw)
        out[name] = {"h_eV_A": r["h_eV_A"], "rel_change_h": r["h_eV_A"] / base["h_eV_A"] - 1.0,
                     "max_abs_change_delta_0_to_3": max(abs(r["delta"][l] - base["delta"][l]) for l in range(4)),
                     "friedel_sum": r["friedel_sum"], "iterations": r["numerics"]["iterations"],
                     "converged": r["numerics"]["converged"]}
    return out


# Numbers this calculation is compared with. Sources are given in jellium-notes.md.
REFEREE_PROTON = {  # reviews/scattering-theorist.md, finding 4: (delta0, delta1, h in eV A), PW92 correlation
    1.0: (0.630, 0.159, 25.8), 1.3: (0.766, 0.156, 30.4), 1.5: (0.852, 0.150, 32.8), 2.0: (1.047, 0.124, 36.3),
    2.07: (1.073, 0.120, 36.6), 2.5: (1.218, 0.091, 36.6), 3.0: (1.364, 0.055, 34.6), 4.0: (1.592, -0.013, 27.6),
    5.0: (1.748, -0.069, 21.0), 6.0: (1.850, -0.112, 16.3)}
REFEREE_BORON = {1.5: 148.8, 1.8: 120.7, 2.07: 101.1}  # h in eV A, same report, finding 3 and check 7
NAGY_ZAWADOWSKI = {"rs": 2.5, "delta0": 1.2213, "delta1": 0.0894}  # arXiv:0711.2575
OTHER_XC = ("pw92", "gl", "hl", "x", "hartree")
# Friction coefficient Q (a.u.) of a slow atom, Q = n0 k_F sigma_tr = k_F h / (3 pi), for Z = 1..6.
GERRITS_Q = {  # Gerrits, Juaristi and Meyer, Phys. Rev. B 102, 155130 (2020), data files; PZ-LDA, unpolarised
    1.5: (0.3107747628, 0.757697773, 0.9126072136, 1.105566359, 1.4081861533, 1.6897392286),
    2.0: (None, 0.4293148572, 0.4402121952, 0.5566089648, 0.74649094, 0.8777849525),
    2.5: (0.2075175302, 0.2399371683, 0.2226220501, 0.3077899564, 0.4494245063, 0.5191281872),
    3.5: (0.1270637556, 0.0792, 0.0634, 0.1339361302, 0.2330458677, 0.2569774621),
    5.0: (0.0599, 0.0187, 0.0228, 0.0842, 0.1313188297, 0.1352066438)}
NAGY_ALDAZABAL_Q = {  # arXiv:2006.05903, Table 1, Q^(1); Z >= 2 from the phase shifts of Puska and Nieminen
    1.5: (0.305, 0.755, 0.912, 1.112, 1.417, 1.692),  # (1983), Echenique et al. (1986), Nagy et al. (1989)
    2.0: (0.255, 0.427, 0.439, 0.557, 0.749, 0.874),
    3.0: (0.162, 0.135, 0.117, 0.191, 0.307, 0.346)}
NAGY_ZAWADOWSKI_ANTIPROTON = {"rs": 2.5, "delta0": -0.7729, "delta1": -0.2003}  # same paper, Z = -1
WHITMORE = {  # Whitmore, Carbotte and Shukla (1979) as reproduced in arXiv:2606.23065, Table II: delta_0..5 as
    # printed by Whitmore et al., and as recomputed there from Whitmore's fitted potential
    0.6: ((0.4518, 0.1445, 0.0578, 0.0257, 0.0123, 0.0062), (0.4244, 0.1435, 0.0577, 0.0257, 0.0123, 0.0062)),
    0.8: ((0.5874, 0.1572, 0.0541, 0.0209, 0.0087, 0.0038), (0.5311, 0.1559, 0.0541, 0.0209, 0.0087, 0.0038)),
    1.0: ((0.7241, 0.1616, 0.0488, 0.0167, 0.0062, 0.0024), (0.6267, 0.1602, 0.0488, 0.0167, 0.0062, 0.0025))}
ALMBLADH_N0 = {2.0: 0.522, 6.0: 0.335}  # density at the proton (1/bohr^3), HL functional, quoted in arXiv:2606.23065
TAKADA_RS4 = {"binding_Ha": 0.0115, "delta0_over_pi": 0.507, "bound_state_from_rs": 1.97}  # arXiv:1507.06432, LSDA


def fermi_level_product(rs, xc):
    """Maximum and first zero of R_0(k_F, r) R_1(k_F, r) for a proton, where R_l = u_l / (k_F r)
    has unit asymptotic amplitude. Nagy and Zawadowski quote both for r_s = 2.5."""
    f = {}
    rec = solve(1, rs, xc=xc, fields=f)
    kf = rec["kF_per_bohr"]
    mesh = Mesh(rec["numerics"]["R_bohr"], rec["numerics"]["dx"])
    nc, r, vs = mesh.nc, mesh.r, f["v_smooth"]
    v2 = np.zeros(mesh.N + 1)
    v2[: nc + 1] = -8.0 + 8.0 * r[: nc + 1] * vs
    l = np.array([0, 1])
    wv = Waves(mesh, l, np.full(2, kf * kf))
    w, t = wv.solve(v2, 1.0, float(vs[0]), float((vs[8] - vs[0]) / r[8]), free_logscale(l, kf))
    sa, ca = riccati(l, kf * r[mesh.na])
    sb, cb = riccati(l, kf * r[mesh.nb])
    ua, ub = wv.u_at(w, t, mesh.na), wv.u_at(w, t, mesh.nb)
    amp = np.hypot(ua * cb - ub * ca, sa * ub - sb * ua) / np.abs(sa * cb - sb * ca)
    big_r = np.sqrt(2.0 * mesh.x[1:nc, None]) * w[1:nc] / t[1:nc] / amp / (kf * r[1:nc, None])
    prod = big_r[:, 0] * big_r[:, 1]
    i = int(np.argmax(prod))
    zero = int(np.nonzero(prod[i:] < 0.0)[0][0]) + i
    return {"xc": xc, "maximum": float(prod[i]), "r_of_maximum": float(r[1 + i]), "first_zero": float(r[1 + zero]),
            "R0_at_origin": float(abs(big_r[0, 0]))}


def literature(run):
    """Comparison with published numbers. run(z, rs, xc) returns a (cached) record."""
    out = {}
    nz = []
    for z, ref in ((1, NAGY_ZAWADOWSKI), (-1, NAGY_ZAWADOWSKI_ANTIPROTON)):
        for xc in ("pz81", "pw92", "gl", "hl"):
            r = run(z, 2.5, xc)
            nz.append({"Z": z, "xc": xc, "delta0": r["delta0"], "delta1": r["delta1"],
                       "published_delta0": ref["delta0"], "published_delta1": ref["delta1"]})
    out["nagy_zawadowski_rs2.5"] = nz
    out["nagy_zawadowski_R0R1"] = {"published": {"maximum": 0.393, "r_of_maximum": 0.7, "first_zero": 2.53},
                                   "this_work": [fermi_level_product(2.5, xc) for xc in ("pz81", "gl")]}
    out["bound_state_onset_pz81"] = [
        {"rs": rs, "E_1s_Ha": (run(1, rs)["bound_states"] or [{"E_Ha": None}])[0]["E_Ha"],
         "friedel_sum": run(1, rs)["friedel_sum"]} for rs in (1.88, 1.90, 1.92, 1.94, 1.96)]
    out["gerrits_2020_Q"] = [
        {"rs": rs, "Z": z, "published": q, "pz81": run(z, rs, "pz81")["friction_Q_au"]}
        for rs, row in GERRITS_Q.items() for z, q in enumerate(row, start=1)]
    out["nagy_aldazabal_2020_Q"] = [
        {"rs": rs, "Z": z, "published": q, "pz81": run(z, rs, "pz81")["friction_Q_au"],
         "gl": run(z, rs, "gl")["friction_Q_au"]}
        for rs, row in NAGY_ALDAZABAL_Q.items() for z, q in enumerate(row, start=1)]
    out["whitmore_1979_phase_shifts"] = [
        {"rs": rs, "printed": list(a), "recomputed_in_arXiv_2606.23065": list(b),
         "pz81": run(1, rs, "pz81")["delta"][:6], "hl": run(1, rs, "hl")["delta"][:6]}
        for rs, (a, b) in WHITMORE.items()]
    out["almbladh_1976_density_at_proton"] = [
        {"rs": rs, "published_hl": n, "hl": run(1, rs, "hl")["density_at_nucleus_per_bohr3"],
         "pz81": run(1, rs, "pz81")["density_at_nucleus_per_bohr3"]} for rs, n in ALMBLADH_N0.items()]
    r4 = run(1, 4.0, "pz81")
    out["takada_2015_rs4"] = dict(TAKADA_RS4, pz81_binding_Ha=-r4["bound_states"][0]["E_Ha"],
                                  pz81_delta0_over_pi=r4["delta0"] / PI)
    return out


def table(headers, rows):
    out = ["| " + " | ".join(headers) + " |", "|" + "|".join("---" for _ in headers) + "|"]
    out += ["| " + " | ".join(str(c) for c in row) + " |" for row in rows]
    return "\n".join(out)


def states_text(rec):
    return "; ".join(f"{st['n']}{'spdf'[st['l']]} {st['E_Ha']:+.5f} Ha" for st in rec["bound_states"]) or "none"


def main_table(records):
    rows = []
    for r in records:
        rows.append([r["rs"], f"{r['kF_per_A']:.3f}", f"{r['delta0']:.4f}", f"{r['delta1']:+.4f}", f"{r['delta2']:+.4f}",
                     f"{r['delta3']:+.4f}", f"{r['friedel_sum']:.4f}", f"{r['friedel_sum_continuum']:+.4f}",
                     states_text(r), f"{r['h_eV_A']:.2f}", f"{r['h_over_e2']:.3f}", f"{r['h_unitarity_eV_A']:.1f}"])
    return table(["r_s", "k_F (1/A)", "delta_0", "delta_1", "delta_2", "delta_3", "Friedel sum", "continuum part",
                  "bound states", "h (eV A)", "h/e^2", "4 k_F e^2 (eV A)"], rows)


def show(rec):
    nm = rec["numerics"]
    print(f"Z={rec['Z']} rs={rec['rs']:<5} d0..3 = {rec['delta0']:.4f} {rec['delta1']:+.4f} {rec['delta2']:+.4f} "
          f"{rec['delta3']:+.4f}  Friedel {rec['friedel_sum']:.4f}  h = {rec['h_eV_A']:.2f} eV A "
          f"({rec['h_over_e2']:.3f} e2; unitarity {rec['h_unitarity_eV_A']:.1f})  bound: {states_text(rec)}  "
          f"[{nm['iterations']} it, {nm['seconds']} s{'' if nm['converged'] else ', NOT CONVERGED'}]", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--quick", action="store_true", help="tables only; no checks, nothing written")
    args = ap.parse_args()
    t0 = time.time()
    packing = {}
    if (HERE / "packing.json").exists():
        packing = {g["rs"]: g["h_eV_A"] for g in json.loads((HERE / "packing.json").read_text())["electron_gas"]}

    cache = {}

    def run(z, rs, xc="pz81"):
        if (z, rs, xc) not in cache:
            cache[(z, rs, xc)] = solve(z, rs, xc=xc)
        return cache[(z, rs, xc)]

    records = []
    for z, densities in ((1, RS_PROTON), (5, RS_BORON)):
        for rs in densities:
            rec = run(z, rs)
            show(rec)
            if not args.quick:
                rec["convergence"] = convergence(z, rs, rec)
                if z == 1:
                    rec["other_functionals"] = {xc: brief(run(z, rs, xc)) for xc in OTHER_XC}
                    rec["yukawa_model"] = yukawa_model(rs)
                    rec["yukawa_model"]["h_eV_A_packing_json"] = packing.get(rs)
                    rec["referee"] = dict(zip(("delta0", "delta1", "h_eV_A"), REFEREE_PROTON[rs])) if rs in REFEREE_PROTON else None
                else:
                    rec["other_functionals"] = {"pw92": brief(run(z, rs, "pw92"))}
                    rec["referee"] = {"h_eV_A": REFEREE_BORON[rs]}
            records.append(rec)
    protons = [r for r in records if r["Z"] == 1]
    borons = [r for r in records if r["Z"] == 5]
    print("\nZ = 1, PZ81\n" + main_table(protons))
    print("\nZ = 5, PZ81\n" + main_table(borons))
    if args.quick:
        print(f"\n{time.time() - t0:.0f} s")
        return

    print("\nConvergence: relative change of h in parts per million when one parameter is changed")
    names = ["R_x2", "nk_x2", "lmax_x2", "dx_half", "tol_tenth", "VH_zero_at_R"]
    rows = [[r["Z"], r["rs"]] + [f"{1e6 * r['convergence'][n]['rel_change_h']:+.1f}" for n in names]
            + [f"{r['convergence']['VH_zero_at_R']['friedel_sum']:.4f}"] for r in records]
    print(table(["Z", "r_s"] + names + ["Friedel sum with V_H(R) = 0"], rows))

    print("\nChecks on the converged solutions")
    rows = [[r["Z"], r["rs"], f"{r['friedel_sum']:.5f}", f"{r['friedel_sum_all_l']:.5f}", r["l_all"],
             f"{r['charge_in_box_period_mean']:.5f}", f"{r['levinson_delta0_over_pi_at_k0']:.4f}",
             sum(1 for st in r["bound_states"] if st["l"] == 0), len(r["unexpected_bound_states"]),
             f"{r['density_at_nucleus_per_bohr3']:.4f}", r["numerics"]["iterations"],
             f"{r['numerics']['residual_Ha']:.1e}"] for r in records]
    print(table(["Z", "r_s", "Friedel sum l<=12", "Friedel sum all l", "l_all", "charge inside R (period mean)",
                 "delta_0(k->0)/pi", "bound s states", "bound l>0 states missed", "n(0) (1/bohr^3)", "iterations",
                 "residual (Ha)"], rows))

    print("\nZ = 1: other functionals, h in eV A (delta_0, delta_1)")
    rows = [[r["rs"], f"{r['h_eV_A']:.2f} ({r['delta0']:.4f}, {r['delta1']:+.4f})"]
            + [f"{r['other_functionals'][xc]['h_eV_A']:.2f} ({r['other_functionals'][xc]['delta0']:.4f}, "
               f"{r['other_functionals'][xc]['delta1']:+.4f})" for xc in OTHER_XC] for r in protons]
    print(table(["r_s", "pz81"] + list(OTHER_XC), rows))

    print("\nZ = 1 against the referee's table (PW92) and the Yukawa model")
    rows = []
    for r in protons:
        ref, pw, yk = r["referee"], r["other_functionals"]["pw92"], r["yukawa_model"]
        rows.append([r["rs"], f"{r['h_eV_A']:.2f}", f"{pw['h_eV_A']:.2f}", "" if ref is None else ref["h_eV_A"],
                     "" if ref is None else f"{100 * (pw['h_eV_A'] / ref['h_eV_A'] - 1):+.2f}",
                     "" if ref is None else f"{pw['delta0'] - ref['delta0']:+.4f}",
                     "" if ref is None else f"{pw['delta1'] - ref['delta1']:+.4f}",
                     f"{yk['h_eV_A']:.2f}", "" if yk["h_eV_A_packing_json"] is None else f"{yk['h_eV_A_packing_json']:.2f}",
                     f"{100 * (yk['h_eV_A'] / r['h_eV_A'] - 1):+.1f}", f"{r['h_eV_A'] / yk['h_eV_A']:.3f}"])
    print(table(["r_s", "h pz81", "h pw92", "h referee", "pw92 vs referee (%)", "d delta_0", "d delta_1",
                 "h Yukawa (this code)", "h Yukawa (packing.json)", "Yukawa vs pz81 (%)", "pz81 / Yukawa"], rows))

    lit = literature(run)
    print("\nNagy and Zawadowski, r_s = 2.5: proton 1.2213, 0.0894; antiproton -0.7729, -0.2003")
    rows = [[e["Z"], e["xc"], f"{e['delta0']:.4f}", f"{e['delta1']:.4f}", f"{e['delta0'] - e['published_delta0']:+.4f}",
             f"{e['delta1'] - e['published_delta1']:+.4f}"] for e in lit["nagy_zawadowski_rs2.5"]]
    print(table(["Z", "functional", "delta_0", "delta_1", "d delta_0", "d delta_1"], rows))

    wf = lit["nagy_zawadowski_R0R1"]
    print("R_0 R_1 at k_F: published maximum 0.393 at r = 0.7, zero at r = 2.53; this work: " + "; ".join(
        f"{e['xc']} maximum {e['maximum']:.4f} at r = {e['r_of_maximum']:.3f}, zero at r = {e['first_zero']:.3f}"
        for e in wf["this_work"]))
    print("Onset of the bound state (pz81): " + ", ".join(
        f"r_s = {e['rs']}: " + ("none" if e["E_1s_Ha"] is None else f"{e['E_1s_Ha']:.1e} Ha")
        for e in lit["bound_state_onset_pz81"]))

    print("\nFriction coefficient Q (a.u.) against Gerrits, Juaristi and Meyer (PZ-LDA): this work (published, % difference)")
    rows = []
    for rs in GERRITS_Q:
        cells = [e for e in lit["gerrits_2020_Q"] if e["rs"] == rs]
        rows.append([rs] + [f"{e['pz81']:.4f}" + ("" if e["published"] is None else
                            f" ({e['published']:.4f}, {100 * (e['pz81'] / e['published'] - 1):+.2f})") for e in cells])
    print(table(["r_s"] + [f"Z = {z}" for z in range(1, 7)], rows))

    print("\nQ (a.u.) against Nagy and Aldazabal, Table 1: published; this work pz81 (%); this work gl (%)")
    rows = []
    for rs in NAGY_ALDAZABAL_Q:
        cells = [e for e in lit["nagy_aldazabal_2020_Q"] if e["rs"] == rs]
        rows.append([rs] + [f"{e['published']:.3f}; {e['pz81']:.4f} ({100 * (e['pz81'] / e['published'] - 1):+.2f}); "
                            f"{e['gl']:.4f} ({100 * (e['gl'] / e['published'] - 1):+.2f})" for e in cells])
    print(table(["r_s"] + [f"Z = {z}" for z in range(1, 7)], rows))

    print("\nProton at high density against Whitmore et al. (delta_0 .. delta_5)")
    rows = []
    for e in lit["whitmore_1979_phase_shifts"]:
        for name in ("printed", "recomputed_in_arXiv_2606.23065", "pz81", "hl"):
            d = e[name]
            rows.append([e["rs"], name] + [f"{x:.4f}" for x in d]
                        + [f"{2.0 / PI * sum((2 * l + 1) * d[l] for l in range(6)):.4f}"])
    print(table(["r_s", "source", "d0", "d1", "d2", "d3", "d4", "d5", "Friedel sum l<=5"], rows))

    print("\nOther published numbers")
    rows = [[f"n(0) at r_s = {e['rs']} (1/bohr^3), Almbladh et al. (HL)", e["published_hl"],
             f"hl {e['hl']:.4f}, pz81 {e['pz81']:.4f}"] for e in lit["almbladh_1976_density_at_proton"]]
    t4 = lit["takada_2015_rs4"]
    rows.append(["1s binding at r_s = 4 (Ha), Takada et al. (LSDA)", t4["binding_Ha"], f"pz81 {t4['pz81_binding_Ha']:.5f}"])
    rows.append(["delta_0 / pi at r_s = 4, Takada et al. (LSDA)", t4["delta0_over_pi"], f"pz81 {t4['pz81_delta0_over_pi']:.4f}"])
    print(table(["quantity", "published", "this work"], rows))

    print("\nZ = 5 against the referee (PW92)")
    rows = [[r["rs"], f"{r['h_eV_A']:.2f}", f"{r['other_functionals']['pw92']['h_eV_A']:.2f}", r["referee"]["h_eV_A"],
             f"{100 * r['sum_terms'][1] / r['sum']:.1f}", f"{r['friedel_sum']:.4f}", states_text(r)] for r in borons]
    print(table(["r_s", "h pz81", "h pw92", "h referee", "l = 1 term (% of sum)", "Friedel sum", "bound states"], rows))

    print("\nWeak charge (Z = +-0.001) against analytic linear response: induced density per unit charge")
    linear = [linear_response_check(rs) for rs in (2.0, 4.0)]
    rows = [[c["rs"], f"{pt['r_bohr']:.3f}", f"{pt['scf']:+.6e}", f"{pt['linear_response']:+.6e}",
             f"{pt['scf'] / pt['linear_response']:.5f}"] for c in linear for pt in c["points"]]
    print(table(["r_s", "r (bohr)", "self-consistent", "linear response", "ratio"], rows))

    print("\nr_s = 1: join of the PZ81 branches")
    join = pz_join_check()
    rows = [[name, f"{v['h_eV_A']:.4f}", f"{v['delta0']:.5f}", f"{v['delta1']:.5f}", f"{v['friedel_sum']:.4f}",
             v["iterations"], "yes" if v["converged"] else f"no, residual {v['residual_Ha']:.1e} Ha"]
            for name, v in join.items()]
    print(table(["run", "h (eV A)", "delta_0", "delta_1", "Friedel sum", "iterations", "converged"], rows))

    out = {
        "description": "Kohn-Sham LDA (Perdew-Zunger 1981) screening of a point charge Z in jellium; see jellium.py",
        "units": {"energy": "hartree unless marked", "length": "bohr unless marked", "h": "eV angstrom",
                  "phase shifts": "radian, absolute (zero at infinite energy)"},
        "friedel_counting": "friedel_sum = (2/pi) sum_{l<=lmax} (2l+1) delta_l(k_F) with absolute phase shifts, "
                            "which by Levinson's theorem include pi for each bound state of that l, so it should "
                            "equal Z. friedel_sum_continuum = friedel_sum - 2 * bound_orbitals is the charge "
                            "displaced in the continuum and should equal Z minus two electrons per bound orbital "
                            "(a bound s state counts as one orbital, two electrons).",
        "e2_eV_A": E2,
        "defaults": dict(DEFAULTS, pz_join_half_width=PZ_JOIN),
        "literature": lit,
        "linear_response_check": linear,
        "pz81_join_check_rs1": join,
        "records": records,
        "seconds": round(time.time() - t0, 1),
    }
    (HERE / "jellium.json").write_text(json.dumps(out, indent=1) + "\n")
    print(f"\nwrote {HERE / 'jellium.json'} in {time.time() - t0:.0f} s")


if __name__ == "__main__":
    main()