"""Scattering strength of one proton in an electron gas, and the ceiling it implies.

For electrons scattering from a single proton the Fermi-surface average of the squared
force matrix element is fixed by the transport cross-section, which gives

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

with delta_l the phase shifts at the Fermi momentum. The proton is modelled as a screened
Coulomb potential -e^2 exp(-kappa r) / r whose screening length is fixed by the Friedel sum
rule (2 / pi) sum_l (2l + 1) delta_l = 1. Phase shifts come from the variable-phase equation

    d delta_l / dr = -(1 / k) U(r) [cos(delta_l) j_l(kr) - sin(delta_l) n_l(kr)]^2

with Riccati-Bessel functions and U = 2 m V / hbar^2. Atomic units inside; eV and
angstrom in the output.
"""

import json
import math
from pathlib import Path

import numpy as np
from scipy.integrate import solve_ivp
from scipy.optimize import brentq
from scipy.special import spherical_jn, spherical_yn

HERE = Path(__file__).parent
E2 = 14.399645  # e^2 in eV angstrom = 1 hartree bohr
BOHR = 0.529177211
L_MAX = 8
M_H_OMEGA2_1000K = 1.7904  # eV/A^2, from requirements.py


def phase_shift(l, k, kappa):
    def rhs(r, d):
        x = k * r
        jl = x * spherical_jn(l, x)
        nl = x * spherical_yn(l, x)
        u = -2.0 * math.exp(-kappa * r) / r
        return [-(u / k) * (math.cos(d[0]) * jl - math.sin(d[0]) * nl) ** 2]

    r_max = 40.0 / kappa
    sol = solve_ivp(rhs, [1e-7, r_max], [0.0], method="LSODA", rtol=1e-10, atol=1e-12)
    return sol.y[0, -1]


def shifts(k, kappa):
    return np.array([phase_shift(l, k, kappa) for l in range(L_MAX + 1)])


def friedel(k, kappa):
    d = shifts(k, kappa)
    return (2.0 / math.pi) * sum((2 * l + 1) * d[l] for l in range(L_MAX + 1))


def strength(d):
    return sum((l + 1) * math.sin(d[l] - d[l + 1]) ** 2 for l in range(L_MAX))


def proton_in_gas(rs):
    kf = (9.0 * math.pi / 4.0) ** (1.0 / 3.0) / rs
    kappa = brentq(lambda q: friedel(kf, q) - 1.0, 0.2, 6.0, xtol=1e-6)
    d = shifts(kf, kappa)
    s = strength(d)
    k_tf = math.sqrt(4.0 * kf / math.pi)
    # First Born approximation with Thomas-Fermi screening, for comparison:
    # sigma_tr k^2 / (4 pi) = (1 / (4 k^2)) * integral_0^{2k} q^3 |v(q)|^2 dq / (2 pi ... )
    q = np.linspace(1e-6, 2 * kf, 20001)
    vq = 4.0 * math.pi / (q**2 + k_tf**2)
    born_sum = np.trapezoid(q**3 * vq**2, q) / (16.0 * math.pi**2 * kf**2)
    return {
        "rs": rs,
        "kF_per_A": kf / BOHR,
        "kappa_per_A": kappa / BOHR,
        "kappa_TF_per_A": k_tf / BOHR,
        "delta0": d[0],
        "delta1": d[1],
        "delta2": d[2],
        "sum": s,
        "h_eV_A": 4.0 * kf * s * E2,
        "h_unitarity_eV_A": 4.0 * kf * E2,
        "h_born_TF_eV_A": 4.0 * kf * born_sum * E2,
    }


def ceiling(eta):
    """Asymptotic Tc in kelvin for a hydrogen Hopfield parameter in eV/A^2."""
    return 182.7 * math.sqrt(eta / M_H_OMEGA2_1000K)


if __name__ == "__main__":
    gas = [proton_in_gas(rs) for rs in (1.0, 1.5, 2.0, 2.07, 2.5, 3.0, 4.0)]
    out = {"electron_gas": gas, "e2_eV_A": E2, "ceiling_K_per_sqrt_eta": ceiling(1.0)}
    (HERE / "packing.json").write_text(json.dumps(out, indent=1))
    print(f"{'rs':>5} {'kF':>6} {'kappa':>6} {'kTF':>6} {'d0':>6} {'d1':>6} {'d2':>6} {'sum':>6} {'h':>6} {'h/e2':>5} {'unit.':>6} {'Born':>6}")
    for g in gas:
        print(f"{g['rs']:5.2f} {g['kF_per_A']:6.3f} {g['kappa_per_A']:6.3f} {g['kappa_TF_per_A']:6.3f} {g['delta0']:6.3f} {g['delta1']:6.3f} {g['delta2']:6.3f} {g['sum']:6.3f} {g['h_eV_A']:6.1f} {g['h_eV_A']/E2:5.2f} {g['h_unitarity_eV_A']:6.1f} {g['h_born_TF_eV_A']:6.1f}")
    print(f"\nceiling = {ceiling(1.0):.1f} K * sqrt(eta_H in eV/A^2) = {ceiling(1.0):.1f} K * sqrt(rho_H * h)")
    for rho in (0.058, 0.081, 0.090, 0.100, 0.118, 0.231):
        row = [f"{ceiling(rho * h):5.0f}" for h in (29, 44, 64)]
        print(f"rho_H = {rho:.3f} per A^3: ceiling for h = 29, 44, 64 eV A: {' '.join(row)} K")