"""Teaching values for the dark-matter-small-scales readings.

Written separately from the Lab engine (TypeScript) so the two can be compared.
Units: kpc, solar masses, km/s, Gyr. Run:
    python3 -I compute_teaching_values.py reference-grid.json linear-output.json teaching-values.json
"""
import json
import math
import sys

import numpy as np
import scipy
from scipy.integrate import solve_ivp
from scipy.optimize import brentq

G = 4.30091e-6                 # kpc (km/s)^2 / Msun
KPC_PER_KMS_GYR = 1.02271      # distance covered at 1 km/s in 1 Gyr
GYR_PER_KPC_OVER_KMS = 0.977792
CM2_PER_G = 2.0884e-10         # 1 cm^2/g in kpc^2/Msun


def load(path):
    with open(path, encoding='utf-8') as handle:
        return json.load(handle)


grid_path, linear_path, output_path = sys.argv[1:4]
grid = load(grid_path)
linear = load(linear_path)
h = grid['parameters']['h']
omega_m = grid['omegaMatter']
omega_cdm = grid['omegaCdm'] / h ** 2
rho_crit = 277.537 * h ** 2
x = omega_m - 1
delta_vir = 18 * math.pi ** 2 + 82 * x - 39 * x * x


def nfw(mass):
    concentration = 10 ** (1.025 - 0.097 * math.log10(mass * h / 1e12))
    r_vir = (3 * mass / (4 * math.pi * delta_vir * rho_crit)) ** (1 / 3)
    r_s = r_vir / concentration
    mu = math.log(1 + concentration) - concentration / (1 + concentration)
    rho_s = mass / (4 * math.pi * r_s ** 3 * mu)
    return dict(mass=mass, c=concentration, r_vir=r_vir, r_s=r_s, rho_s=rho_s)


def nfw_density(halo, r):
    y = r / halo['r_s']
    return halo['rho_s'] / (y * (1 + y) ** 2)


def nfw_mass(halo, r):
    y = r / halo['r_s']
    return 4 * math.pi * halo['rho_s'] * halo['r_s'] ** 3 * (math.log1p(y) - y / (1 + y))


def nfw_slope(halo, r):
    y = r / halo['r_s']
    return -1 - 2 * y / (1 + y)


halo = nfw(1e10)
v1 = math.sqrt(G * nfw_mass(halo, 1.0) / 1.0)
x_max = brentq(lambda y: math.log1p(y) - y / (1 + y) - y * y / (1 + y) ** 2, 1, 4)
v_max = math.sqrt(G * nfw_mass(halo, x_max * halo['r_s']) / (x_max * halo['r_s']))

# A uniform core of the self-interacting central density, for the rotation-curve comparison.
uniform_core_density = 4.0e7  # Msun / kpc^3 = 0.04 Msun / pc^3
uniform_core_v_half_kpc = math.sqrt(G * 4 * math.pi / 3 * uniform_core_density * 0.5 ** 2)

# Warm: Viel et al. (2005) eq. 7 and Schneider et al. (2012) eqs. 28 and 39.
NU = 1.12


def viel_alpha(mass_kev, omega_warm, little_h):
    return 0.049 * mass_kev ** -1.11 * (omega_warm / 0.25) ** 0.11 * (little_h / 0.7) ** 1.22


def half_mode_k(mass_kev, omega_warm, little_h):
    return (2 ** (NU / 5) - 1) ** (1 / (2 * NU)) / viel_alpha(mass_kev, omega_warm, little_h)


def half_mode_mass(k_h, little_h, matter):
    radius = math.pi / k_h * 1000 / little_h
    return 4 * math.pi / 3 * matter * 277.537 * little_h ** 2 * radius ** 3


k_hm_3 = half_mode_k(3, omega_cdm, h)
m_hm_3 = half_mode_mass(k_hm_3, h, omega_m)
warm = dict(
    alpha3keV=viel_alpha(3, omega_cdm, h), halfModeK3keV=k_hm_3, halfModeMass3keV=m_hm_3,
    abundance3keVAt1e9=(1 + m_hm_3 / 1e9) ** -1.16, abundance3keVAt1e8=(1 + m_hm_3 / 1e8) ** -1.16,
    concentrationRatio3keVAt1e9=(1 + 15 * m_hm_3 / 1e9) ** -0.3,
    classHalfModeK3keV=next(row['classHalfModeK'] for row in linear['warmHalfModes'] if row['massKeV'] == 3.0),
    halfModeMass3keVNadlerCosmology=half_mode_mass(half_mode_k(3, 0.24, 0.7), 0.7, 0.286),
)

# Nadler et al. (2021) eqs. 4, 5 and 7, read from the published letter.
nadler = dict(
    massFromHalfMode3p8e7=3 * (3.8e7 / 5e8) ** (-3 / 10),
    massFromHalfMode3p0e7=3 * (3.0e7 / 5e8) ** (-3 / 10),
    fuzzyMassFromM0_1p8e8=(1.8e8 / 1.6e10) ** (-3 / 4),
    fuzzyMassFromM0_1p4e8=(1.4e8 / 1.6e10) ** (-3 / 4),
    subhaloSuppressionAt1e8=(1 + 2.7 * 3.8e7 / 1e8) ** -0.99,
    subhaloSuppressionAt1e9=(1 + 2.7 * 3.8e7 / 1e9) ** -0.99,
)

# Fuzzy: Hu et al. (2000) eqs. 8-9, Schive et al. (2014) eqs. 3 and 7, Schive et al. (2016) eq. 7.


def soliton(m22, mass, r):
    r_c = 1.6 / m22 * (mass / 1e9) ** (-1 / 3)
    rho_0 = 1.9e9 / (10 * m22) ** 2 / r_c ** 4
    return r_c, rho_0, rho_0 / (1 + 0.091 * (r / r_c) ** 2) ** 8


r_c_158, rho_0_158, rho_150_158 = soliton(10 ** 0.2, 1e10, 0.15)
r_c_1, rho_0_1, _ = soliton(1.0, 1e10, 0.15)
fuzzy = dict(
    m22=10 ** 0.2, coreRadius=r_c_158, centralDensity=rho_0_158, densityAt150pc=rho_150_158,
    coreRadiusM22Is1=r_c_1, centralDensityM22Is1=rho_0_1,
    powerHalfPerMpcM22Is1=4.5, jeansAtEqualityPerMpcM22Is1=9.0,
    abundanceM22Is1At1e9=(1 + (1e9 / 1.6e10) ** -1.1) ** -2.2,
    abundanceM22Is1At1e10=(1 + (1e10 / 1.6e10) ** -1.1) ** -2.2,
    nfwDensityAt150pc=nfw_density(halo, 0.15),
)

# Self-interacting: isothermal sphere matched at r1 (Kaplinghat, Tulin & Yu 2016).


def isothermal(xi, state):
    potential, slope = state
    return [slope, 0.0 if xi == 0 else math.exp(-potential) - 2 * slope / xi]


start = 1e-6
iso = solve_ivp(isothermal, (start, 60), [start ** 2 / 6, start / 3], dense_output=True, rtol=1e-11, atol=1e-13, max_step=0.01)


def iso_ratio(xi):
    potential, slope = iso.sol(xi)
    return xi * xi * slope / (xi ** 3 * math.exp(-potential))


xi_grid = np.linspace(0.01, 60, 60000)
ratios = np.array([iso_ratio(value) for value in xi_grid])
first_max = int(np.argmax(np.diff(ratios) < 0))


def match_sidm(halo_model, cross_section):
    sigma = cross_section * CM2_PER_G
    dispersion = math.sqrt(G * nfw_mass(halo_model, halo_model['r_s']) / halo_model['r_s'] / 2)
    for _ in range(400):
        target_density = 1 / (sigma * 4 / math.sqrt(math.pi) * dispersion * KPC_PER_KMS_GYR * 10)
        r1 = brentq(lambda r: nfw_density(halo_model, r) - target_density, 1e-6, 1e3 * halo_model['r_s'])
        target = nfw_mass(halo_model, r1) / (4 * math.pi * r1 ** 3 * nfw_density(halo_model, r1))
        xi1 = brentq(lambda xi: iso_ratio(xi) - target, xi_grid[0], xi_grid[first_max])
        scale = r1 / xi1
        rho_0 = nfw_density(halo_model, r1) * math.exp(iso.sol(xi1)[0])
        new = scale * math.sqrt(4 * math.pi * G * rho_0)
        if abs(new / dispersion - 1) < 1e-12:
            break
        dispersion = 0.5 * (dispersion + new)
    return dict(r1=r1, xi1=xi1, scale=scale, centralDensity=rho_0, dispersion=new,
                densityAt150pc=rho_0 * math.exp(-iso.sol(0.15 / scale)[0]),
                speedAt1kpc=math.sqrt(G * 4 * math.pi * rho_0 * scale ** 3 * (1 / scale) ** 2 * iso.sol(1 / scale)[1] / 1.0),
                maximumRatio=float(ratios[first_max]))


sidm = match_sidm(halo, 2.0)

# Gas feedback: coreNFW (Read et al. 2016, eqs. 16-21).
t_dyn = 2 * math.pi * math.sqrt(halo['r_s'] ** 3 / (G * nfw_mass(halo, halo['r_s']))) * GYR_PER_KPC_OVER_KMS
n = math.tanh(0.04 * 10 / t_dyn)
r_core = 1.75 * 1.5


def core_nfw_density(r):
    f = math.tanh(r / r_core)
    return f ** n * nfw_density(halo, r) + n * f ** (n - 1) * (1 - f * f) * nfw_mass(halo, r) / (4 * math.pi * r * r * r_core)


feedback = dict(dynamicalTime=t_dyn, exponent=n, coreRadius=r_core, densityAt150pc=core_nfw_density(0.15),
                speedAt1kpc=math.sqrt(G * nfw_mass(halo, 1.0) * math.tanh(1.0 / r_core) ** n / 1.0))

values = dict(
    description='Independent teaching values for the dark-matter-small-scales readings (Python, not the Lab engine).',
    units=dict(length='kpc', mass='Msun', speed='km/s', time='Gyr', density='Msun/kpc^3'),
    cosmology=dict(h=h, omegaMatter=omega_m, omegaCold=omega_cdm, criticalDensity=rho_crit, virialOverdensity=delta_vir),
    cold=dict(virialMass=1e10, concentration=halo['c'], virialRadius=halo['r_vir'], scaleRadius=halo['r_s'],
              scaleDensity=halo['rho_s'], densityAt150pc=nfw_density(halo, 0.15), slopeAt300pc=nfw_slope(halo, 0.3),
              massInside1kpc=nfw_mass(halo, 1.0), speedAt1kpc=v1, speedAtHalfKpc=math.sqrt(G * nfw_mass(halo, 0.5) / 0.5), maximumSpeed=v_max, maximumSpeedRadius=x_max * halo['r_s']),
    uniformCore=dict(density=uniform_core_density, speedAtHalfKpc=uniform_core_v_half_kpc),
    warm=warm, nadler=nadler, fuzzy=fuzzy, selfInteracting=sidm, feedback=feedback,
    versions=dict(python=sys.version.split()[0], numpy=np.__version__, scipy=scipy.__version__),
)
with open(output_path, 'w', encoding='utf-8') as handle:
    json.dump(values, handle, indent=1)
    handle.write('\n')
print(json.dumps(values, indent=1))
