"""Compare the Lab's simplified model with frozen reference files shipped in
CLASS v3.3.4 (commit e85808324f51fc694d12e3ed7439552a3c3f9540).

Sandbox Physics does not run CLASS. This script reads four data files from a
CLASS checkout (or from the raw URLs listed in README.md), checks their SHA256,
and writes a JSON comparison:
  1. the full HyRec-2 ionization history against Saha and the three-level atom;
  2. the energy carried by each CosmoTherm Green's function;
  3. a re-run of CLASS's Gram-Schmidt split of those Green's functions on the
     COBE/FIRAS frequency grid, against FIRAS_branching_ratios.dat;
  4. the soft branching approximation used by the Lab against that table.

Usage:  python3 -I compare_class_files.py path/to/class_public my-class-comparison.json
"""
import hashlib
import json
import math
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ac75_model as model  # noqa: E402

FILES = {
    'hyrec': ('external/HyRec2020/output_xe.dat', '1ac2ddad6aa0e3c4080883c71479e2e398669468e4813075ec61bca86b7afc40'),
    'greens': ('external/distortions/Greens_data.dat', '03972c42a530315a768c973a45c343ee744a8d64da93493c9c68057f92327bed'),
    'branching': ('external/distortions/FIRAS_branching_ratios.dat', '966daa8fa99d7e15de6961f33e0fac8800ae5d64464c21d3526c52cd8fa42864'),
    'firas_noise': ('external/distortions/FIRAS_nu_delta_I.dat', '92f583021393c543b512c038a9a829b766027ab18938a2b527a1d8fdac889aa7'),
}
INPUTS = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'inputs.json')


def read_checked(root, key):
    relative, expected = FILES[key]
    raw = open(os.path.join(root, relative), 'rb').read()
    digest = hashlib.sha256(raw).hexdigest()
    if digest != expected:
        raise SystemExit('%s changed: %s' % (relative, digest))
    return raw.decode('ascii')


def numeric_lines(text):
    return [line for line in text.splitlines() if line.strip() and not line.lstrip().startswith('#')]


def compare_hyrec(text, cosmology, solver):
    rows = sorted((tuple(float(v) for v in line.split()) for line in numeric_lines(text)), key=lambda row: row[0])
    reference_z = [row[0] for row in rows]
    reference_x = [row[1] for row in rows]
    redshifts, fractions = model.solve_ionization_history(
        cosmology, start_z=solver['start_z'], end_z=solver['end_z'], step=solver['step_in_ln_1_plus_z'],
        fudge=solver['fudge_factor'])
    z_sorted = list(reversed(redshifts))
    x_sorted = list(reversed(fractions))
    worst = {'z': None, 'relative': 0.0}
    table = []
    for z in range(1600, 49, -1):
        reference = model.interpolate(reference_z, reference_x, z)
        relative = (model.interpolate(z_sorted, x_sorted, z) - reference) / reference
        if abs(relative) > abs(worst['relative']):
            worst = {'z': z, 'relative': relative}
        if z % 100 == 0:
            table.append({'z': z, 'hyrec': reference, 'three_level': model.interpolate(z_sorted, x_sorted, z),
                          'saha': model.saha_fraction(cosmology, z), 'relative': relative})
    helium_plateau = model.interpolate(reference_z, reference_x, 8000.0)
    helium_fraction = cosmology.helium_mass_fraction / (4.0026 / 1.00794 * (1.0 - cosmology.helium_mass_fraction))
    return {'rows': len(rows), 'z_range': [reference_z[0], reference_z[-1]], 'every_100': table,
            'largest_relative_difference_z_50_to_1600': worst,
            'x_e_at_8000': helium_plateau, 'one_plus_two_f_He': 1.0 + 2.0 * helium_fraction}


def parse_greens(text):
    lines = numeric_lines(text)
    count_z, count_x = (int(v) for v in lines[0].split())
    redshifts = [float(v) for v in lines[1].split()]
    initial_t = [float(v) for v in lines[2].split()]
    last_t = [float(v) for v in lines[3].split()]
    injected = [float(v) for v in lines[4].split()]
    rows = sorted((tuple(float(v) for v in line.split()) for line in lines[5:5 + count_x]), key=lambda row: row[0])
    assert len(redshifts) == count_z and len(rows) == count_x
    return redshifts, initial_t, last_t, injected, rows


def trapezoid(xs, ys):
    return sum(0.5 * (ys[i] + ys[i + 1]) * (xs[i + 1] - xs[i]) for i in range(len(xs) - 1))


def greens_energy(greens):
    redshifts, _, _, _, rows = greens
    xs = [row[0] for row in rows]
    blackbody = [row[-1] for row in rows]
    energy_b = trapezoid(xs, blackbody)
    picks = [0, 20, 40, 60, 70, 80, 90, 100, 110, 117]
    return [{'z': redshifts[k], 'energy_over_blackbody_energy': trapezoid(xs, [row[1 + k] for row in rows]) / energy_b}
            for k in picks]


def dot(a, b):
    return sum(p * q for p, q in zip(a, b))


def project_like_class(greens, noise_text, cosmology):
    """Same steps as external/distortions/generate_PCA_files.py, evaluated only at
    the Green's file redshifts (no spline in z) and with linear x interpolation."""
    redshifts, initial_t, last_t, injected, rows = greens
    xs = [row[0] for row in rows]
    lines = numeric_lines(noise_text)
    frequencies = [float(line.split()[0]) for line in lines[1:]]
    x_to_nu = model.K_BOLTZMANN * cosmology.temperature_today / model.H_PLANCK / 1.0e9
    grid = [nu / x_to_nu for nu in frequencies]
    unit = rows[len(rows) // 2][-1] / model.planck_shape(xs[len(rows) // 2])
    shapes = {'G': [model.temperature_shape(x) for x in grid], 'Y': [model.compton_y_shape(x) for x in grid],
              'M': [model.chemical_potential_shape(x) for x in grid]}
    norm = lambda v: math.sqrt(dot(v, v))  # noqa: E731
    e_y = [v / norm(shapes['Y']) for v in shapes['Y']]
    m_y, g_y = dot(e_y, shapes['M']), dot(e_y, shapes['G'])
    m_perp = [m - m_y * e for m, e in zip(shapes['M'], e_y)]
    e_m = [v / norm(m_perp) for v in m_perp]
    g_m = dot(e_m, shapes['G'])
    g_perp = [g - g_y * a - g_m * b for g, a, b in zip(shapes['G'], e_y, e_m)]
    e_g = [v / norm(g_perp) for v in g_perp]
    z_th = cosmology.thermalization_redshift()
    result = []
    for k, z in enumerate(redshifts):
        visibility = math.exp(-((z / z_th) ** 2.5))
        column = [row[1 + k] / unit for row in rows]
        response = []
        for x in grid:
            value = model.interpolate(xs, column, x) * visibility
            shifted = initial_t[k] / last_t[k] * x
            value += x ** 3 * (1.0 / math.expm1(shifted) - 1.0 / math.expm1(x)) / injected[k]
            response.append(value)
        f_g = dot(response, e_g) / norm(g_perp)
        f_mu = (dot(response, e_m) - g_m * f_g) / norm(m_perp)
        f_y = (dot(response, e_y) - m_y * f_mu - g_y * f_g) / norm(shapes['Y'])
        result.append({'z': z, 'f_T': f_g, 'f_y': f_y, 'f_mu': f_mu})
    return result, {'frequencies_GHz': [frequencies[0], frequencies[-1]], 'count': len(frequencies), 'x_to_nu_GHz': x_to_nu}


def compare_branching(text, projection, cosmology):
    rows = [tuple(float(v) for v in line.split()) for line in numeric_lines(text)[1:]]
    table_z = [row[0] for row in rows]
    z_th = cosmology.thermalization_redshift()
    picks = [1.1e3, 5e3, 1e4, 2e4, 5e4, 1e5, 2e5, 5e5, 1e6, 2e6, 3e6]
    comparison = []
    for z in picks:
        soft = model.branching_soft(z, z_th)
        index = min(range(len(projection)), key=lambda i: abs(math.log(projection[i]['z'] / z)))
        own = projection[index]
        comparison.append({
            'z': z,
            'table': {name: model.interpolate(table_z, [row[column] for row in rows], z)
                      for name, column in (('f_T', 1), ('f_y', 2), ('f_mu', 3))},
            'soft_used_by_lab': {'f_T': 0.25 * (1.0 - soft['thermalization']), 'f_y': 0.25 * soft['y'],
                                 'f_mu': 1.401 * soft['mu'] * soft['thermalization']},
            'own_projection_at_nearest_greens_z': own,
        })
    return {'rows': len(rows), 'z_range': [table_z[0], table_z[-1]], 'comparison': comparison}


def main():
    root, output_path = sys.argv[1:3]
    inputs = model.load_inputs(INPUTS)
    cosmology = model.Cosmology(inputs['cosmology'])
    greens = parse_greens(read_checked(root, 'greens'))
    projection, grid = project_like_class(greens, read_checked(root, 'firas_noise'), cosmology)
    output = {
        'format': 'sandboxphysics-ac75-class-comparison-v1',
        'python': sys.version.split()[0],
        'class_version': 'v3.3.4, commit e85808324f51fc694d12e3ed7439552a3c3f9540 (files read, CLASS not executed)',
        'sha256': {key: value[1] for key, value in FILES.items()},
        'hyrec': compare_hyrec(read_checked(root, 'hyrec'), cosmology, inputs['solver']),
        'greens_energy': greens_energy(greens),
        'firas_grid': grid,
        'branching': compare_branching(read_checked(root, 'branching'), projection, cosmology),
    }
    with open(output_path, 'w', encoding='utf-8') as handle:
        json.dump(output, handle, indent=1)
        handle.write('\n')
    print('wrote', output_path)
    print('HyRec largest relative difference (z 50-1600):', output['hyrec']['largest_relative_difference_z_50_to_1600'])
    for row in output['branching']['comparison']:
        print('z=%.1e table T/y/mu %.3f %.3f %.3f | soft %.3f %.3f %.3f | own %.3f %.3f %.3f' % (
            row['z'], row['table']['f_T'], row['table']['f_y'], row['table']['f_mu'],
            row['soft_used_by_lab']['f_T'], row['soft_used_by_lab']['f_y'], row['soft_used_by_lab']['f_mu'],
            row['own_projection_at_nearest_greens_z']['f_T'], row['own_projection_at_nearest_greens_z']['f_y'],
            row['own_projection_at_nearest_greens_z']['f_mu']))


if __name__ == '__main__':
    main()