"""AC78 isolated-halo benchmark: the fuzzy dark matter ground state and the isothermal core.

Needs : Python 3.9+, NumPy, SciPy. No input file.
Run   : python3 benchmark_isolated_soliton.py isolated-output.json

Part A. Schrodinger-Poisson ground state (the "soliton").
  In units with hbar / m = 1 and G = 1 the stationary equations are
      psi'' + (2/r) psi' = 2 u psi,      u'' + (2/r) u' = 4 pi psi^2,
  where u = Phi - E. We fix psi(0) = 1 and shoot on u(0) until psi decays
  without a node. The equations have the scaling symmetry
      r -> r / lam,  rho -> lam^4 rho,
  so rho_0 r_c^4 G m^2 / hbar^2 is one pure number for every ground state.
  Schive et al. (2014, PRL 113, 261302, arXiv:1407.7762v2, eq. 3) give
      rho_0 = 1.9 (m / 1e-23 eV)^-2 (r_c / kpc)^-4  Msun / pc^3   at z = 0,
  which corresponds to 1.9e9 / 8.5454e9 = 0.2223 in these units
  (hbar^2 / (G m^2) = 8.5454e9 Msun kpc for m = 1e-23 eV).
  We also test their shape [1 + 0.091 (r / r_c)^2]^-8, their statement that
  the fit is accurate to about 2 per cent inside 3 r_c, and the mass fractions.

Part B. Isothermal core used by the self-interacting model.
  The dimensionless isothermal sphere psi'' + (2/x) psi' = exp(-psi),
  psi(0) = psi'(0) = 0, is compared with its series
  psi = x^2/6 - x^4/120 + x^6/1890 near the centre and with the singular
  isothermal sphere exp(-psi) -> 2 / x^2 far out.

Every number is compared at two resolutions so the reader sees convergence.
"""
import json
import sys

import numpy as np
from scipy.integrate import solve_ivp, cumulative_trapezoid

HBAR2_OVER_G_M2_MSUN_KPC = 8.5454e9  # for m = 1e-23 eV


def shoot(u0, r_max, tolerance):
    def rhs(r, y):
        psi, dpsi, u, du = y
        return [dpsi, -2 * dpsi / r + 2 * u * psi, du, -2 * du / r + 4 * np.pi * psi * psi]

    r0 = 1e-6
    start = [1.0, 2 * u0 * r0 / 3, u0 + 4 * np.pi * r0 * r0 / 6, 4 * np.pi * r0 / 3]

    def node(r, y):
        return y[0]
    node.terminal = True

    def turn_up(r, y):
        return y[1]
    turn_up.terminal = True
    turn_up.direction = 1
    solution = solve_ivp(rhs, (r0, r_max), start, method='DOP853', rtol=tolerance, atol=tolerance * 1e-3,
                         events=[node, turn_up], dense_output=True)
    crossed = len(solution.t_events[0]) > 0
    return solution, crossed


def ground_state(tolerance, r_max=6.0):
    low, high = -5.0, 0.0  # very negative u0 -> node; too high -> psi turns back up
    for _ in range(80):
        middle = (low + high) / 2
        _, crossed = shoot(middle, r_max, tolerance)
        if crossed:
            low = middle
        else:
            high = middle
    solution, _ = shoot(low, r_max, tolerance)
    radius = np.linspace(1e-6, solution.t[-1], 40001)
    density = solution.sol(radius)[0] ** 2
    return low, radius, density


def soliton_numbers(tolerance):
    u0, radius, density = ground_state(tolerance)
    core = float(np.interp(-0.5, -density, radius))
    enclosed = cumulative_trapezoid(4 * np.pi * radius ** 2 * density, radius, initial=0.0)
    # The profile is cut where psi has fallen by many e-folds; add nothing beyond it.
    total = float(enclosed[-1])
    scaled = radius / core
    fit = (1 + 0.091 * scaled ** 2) ** -8
    inside = scaled <= 3
    fit_error = float(np.max(np.abs(fit[inside] - density[inside]) / density[inside]))
    shape_coefficient = 2 ** (1 / 8) - 1
    mass_core = float(np.interp(core, radius, enclosed))
    return {
        'u0': u0,
        'invariantRho0Rc4': core ** 4,
        'schiveInvariant': 1.9e9 / HBAR2_OVER_G_M2_MSUN_KPC,
        'maxFitErrorInside3Rc': fit_error,
        'exactHalfDensityCoefficient': shape_coefficient,
        'massFractionInside3Rc': float(np.interp(3 * core, radius, enclosed)) / total,
        'halfMassRadiusOverRc': float(np.interp(0.5 * total, enclosed, radius)) / core,
        'coreMassTimesRcMsunKpcAt1e-23eV': mass_core * core * HBAR2_OVER_G_M2_MSUN_KPC,
        'totalMassOverCoreMass': total / mass_core,
        'profileRadiusOverRcCut': float(radius[-1] / core),
    }


def isothermal(tolerance):
    def rhs(x, y):
        return [y[1], np.exp(-y[0]) - 2 * y[1] / x]

    x0 = 1e-5
    solution = solve_ivp(rhs, (x0, 2000.0), [x0 * x0 / 6, x0 / 3], method='DOP853', rtol=tolerance,
                         atol=tolerance * 1e-3, dense_output=True)
    small = np.array([0.05, 0.1, 0.2, 0.4])
    series = small ** 2 / 6 - small ** 4 / 120 + small ** 6 / 1890
    far = np.array([200.0, 600.0, 2000.0])
    density_far = np.exp(-solution.sol(far)[0])
    half_density_radius = float(np.interp(-0.5, -np.exp(-solution.sol(np.linspace(x0, 10, 20001))[0]),
                                          np.linspace(x0, 10, 20001)))
    return {
        'seriesMaxError': float(np.max(np.abs(solution.sol(small)[0] - series))),
        'farDensityTimesXSquaredOver2': [float(v) for v in density_far * far ** 2 / 2],
        'halfDensityRadius': half_density_radius,
    }


def main(output_path):
    coarse, fine = soliton_numbers(1e-7), soliton_numbers(1e-10)
    iso_coarse, iso_fine = isothermal(1e-7), isothermal(1e-11)
    result = {
        'soliton': {'tolerance1e-7': coarse, 'tolerance1e-10': fine},
        'isothermal': {'tolerance1e-7': iso_coarse, 'tolerance1e-11': iso_fine},
        'versions': {'numpy': np.__version__, 'scipy': __import__('scipy').__version__, 'python': sys.version.split()[0]},
    }
    with open(output_path, 'w', encoding='utf-8') as handle:
        json.dump(result, handle, indent=1)
        handle.write('\n')
    for key, value in fine.items():
        print('soliton %-36s %.6g (coarse %.6g)' % (key, value, coarse[key]))
    print('isothermal series error', iso_fine['seriesMaxError'], 'far rho x^2/2', iso_fine['farDensityTimesXSquaredOver2'],
          'half-density x', iso_fine['halfDensityRadius'])


if __name__ == '__main__':
    main(sys.argv[1] if len(sys.argv) > 1 else 'isolated-output.json')