"""Generate the frozen linear reference grid used by the AC78 tutorial and Lab.

This is the only script in the package that needs the CLASS Boltzmann code
(Python wrapper `classy`). Everything else reads the frozen JSON it writes.

Run:  python3 generate_frozen_reference_grid.py reference-grid.json

What it computes (redshift z = 0, linear theory):
  * P_CDM(k) for a flat LambdaCDM model with the Planck 2018
    TT,TE,EE+lowE+lensing+BAO central values (arXiv:1807.06209v4, Table 2,
    last column). Neutrinos are massless here, a stated simplification.
  * The amplitude ratio T(k) = sqrt(P_WDM / P_CDM) for thermal-relic warm
    dark matter of mass 2, 3 and 6.5 keV, with CLASS's non-cold dark matter
    species carrying all of the dark matter. The relic temperature is tuned
    so that omega_ncdm equals the CDM omega_c.

Units: k in h/Mpc, P in (Mpc/h)^3.
"""
import hashlib
import json
import sys

import classy
import numpy as np
from classy import Class

HUBBLE = 0.6766
OMEGA_B = 0.02242
OMEGA_C = 0.11933
BASE = {
    'output': 'mPk', 'P_k_max_h/Mpc': 320.0, 'z_pk': 0,
    'h': HUBBLE, 'omega_b': OMEGA_B, 'A_s': 2.105e-9, 'n_s': 0.9665, 'tau_reio': 0.0561,
}
WAVENUMBERS = np.logspace(np.log10(0.05), np.log10(300.0), 141)
WDM_MASSES_KEV = [2.0, 3.0, 6.5]


def linear_power(cosmology):
    return np.array([cosmology.pk_lin(k * HUBBLE, 0.0) * HUBBLE ** 3 for k in WAVENUMBERS])


def warm_transfer(mass_kev, cold_power):
    # Thermal relic: (m / 94.07 eV) (T_x / T_nu)^3 = omega_x, with T_nu / T_cmb = 0.71611.
    temperature = 0.71611 * (94.07 * OMEGA_C / (mass_kev * 1000.0)) ** (1.0 / 3.0)
    for _ in range(6):
        warm = Class()
        warm.set(dict(BASE, omega_cdm=1e-10, N_ncdm=1, m_ncdm=mass_kev * 1000.0, T_ncdm=temperature,
                      ncdm_fluid_approximation=3, l_max_ncdm=50))
        warm.compute()
        omega_ncdm = warm.Omega0_m() * HUBBLE ** 2 - OMEGA_B - 1e-10
        if abs(omega_ncdm / OMEGA_C - 1.0) < 2e-5:
            break
        temperature *= (OMEGA_C / omega_ncdm) ** (1.0 / 3.0)
        warm.struct_cleanup()
    power = linear_power(warm)
    return {'massKeV': mass_kev, 'temperatureNcdm': temperature, 'omegaNcdm': omega_ncdm,
            'transfer': [round(float(value), 7) for value in np.sqrt(power / cold_power)]}


def main(output_path):
    cold = Class()
    cold.set(dict(BASE, omega_cdm=OMEGA_C))
    cold.compute()
    cold_power = linear_power(cold)
    record = {
        'description': 'Frozen linear reference grid for AC78 (z = 0).',
        'generator': 'CLASS ' + classy.__version__ + ' through the classy Python wrapper',
        'parameters': {key: value for key, value in BASE.items() if key != 'output'},
        'omegaCdm': OMEGA_C, 'omegaMatter': cold.Omega0_m(), 'sigma8': cold.sigma8(),
        'units': {'k': 'h/Mpc', 'power': '(Mpc/h)^3'},
        'k': [round(float(value), 7) for value in WAVENUMBERS],
        'coldPower': [float('%.7g' % value) for value in cold_power],
        'warm': [warm_transfer(mass, cold_power) for mass in WDM_MASSES_KEV],
    }
    text = json.dumps(record, indent=1)
    with open(output_path, 'w', encoding='utf-8') as handle:
        handle.write(text + '\n')
    print(output_path, hashlib.sha256((text + '\n').encode()).hexdigest())


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