"""Utility functions for tests.

Regenerate the comparison files only when changing the precision of the models.
"""
import json
import os

import numpy as np

from .context import petitRADTRANS

version = "2.4.5"  # petitRADTRANS.version.version used to generate last tests

tests_data_directory = os.path.join(os.path.dirname(__file__), 'data')
tests_error_directory = os.path.join(os.path.dirname(__file__), 'errors')
tests_results_directory = os.path.join(os.path.dirname(__file__), 'results')
reference_filenames = {
    'config_test_radtrans':
        'config_test_radtrans',
    'guillot_2010':
        'radtrans_guillot_2010_temperature_profile_ref',
    'correlated_k_transmission':
        'radtrans_correlated_k_transmission_ref',
    'correlated_k_transmission_cloud_power_law':
        'radtrans_correlated_k_transmission_cloud_power_law_ref',
    'correlated_k_transmission_gray_cloud':
        'radtrans_correlated_k_transmission_gray_cloud_ref',
    'correlated_k_transmission_rayleigh':
        'radtrans_correlated_k_transmission_rayleigh_ref',
    'correlated_k_transmission_cloud_fixed_radius':
        'radtrans_correlated_k_transmission_cloud_fixed_radius_ref',
    'correlated_k_transmission_cloud_calculated_radius':
        'radtrans_correlated_k_transmission_cloud_calculated_radius_ref',
    'correlated_k_transmission_cloud_calculated_radius_scattering':
        'radtrans_correlated_k_transmission_cloud_calculated_radius_scattering_ref',
    'correlated_k_transmission_contribution_cloud_calculated_radius':
        'radtrans_correlated_k_transmission_contribution_cloud_calculated_radius_ref',
    'correlated_k_emission':
        'radtrans_correlated_k_emission_ref',
    'correlated_k_emission_cloud_calculated_radius':
        'radtrans_correlated_k_emission_cloud_calculated_radius_ref',
    'correlated_k_emission_cloud_calculated_radius_scattering':
        'radtrans_correlated_k_emission_cloud_calculated_radius_scattering_ref',
    'correlated_k_emission_cloud_calculated_radius_scattering_planetary_ave':
        'radtrans_correlated_k_emission_cloud_calculated_radius_scattering_average_ref',
    'correlated_k_emission_cloud_calculated_radius_scattering_dayside_ave':
        'radtrans_correlated_k_emission_cloud_calculated_radius_scattering_dayside_ref',
    'correlated_k_emission_cloud_calculated_radius_scattering_non-isotropic':
        'radtrans_correlated_k_emission_cloud_calculated_radius_scattering_non-isotropic_ref',
    'correlated_k_emission_contribution_cloud_calculated_radius':
        'radtrans_correlated_k_emission_contribution_cloud_calculated_radius_ref',
    'correlated_k_emission_cloud_hansen_radius':
        'radtrans_correlated_k_emission_cloud_hansen_radius_ref',
    'correlated_k_emission_surface_scattering':
        'radtrans_correlated_k_emission_surface_scattering_ref',
    'line_by_line_downsampled_transmission':
        'radtrans_line_by_line_downsampled_transmission_ref',
    'line_by_line_downsampled_emission':
        'radtrans_line_by_line_downsampled_emission_ref',
    'line_by_line_transmission':
        'radtrans_line_by_line_transmission_ref',
    'line_by_line_emission':
        'radtrans_line_by_line_emission_ref',
    'mass_fractions_atmosphere':
        'mass_fractions_atmosphere_ref',
    'mass_fractions_atmosphere_quench':
        'mass_fractions_atmosphere_quench_ref',
    'mass_fractions_c_o_ratios':
        'mass_fractions_c_o_ratios_ref',
    'mass_fractions_metallicities':
        'mass_fractions_metallicities_ref'
}

# Complete filenames
reference_filenames = {
    key: os.path.join(tests_data_directory, value + '.npz')
    if key != 'config_test_radtrans' else
    os.path.join(tests_data_directory, value + '.json')
    for key, value in reference_filenames.items()
}

reference_filenames['mock_observation_transmission'] = reference_filenames[
        'correlated_k_transmission_cloud_calculated_radius_scattering'
    ].rsplit('.', 1)[0] + '.dat'
reference_filenames['pymultinest_parameter_analysis'] = \
    os.path.join(tests_data_directory, 'test_stats' + '.json')


# Make directories if needed
if not os.path.isdir(tests_results_directory):
    os.mkdir(tests_results_directory)


# Common parameters
def create_test_radtrans_config_file(filename):
    with open(os.path.join(filename), 'w') as f:
        json.dump(
            obj={
                'header': f'File generated by tests.utils function\n'
                          f'wavelength units: um\n'
                          f'pressure units: log10(bar), generate using numpy.logspace\n'
                          f'planet radius units: R_jup\n'
                          f'star radius units: R_sun\n'
                          f'metallicity units: log10(solar metallicity)\n'
                          f'angle units: deg\n'
                          f'other units: cgs',
                'prt_version': f'{version}',
                'pressures': {
                    'start': -6,
                    'stop': 2,
                    'stop_thin_atmosphere': 0,
                    'num': 27
                },
                'mass_fractions': {
                    'H2': 0.74,
                    'He': 0.24,
                    'H2O_HITEMP': 0.001,
                    'H2O_main_iso': 0.001,
                    'CH4': 0.001,
                    'CO_all_iso': 0.1,
                    'Mg2SiO4(c)': 0.0
                 },
                'mean_molar_mass': 2.33,  # (g.cm-3)
                'chemical_parameters': {
                    'metallicities': [-1.9, 0.0, 3.0],  # (log_10 solar)
                    'c_o_ratios': [0.1, 0.55, 1.6],
                    'pressure': 0.1,  # (bar)
                    'temperature': 900.0,  # (K)
                    'pressure_quench_carbon': 1.0  # (bar)
                },
                'temperature_isothermal': 1200.0,  # (K)
                'temperature_guillot_2010_parameters': {
                    'infrared_mean_opacity': 0.01,
                    'gamma': 0.4,
                    'intrinsic_temperature': 200.0,  # (K)
                    'equilibrium_temperature': 1500.0  # (K)
                },
                'planetary_parameters': {
                    'reference_pressure': 0.01,  # (bar)
                    'radius': 1.838,  # (R_jup)
                    'surface_gravity': 1e1 ** 2.45,  # (cm.s-2)
                    'eddy_diffusion_coefficient': 10 ** 7.5,
                    'orbit_semi_major_axis': 7.5e11,  # (cm)
                    'surface_reflectance': 0.3
                },
                'stellar_parameters': {
                    'effective_temperature': 5778.0,  # (K)
                    'radius': 1.0,  # (R_sun)
                    'incidence_angle': 30  # (deg)
                },
                'spectrum_parameters': {
                    'line_by_line_opacity_sampling': 4,
                    'line_species_correlated_k': [
                        'H2O_HITEMP',
                        'CH4'
                    ],
                    'line_species_line_by_line': [
                        'H2O_main_iso',
                        'CO_all_iso'
                    ],
                    'rayleigh_species': ['H2', 'He'],
                    'continuum_opacities': ['H2-H2', 'H2-He'],
                    'wavelength_range_correlated_k': [0.9, 1.2],
                    'wavelength_range_line_by_line': [2.3000, 2.3025]
                },
                'cloud_parameters': {
                   'kappa_zero': 0.01,
                   'gamma_scattering': -4.0,
                   'cloud_pressure': 0.01,
                   'haze_factor': 10.0,
                   'cloud_species': {
                       'Mg2SiO4(c)_cd': {
                           'mass_fraction': 5e-7,
                           'radius': 5e-5,  # (cm)
                           'f_sed': 2.0,
                           'sigma_log_normal': 1.05,
                           'b_hansen': 0.01
                       },
                   }
                },
                'retrieval_parameters': {
                    'planetary_radius_bounds': (1.8, 2.0),
                    'intrinsic_temperature_bounds': (500, 1500),
                    'log10_cloud_pressure_bounds': (-6, 2),
                    'log10_species_mass_fractions_bounds': (-6, 0),
                    'sample_spectrum_output': False,
                    'ultranest': False,
                    'sampling_efficiency': 0.8,
                    'n_live_points': 50,
                    'const_efficiency_mode': False,
                    'resume': False
                },
                'mock_observation_parameters': {
                    'resolution_power': 60,
                    'relative_error': 1e-2
                }
            },
            fp=f,
            indent=4
        )


def init_radtrans_parameters(recreate_parameter_file=False):
    """
    Initialize various parameters used both to perform the tests and generate the reference files.
    Do not change these parameters when comparing with a previous version.
    """
    if not os.path.isfile(reference_filenames['config_test_radtrans']) or recreate_parameter_file:
        print('Generating Radtrans test parameters file...')
        create_test_radtrans_config_file(filename=reference_filenames['config_test_radtrans'])

    with open(reference_filenames['config_test_radtrans'], 'r') as f:
        parameters = json.load(f)

    parameters['pressures_thin_atmosphere'] = np.logspace(
        parameters['pressures']['start'],
        parameters['pressures']['stop_thin_atmosphere'],
        parameters['pressures']['num']
    )

    parameters['pressures'] = np.logspace(
        parameters['pressures']['start'],
        parameters['pressures']['stop'],
        parameters['pressures']['num']
    )

    for key in parameters['mass_fractions']:
        parameters['mass_fractions'][key] *= np.ones_like(parameters['pressures'])

    parameters['mean_molar_mass'] *= np.ones_like(parameters['pressures'])
    parameters['planetary_parameters']['eddy_diffusion_coefficient'] *= np.ones_like(parameters['pressures'])

    return parameters


radtrans_parameters = init_radtrans_parameters()


# Useful functions
def check_cloud_mass_fractions():
    """
    Check if cloud mass fraction is set to 0 by default.
    This is necessary to correctly assess the effect of the different clear and cloud models.
    """
    for species, mmr in radtrans_parameters['mass_fractions'].items():
        if '(c)' in species or '(l)' in species or '(s)' in species or '(cr)' in species:  # condensed species
            if not np.all(mmr == 0):
                raise ValueError(
                    f"cloud {species} has a default mass fraction different of 0, cannot perform test\n"
                    f"mass fraction was: {mmr}"
                )


def compare_from_reference_file(reference_file, comparison_dict, relative_tolerance, absolute_tolerance=0):
    reference_data = np.load(reference_file)
    print(f"Comparing generated spectrum to result from petitRADTRANS-{reference_data['prt_version']}...")

    for reference_file_key in comparison_dict:
        try:
            assert np.allclose(
                comparison_dict[reference_file_key],
                reference_data[reference_file_key],
                rtol=relative_tolerance,
                atol=absolute_tolerance
            )
        except AssertionError:
            # Save data for diagnostic
            if not os.path.isdir(tests_error_directory):
                os.mkdir(tests_error_directory)

            error_file = os.path.join(
                tests_error_directory,
                f"{os.path.basename(reference_file).rsplit('.', 1)[0]}_error_{reference_file_key}"
            )
            print(f"Saving assertion error data in file '{error_file}' for diagnostic...")

            np.savez_compressed(
                error_file,
                test_result=comparison_dict[reference_file_key],
                data=reference_data[reference_file_key],
                relative_tolerance=relative_tolerance,
                absolute_tolerance=absolute_tolerance
            )

            # Raise the AssertionError
            raise


# Initializations
def init_guillot_2010_temperature_profile():
    temperature_guillot = petitRADTRANS.nat_cst.guillot_global(
        P=radtrans_parameters['pressures'],
        kappa_IR=radtrans_parameters['temperature_guillot_2010_parameters']['infrared_mean_opacity'],
        gamma=radtrans_parameters['temperature_guillot_2010_parameters']['gamma'],
        grav=radtrans_parameters['planetary_parameters']['surface_gravity'],
        T_int=radtrans_parameters['temperature_guillot_2010_parameters']['intrinsic_temperature'],
        T_equ=radtrans_parameters['temperature_guillot_2010_parameters']['equilibrium_temperature']
    )

    return temperature_guillot


def init_radtrans_test():
    check_cloud_mass_fractions()

    tp_iso = radtrans_parameters['temperature_isothermal'] * np.ones_like(radtrans_parameters['pressures'])
    tp_guillot_2010 = init_guillot_2010_temperature_profile()

    return tp_iso, tp_guillot_2010


temperature_isothermal, temperature_guillot_2010 = init_radtrans_test()
