"""Compare Saha, the effective three-level atom and the frozen HyRec-2 extract.

Usage:  python3 -I run_recombination.py inputs.json hyrec-reference-extract.json my-output.json
"""
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

SAMPLE_REDSHIFTS = [1700, 1600, 1500, 1400, 1300, 1200, 1100, 1000, 900, 800, 600, 400, 200, 100]
PLANCK_2018_Z_STAR = {'value': 1089.92, 'sigma': 0.25,
                      'source': 'Planck 2018 VI, Table 2, TT,TE,EE+lowE+lensing column'}


def history(cosmology, **options):
    redshifts, fractions = model.solve_ionization_history(cosmology, **options)
    z_sorted, x_sorted, taus, visibility = model.optical_depth_and_visibility(cosmology, redshifts, fractions)
    return {'z': z_sorted, 'x': x_sorted, 'tau': taus, 'g': visibility,
            'summary': model.summarize_visibility(z_sorted, taus, visibility),
            'raw': (redshifts, fractions)}


def sample(run, z):
    return model.interpolate(run['z'], run['x'], z)


def main():
    inputs_path, extract_path, output_path = sys.argv[1:4]
    inputs = model.load_inputs(inputs_path)
    with open(extract_path, encoding='utf-8') as handle:
        extract = json.load(handle)
    cosmology = model.Cosmology(inputs['cosmology'])
    solver = inputs['solver']
    options = dict(start_z=solver['start_z'], end_z=solver['end_z'], step=solver['step_in_ln_1_plus_z'])
    baseline = history(cosmology, fudge=solver['fudge_factor'], **options)
    no_fudge = history(cosmology, fudge=1.0, **options)
    half_step = history(cosmology, fudge=solver['fudge_factor'], **dict(options, step=options['step'] / 2.0))
    hyrec_z = list(reversed(extract['z']))
    hyrec_x = list(reversed(extract['x_e']))
    rows = []
    for z in SAMPLE_REDSHIFTS:
        reference = model.interpolate(hyrec_z, hyrec_x, z)
        peebles = sample(baseline, z)
        rows.append({'z': z, 'saha': model.saha_fraction(cosmology, z), 'three_level_F1_14': peebles,
                     'three_level_F1_00': sample(no_fudge, z), 'hyrec_extract': reference,
                     'three_level_minus_hyrec_over_hyrec': (peebles - reference) / reference})
    # The extract is also integrated for its own visibility, below z = 1600 where
    # helium is neutral, so the hydrogen-only optical depth is comparable.
    log_reference = [math.log(x) for x in hyrec_x]
    reference_run_z = [1600.0 - step for step in range(0, 1551)]
    reference_run_x = [math.exp(model.interpolate(hyrec_z, log_reference, z)) for z in reference_run_z]
    reference_tau = model.optical_depth_and_visibility(cosmology, reference_run_z, reference_run_x)
    failures = {}
    try:
        explicit = model.solve_ionization_history(cosmology, method='explicit-euler', start_z=1800.0,
                                                  end_z=1500.0, step=5.0e-3)
        final = explicit[1][-1]
        failures['explicit_euler_step_0_005'] = {'x_at_1500': final if math.isfinite(final) else repr(final),
                                                 'note': 'Stiff equation: the explicit step overshoots Saha.'}
    except (OverflowError, ValueError) as error:
        failures['explicit_euler_step_0_005'] = {'error': repr(error)}
    late = history(cosmology, start_z=1300.0, end_z=50.0, step=options['step'])
    failures['start_from_saha_at_1300'] = {'x_at_1100': sample(late, 1100), 'baseline_x_at_1100': sample(baseline, 1100),
                                           'note': 'Saha is already far from the true history at z = 1300.'}
    cases = []
    for case in inputs['injection_cases']:
        injection = dict(case, width=inputs['injection_width_in_ln_1_plus_z'])
        run = history(cosmology, injection=injection, fudge=solver['fudge_factor'], **options)
        distortion = model.distortion_parameters(cosmology, injection, *run['raw'])
        cases.append({'id': case['id'], 'center_z': case['center_z'], 'amplitude': case['amplitude'],
                      'visibility': run['summary'], 'x_at_500': sample(run, 500), 'x_at_200': sample(run, 200),
                      'mu': distortion['mu'], 'y': distortion['y'], 'delta_T_over_T': distortion['temperature'],
                      'heat_fraction_of_injection': distortion['heat'] / distortion['injected']})
    output = {
        'format': 'sandboxphysics-ac75-recombination-output-v1',
        'python': sys.version.split()[0],
        'derived_background': {'Omega_radiation': cosmology.omega_radiation, 'Omega_lambda': cosmology.omega_lambda,
                               'n_H_today_m^-3': cosmology.hydrogen_density_today,
                               'z_thermalization': cosmology.thermalization_redshift()},
        'samples': rows,
        'visibility': {'three_level_F1_14': baseline['summary'], 'three_level_F1_00': no_fudge['summary'],
                       'hyrec_extract_below_1600': model.summarize_visibility(*[reference_tau[i] for i in (0, 2, 3)]),
                       'planck_2018_reference': PLANCK_2018_Z_STAR},
        'step_convergence': {'z_peak_change': half_step['summary']['z_peak'] - baseline['summary']['z_peak'],
                             'z_tau_unity_change': half_step['summary']['z_tau_unity'] - baseline['summary']['z_tau_unity'],
                             'x_at_1100_relative_change': (sample(half_step, 1100) - sample(baseline, 1100)) / sample(baseline, 1100)},
        'failure_cases': failures,
        'injection_cases': cases,
    }
    with open(output_path, 'w', encoding='utf-8') as handle:
        json.dump(output, handle, indent=1)
        handle.write('\n')
    print('wrote', output_path)
    print('tau = 1 at z = %.2f (Planck 2018: %.2f +/- %.2f)' % (
        baseline['summary']['z_tau_unity'], PLANCK_2018_Z_STAR['value'], PLANCK_2018_Z_STAR['sigma']))
    worst = max(abs(row['three_level_minus_hyrec_over_hyrec']) for row in rows if row['z'] <= 1500)
    print('largest |three-level - HyRec| / HyRec for z <= 1500: %.3f' % worst)


if __name__ == '__main__':
    main()