"""AC75 reference model: Saha, effective three-level hydrogen, visibility and
small CMB spectral distortions. Standard library only (Python >= 3.8).

Units: SI unless a name says otherwise. Redshift z, s = ln(1 + z).
Every equation used here is listed in README.md with its source.
"""
import json
import math

# CODATA 2018 exact or recommended values
C_LIGHT = 2.99792458e8          # m s^-1
H_PLANCK = 6.62607015e-34       # J s
K_BOLTZMANN = 1.380649e-23      # J K^-1
ELECTRON_MASS = 9.1093837015e-31  # kg
THOMSON = 6.6524587321e-29      # m^2
GRAVITY = 6.67430e-11           # m^3 kg^-1 s^-2
ELECTRON_VOLT = 1.602176634e-19  # J
STEFAN_BOLTZMANN = 5.670374419e-8  # W m^-2 K^-4
RADIATION_CONSTANT = 4.0 * STEFAN_BOLTZMANN / C_LIGHT  # J m^-3 K^-4
MEGAPARSEC = 3.0856775814913673e22  # m
HYDROGEN_MASS = 1.6735575e-27   # kg, hydrogen atom
# Atomic data adopted from Seager, Sasselov & Scott (1999), Section 3
IONIZATION_ENERGY = 13.605693122994 * ELECTRON_VOLT  # J, ground state
LYMAN_ALPHA_WAVELENGTH = 121.5682e-9  # m
TWO_PHOTON_RATE = 8.22458       # s^-1, 2s -> 1s
FUDGE_FACTOR = 1.14             # RECFAST fudge factor F
LYMAN_ALPHA_ENERGY = H_PLANCK * C_LIGHT / LYMAN_ALPHA_WAVELENGTH  # J (10.2 eV)
N2_IONIZATION_ENERGY = IONIZATION_ENERGY - LYMAN_ALPHA_ENERGY      # J (3.4 eV)
BETA_MU = 3.0 * 1.2020569031595942 / (math.pi ** 2 / 6.0)          # 3 zeta(3)/zeta(2)


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


class Cosmology:
    """Flat background. Radiation = photons + N_eff massless neutrinos."""

    def __init__(self, parameters):
        self.h = parameters['h']
        self.temperature_today = parameters['T0_K']
        self.omega_baryon = parameters['Omega_b']
        self.omega_matter = parameters['Omega_cb']
        self.helium_mass_fraction = parameters['Y_He']
        self.effective_neutrinos = parameters['N_eff']
        self.hubble_today = self.h * 100.0e3 / MEGAPARSEC  # s^-1
        self.critical_density = 3.0 * self.hubble_today ** 2 / (8.0 * math.pi * GRAVITY)
        photon_energy_density = RADIATION_CONSTANT * self.temperature_today ** 4
        self.omega_photon = photon_energy_density / (self.critical_density * C_LIGHT ** 2)
        self.omega_radiation = self.omega_photon * (
            1.0 + self.effective_neutrinos * 7.0 / 8.0 * (4.0 / 11.0) ** (4.0 / 3.0))
        self.omega_lambda = 1.0 - self.omega_matter - self.omega_radiation
        self.hydrogen_density_today = ((1.0 - self.helium_mass_fraction) * self.omega_baryon
                                       * self.critical_density / HYDROGEN_MASS)

    def hubble(self, z):
        a = 1.0 + z
        return self.hubble_today * math.sqrt(self.omega_radiation * a ** 4
                                             + self.omega_matter * a ** 3 + self.omega_lambda)

    def temperature(self, z):
        return self.temperature_today * (1.0 + z)

    def hydrogen_density(self, z):
        return self.hydrogen_density_today * (1.0 + z) ** 3

    def photon_energy_per_hydrogen(self, z):
        """rho_gamma / n_H in joules."""
        return RADIATION_CONSTANT * self.temperature(z) ** 4 / self.hydrogen_density(z)

    def thermalization_redshift(self):
        """z_th scaling used in CLASS v3.3.4 source/distortions.c."""
        return (1.98e6 * ((1.0 - self.helium_mass_fraction / 2.0) / 0.8767) ** (-0.4)
                * (self.omega_baryon * self.h ** 2 / 0.02225) ** (-0.4)
                * (self.temperature_today / 2.726) ** 0.2)


def thermal_de_broglie_factor(temperature):
    """(2 pi m_e k T / h^2)^(3/2) in m^-3."""
    return (2.0 * math.pi * ELECTRON_MASS * K_BOLTZMANN * temperature / H_PLANCK ** 2) ** 1.5


def saha_fraction(cosmology, z):
    """Hydrogen-only Saha equilibrium x^2/(1-x) = S."""
    temperature = cosmology.temperature(z)
    exponent = IONIZATION_ENERGY / (K_BOLTZMANN * temperature)
    if exponent > 700.0:
        return 0.0
    right = thermal_de_broglie_factor(temperature) * math.exp(-exponent) / cosmology.hydrogen_density(z)
    # Root of x^2 + S x - S = 0 written without cancellation, so 1 - x stays
    # accurate when S is large.
    return 2.0 / (1.0 + math.sqrt(1.0 + 4.0 / right))


def case_b_recombination(temperature, fudge):
    """Pequignot et al. (1991) fit as quoted by Seager et al. (1999), m^3 s^-1."""
    t = temperature / 1.0e4
    return fudge * 1.0e-19 * 4.309 * t ** (-0.6166) / (1.0 + 0.6703 * t ** 0.5300)


def gaussian_injection(s, center_z, width):
    """Normalized pulse in s = ln(1+z): integral over s equals one."""
    center = math.log1p(center_z)
    return math.exp(-0.5 * ((s - center) / width) ** 2) / (math.sqrt(2.0 * math.pi) * width)


def deposition_fractions(x):
    """Chen & Kamionkowski (2004) eq. 14, after Shull & van Steenberg (1985)."""
    x = min(max(x, 0.0), 1.0)
    return {'ionization': (1.0 - x) / 3.0, 'excitation': (1.0 - x) / 3.0, 'heat': (1.0 + 2.0 * x) / 3.0}


def peebles_rate(cosmology, z, x, injection, fudge):
    """dx/ds with s = ln(1+z). Returns (rate, C factor)."""
    temperature = cosmology.temperature(z)
    hydrogen = cosmology.hydrogen_density(z)
    hubble = cosmology.hubble(z)
    alpha = case_b_recombination(temperature, fudge)
    beta = alpha * thermal_de_broglie_factor(temperature) * math.exp(
        -N2_IONIZATION_ENERGY / (K_BOLTZMANN * temperature))
    redshifting = LYMAN_ALPHA_WAVELENGTH ** 3 / (8.0 * math.pi * hubble)
    neutral = hydrogen * max(1.0 - x, 0.0)
    peebles_c = (1.0 + redshifting * TWO_PHOTON_RATE * neutral) / (
        1.0 + redshifting * (TWO_PHOTON_RATE + beta) * neutral)
    lyman_boltzmann = math.exp(-LYMAN_ALPHA_ENERGY / (K_BOLTZMANN * temperature))
    rate = peebles_c * (x * x * hydrogen * alpha - beta * (1.0 - x) * lyman_boltzmann) / hubble
    if injection and injection['amplitude'] > 0.0:
        s = math.log1p(z)
        fractions = deposition_fractions(x)
        per_hydrogen = cosmology.photon_energy_per_hydrogen(z) * injection['amplitude'] * gaussian_injection(
            s, injection['center_z'], injection['width'])
        rate -= per_hydrogen * (fractions['ionization'] / IONIZATION_ENERGY
                                + (1.0 - peebles_c) * fractions['excitation'] / LYMAN_ALPHA_ENERGY)
    return rate, peebles_c


def solve_ionization_history(cosmology, injection=None, start_z=1800.0, end_z=50.0,
                             step=5.0e-4, fudge=FUDGE_FACTOR, method='backward-euler'):
    """Backward-Euler (default) integration from Saha at start_z down to end_z."""
    s = math.log1p(start_z)
    s_end = math.log1p(end_z)
    x = saha_fraction(cosmology, start_z)
    redshifts = [start_z]
    fractions = [x]
    while s > s_end + 1e-12:
        delta = -min(step, s - s_end)
        s_next = s + delta
        z_next = math.expm1(s_next)
        if method == 'explicit-euler':
            rate, _ = peebles_rate(cosmology, math.expm1(s), x, injection, fudge)
            guess = x + delta * rate
        else:
            guess = x
            for _ in range(60):
                rate, _ = peebles_rate(cosmology, z_next, guess, injection, fudge)
                residual = guess - x - delta * rate
                probe = max(abs(guess) * 1e-7, 1e-14)
                rate_probe, _ = peebles_rate(cosmology, z_next, guess + probe, injection, fudge)
                derivative = 1.0 - delta * (rate_probe - rate) / probe
                update = residual / derivative
                guess = min(max(guess - update, 1e-10), 1.0)
                if abs(update) < 1e-13 * max(guess, 1e-6):
                    break
        x = guess
        s = s_next
        redshifts.append(z_next)
        fractions.append(x)
    return redshifts, fractions


def optical_depth_and_visibility(cosmology, redshifts, fractions, low_step=1.0e-3):
    """tau(z) from z = 0 and g(z) = exp(-tau) dtau/dz. Below the last solved
    redshift the ionized fraction is held at its final value."""
    order = sorted(range(len(redshifts)), key=lambda index: redshifts[index])
    z_sorted = [redshifts[index] for index in order]
    x_sorted = [fractions[index] for index in order]
    tail_fraction = x_sorted[0]

    def dtau_ds(z, x):
        return C_LIGHT * THOMSON * x * cosmology.hydrogen_density(z) / cosmology.hubble(z)

    tau = 0.0
    s = 0.0
    s_low = math.log1p(z_sorted[0])
    while s < s_low - 1e-12:
        ds = min(low_step, s_low - s)
        tau += 0.5 * ds * (dtau_ds(math.expm1(s), tail_fraction) + dtau_ds(math.expm1(s + ds), tail_fraction))
        s += ds
    taus = [tau]
    for index in range(1, len(z_sorted)):
        z0, z1 = z_sorted[index - 1], z_sorted[index]
        ds = math.log1p(z1) - math.log1p(z0)
        tau += 0.5 * ds * (dtau_ds(z0, x_sorted[index - 1]) + dtau_ds(z1, x_sorted[index]))
        taus.append(tau)
    visibility = [math.exp(-taus[index]) * dtau_ds(z, x_sorted[index]) / (1.0 + z)
                  for index, z in enumerate(z_sorted)]
    return z_sorted, x_sorted, taus, visibility


def summarize_visibility(z_sorted, taus, visibility):
    peak = max(range(len(visibility)), key=lambda index: visibility[index])
    half = 0.5 * visibility[peak]
    lower = next(index for index in range(peak, -1, -1) if visibility[index] < half)
    upper = next(index for index in range(peak, len(visibility)) if visibility[index] < half)
    unity = next(index for index in range(len(taus)) if taus[index] >= 1.0)
    z_unity = z_sorted[unity - 1] + (1.0 - taus[unity - 1]) * (z_sorted[unity] - z_sorted[unity - 1]) / (
        taus[unity] - taus[unity - 1])
    area = sum(0.5 * (visibility[index] + visibility[index + 1]) * (z_sorted[index + 1] - z_sorted[index])
               for index in range(len(z_sorted) - 1))
    return {'z_peak': z_sorted[peak], 'z_tau_unity': z_unity,
            'fwhm': z_sorted[upper] - z_sorted[lower], 'area_from_low_edge': area}


def interpolate(xs, ys, x):
    if x <= xs[0]:
        return ys[0]
    if x >= xs[-1]:
        return ys[-1]
    low, high = 0, len(xs) - 1
    while high - low > 1:
        middle = (low + high) // 2
        if xs[middle] <= x:
            low = middle
        else:
            high = middle
    weight = (x - xs[low]) / (xs[high] - xs[low])
    return ys[low] * (1.0 - weight) + ys[high] * weight


# ---------- spectral distortion shapes (dimensionless, multiply by I0) ----------

def planck_shape(x):
    return x ** 3 / math.expm1(x)


def temperature_shape(x):
    return x ** 4 * math.exp(x) / math.expm1(x) ** 2


def compton_y_shape(x):
    return temperature_shape(x) * (x / math.tanh(x / 2.0) - 4.0)


def chemical_potential_shape(x):
    return temperature_shape(x) * (1.0 / BETA_MU - 1.0 / x)


def intensity_unit(temperature_today):
    """I0 = 2 (k T0)^3 / (h c)^2 in W m^-2 Hz^-1 sr^-1."""
    return 2.0 * (K_BOLTZMANN * temperature_today) ** 3 / (H_PLANCK * C_LIGHT) ** 2


def branching_soft(z, z_thermalization):
    """Chluba (2013) eqs. 5-6 with J = exp[-(z/z_th)^(5/2)]."""
    visibility = math.exp(-((z / z_thermalization) ** 2.5))
    j_y = 1.0 / (1.0 + ((1.0 + z) / 6.0e4) ** 2.58)
    j_mu = 1.0 - math.exp(-(((1.0 + z) / 5.8e4) ** 1.88))
    return {'thermalization': visibility, 'y': j_y, 'mu': j_mu}


def distortion_parameters(cosmology, injection, redshifts, fractions, upper_z=1.0e7, step=2.0e-3):
    """mu, y and Delta T / T from the heated part of a Gaussian pulse.
    Above the solved range x = 1 (heat fraction one)."""
    z_thermalization = cosmology.thermalization_redshift()
    z_sorted = sorted(redshifts)
    x_by_z = dict(zip(redshifts, fractions))
    x_sorted = [x_by_z[z] for z in z_sorted]
    s = math.log1p(upper_z)
    s_low = math.log1p(z_sorted[0])
    totals = {'mu': 0.0, 'y': 0.0, 'temperature': 0.0, 'heat': 0.0, 'injected': 0.0}

    def integrand(s_value):
        z = math.expm1(s_value)
        x = 1.0 if z > z_sorted[-1] else interpolate(z_sorted, x_sorted, z)
        pulse = injection['amplitude'] * gaussian_injection(s_value, injection['center_z'], injection['width'])
        heat = deposition_fractions(x)['heat'] * pulse
        branch = branching_soft(z, z_thermalization)
        return {'mu': 1.401 * branch['mu'] * branch['thermalization'] * heat,
                'y': 0.25 * branch['y'] * heat,
                'temperature': 0.25 * (1.0 - branch['thermalization']) * heat,
                'heat': heat, 'injected': pulse}

    previous = integrand(s)
    while s > s_low + 1e-12:
        ds = min(step, s - s_low)
        current = integrand(s - ds)
        for key in totals:
            totals[key] += 0.5 * ds * (previous[key] + current[key])
        previous = current
        s -= ds
    totals['z_thermalization'] = z_thermalization
    return totals