#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Mar  2 23:39:52 2018

@author: erwan
"""

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

def test_compression__fast(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'))
    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


def test_update_transmittance__fast(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 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'))
    assert s.get_vars() == ['abscoeff']
    assert s.is_at_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 # force non equilibrium
    assert not s.is_at_equilibrium()
    # 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 by setting Tvib=2000 
    s2 = s1.copy()
    s2.conditions['Tvib'] = 2000 # force non equilibrium
    assert not s2.is_at_equilibrium()
    s2.update('radiance_noslit') 

    # update s1 now (at equilibrium)
    s1.update('radiance_noslit')
    
    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')
    
    # Compare
    assert s1.compare_with(s2,spectra_only='radiance_noslit', plot=plot)
    

def _run_all_tests(verbose=True, warnings=True, *args, **kwargs):
    test_compression__fast(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_update_transmittance__fast(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_get_recompute(verbose=verbose, warnings=warnings, *args, **kwargs)
    test_recompute_equilibrium(verbose=verbose, warnings=warnings, *args, **kwargs)
    
    return True

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