"""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
from petitRADTRANS.physics import guillot_global

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',
    'co_added_cross_correlation':
        'co_added_cross_correlation_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_transmission_2d':
        'radtrans_line_by_line_transmission_2d_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': {
                    'name': 'PSR B1257+12 b',  # actual planet has no link to below 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)
                    'mass': 1.0,  # (M_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': {
                    'resolving_power': 60,
                    'high_resolution_resolving_power': 1e5,
                    'wavelength_range_high_resolution': (2.3005, 2.3020),
                    'orbital_phase_range': (-1e-2, 1e-2),
                    'system_observer_radial_velocities_range': (-5.0e5, -4.85e5),
                    'n_exposures': 20,
                    'planet_rest_frame_velocity_shift': -1.5e5,
                    'relative_error': 1e-2
                },
                'ccf_analysis_parameters': {
                    'normalize_ccf': True,
                    'calculate_ccf_snr': True,
                    'ccf_sum_axes': (0,),
                    'pixels_per_resolution_element': 2,
                    'peak_lsf_factor': 1.5,
                    'velocity_interval_extension_factor': -0.3,
                    'kp_factor': 1.5
                }
            },
            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 = 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()
