#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""

Summary
-------

Test Spectrum rescaling methods

-------------------------------------------------------------------------------


"""

from __future__ import print_function, absolute_import, division, unicode_literals
import radis
from radis.spectrum.rescale import get_redundant, get_recompute
from radis.tools.database import load_spec
from radis.test.utils import getTestFile
from radis.misc.printer import printm
import numpy as np
import pytest

@pytest.mark.fast
def test_compression(verbose=True, warnings=True, *args, **kwargs):
    ''' Test that redundant quantities are properly infered from already known 
    spectral quantities '''

    # Get spectrum
    s1 = load_spec(getTestFile('CO_Tgas1500K_mole_fraction0.01.spec'), binary=True)
    
    s1.conditions['thermal_equilibrium'] = True
    s1.update()

    # Analyse redundant spectral quantities
    redundant = get_redundant(s1)
    if verbose:
        print(redundant)

    assert redundant == {'emissivity_noslit': True, 'radiance_noslit': True,
                         'emisscoeff': True, 'transmittance_noslit': True,
                         'absorbance': True, 'abscoeff': False}

    return True


@pytest.mark.fast
def test_update_transmittance(verbose=True, warnings=True, *args, **kwargs):
    ''' Test that update can correctly recompute missing quantities '''
    # TODO: add one with radiance too

    # Work with a Spectrum object that was generated by Specair
    s = load_spec(getTestFile('N2C_specair_380nm.spec'))
    w1, T1 = s.get('transmittance_noslit')   # our reference

    if verbose:
        debug_mode = radis.DEBUG_MODE  # shows all the rescale steps taken
        radis.DEBUG_MODE = True

    # Now let's apply some update() steps

    # 1) Make sure updating doesnt change anything
    s.update()
    w2, T2 = s.get('transmittance_noslit')

    # 2) Now recompute from abscoeff
    del s._q['transmittance_noslit']
    s.update()
    w2, T3 = s.get('transmittance_noslit')

    # 3) Now recompute from absorbance
    del s._q['transmittance_noslit']
    del s._q['abscoeff']
    s.update()
    w2, T4 = s.get('transmittance_noslit')

    if verbose:
        radis.DEBUG_MODE = debug_mode

    # Compare
    assert np.allclose(T1, T2)
    assert np.allclose(T1, T3)
    assert np.allclose(T1, T4)

    return True


def test_get_recompute(verbose=True, *args, **kwargs):
    ''' Make sure :func:`~radis.spectrum.rescale.get_recompute` works as expected

    Here, we check which quantities are needed to recompute radiance_noslit'''

    # Equilibrium
    # -----------
    s = load_spec(getTestFile('CO_Tgas1500K_mole_fraction0.01.spec'), binary=True)
    
    assert s.get_vars() == ['abscoeff']
    assert s.conditions['thermal_equilibrium']
    # At equilibrium, everything should be deduced from abscoeff
    assert set(get_recompute(s, ['radiance_noslit'])) == set(
        ('radiance_noslit', 'abscoeff'))

    # Non Equilibrium
    # ----------------
    s.conditions['Tvib'] = 2000  
    # a problem should be detected by is_at_equilibrium()
    with pytest.raises(AssertionError):
        assert not s.is_at_equilibrium(check='error')
    # force non equilibrium
    s.conditions['thermal_equilibrium'] = False
    
    # Now more data is needed:
    assert set(get_recompute(s, ['radiance_noslit'])) == set(
        ('abscoeff', 'emisscoeff', 'radiance_noslit'))


def test_recompute_equilibrium(verbose=True, warnings=True, plot=True,
                               *args, **kwargs):
    ''' Test that spectral quantities recomputed under equilibrium assumption
    yields the same output as with non equilibrium routines when Tvib = Trot '''

    if plot:
        import matplotlib.pyplot as plt
        plt.ion()   # dont get stuck with Matplotlib if executing through pytest

    # Get spectrum
    s1 = load_spec(getTestFile('CO_Tgas1500K_mole_fraction0.01.spec'))
    s1.rescale_path_length(100)  # just for fun

    assert s1.is_at_equilibrium()
    s1.update('emisscoeff')

    # force non equilibrium calculation
    s2 = s1.copy()
    s2.conditions['thermal_equilibrium'] = False
    s2.update()
    assert 'emissivity_noslit' not in s2.get_vars()  # just a check update() was done at nonequilibrium

    # update s1 now (at equilibrium)
    s1.update()
    assert 'emissivity_noslit' in s1.get_vars()   # just a check update() was done at equilibrium

    s1.name = 'scaled with Kirchoff law'
    s2.name = 'scaled from emisscoeff + abscoeff with RTE'

    if verbose:
        print('Checked that scaling at equilibrium with Kirchoff law yields the ' +
              'same radiance as by solving the RTE from emisscoeff and abscoeff')

    # Now Compare both update processes
    assert s1.compare_with(s2, spectra_only='radiance_noslit', plot=plot)



def test_rescale_all_quantities(verbose=True, warnings=True, *args, **kwargs):

    new_mole_fraction = 0.5
    new_path_length = 0.1
    
    # Get spectrum
    s0 = load_spec(getTestFile('CO_Tgas1500K_mole_fraction0.01.spec'), binary=True)
    s0.update('all')     # start with all possible quantities in s0
    sscaled = s0.copy()
    sscaled.rescale_mole_fraction(new_mole_fraction)
    sscaled.rescale_path_length(new_path_length)
    s0.conditions['thermal_equilibrium'] = False   # to prevent rescaling with Kirchoff
    # remove emissivity_no_slit (speciifc to equilibrium)
    del s0._q['emissivity_noslit']
     
    
    # Determine all quantities that can be recomputed
    if verbose >= 2:
        import radis
        DEBUG_MODE = radis.DEBUG_MODE
        radis.DEBUG_MODE = True
    from radis.spectrum.rescale import get_reachable, ordered_keys, _build_update_graph
    # ordered_keys: all spectral quantities that can be rescaled
    can_be_recomputed = get_reachable(s0)
    # can_be_recomputed: all spectra quantities that can be rescaled for this 
    # particular spectrum
    update_paths = _build_update_graph(s0)
    # update_paths: which quantities are needed to recompute the others
    
    rescale_list = [k for k in ordered_keys if can_be_recomputed[k]]
    
    for quantity in rescale_list:
        all_paths = update_paths[quantity]
        if verbose:
            printm('{0} can be recomputed from {1}'.format(quantity, ' or '.join(
                    ['&'.join(combinations) for combinations in all_paths])))
        
        # Now let's test all paths
        for combinations in all_paths:
            if verbose:
                printm('> computing {0} from {1}'.format(quantity, '&'.join(combinations)))
            s = s0.copy()
            # Delete all other quantities
            for k in s.get_vars():
                if k not in combinations:
                    del s._q[k]
            
            s.update(quantity, verbose=verbose)
            
            # Now rescale
            s.rescale_mole_fraction(new_mole_fraction)
            s.rescale_path_length(new_path_length)
            
            # Compare
            assert s.compare_with(sscaled, spectra_only=quantity, plot=False)
            
    if verbose >= 2:
        radis.DEBUG_MODE = DEBUG_MODE
    
def _run_all_tests(verbose=True, warnings=True, *args, **kwargs):
    test_compression(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_update_transmittance(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_get_recompute(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_recompute_equilibrium(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_rescale_all_quantities(verbose=verbose, *args, **kwargs)

    return True


if __name__ == '__main__':
    print(('Testing test_rescale.py:', _run_all_tests(verbose=True)))
