# -*- coding: utf-8 -*-
"""
Created on Tue Jan  9 17:14:01 2018

@author: erwan

Test broadening against HAPI and tabulated data

We're looking at CO(0->1) line 'R1' at 2150.86 cm-1

"""

from __future__ import unicode_literals, print_function, absolute_import, division
from radis.lbl.factory import SpectrumFactory
from radis.spectrum.spectrum import Spectrum
from radis import plot_diff, get_residual_integral, get_residual
from radis.misc.utils import DatabankNotFound
from radis.test.utils import IgnoreMissingDatabase, setup_test_line_databases
from radis.misc.printer import printm
from os.path import join, dirname
from numpy import isclose
import matplotlib.pyplot as plt
import pytest


@pytest.mark.fast
@pytest.mark.needs_connection
def test_broadening_vs_hapi(rtol=1e-2, verbose=True, plot=False, *args, **kwargs):
    '''
    Test broadening against HAPI and tabulated data

    We're looking at CO(0->1) line 'R1' at 2150.86 cm-1
    '''
    from radis.io.hapi import db_begin, fetch, tableList, absorptionCoefficient_Voigt

    if plot:  # Make sure matplotlib is interactive so that test are not stuck in pytest
        plt.ion()

    setup_test_line_databases()  # add HITRAN-CO-TEST in ~/.radis if not there

    # Conditions
    T = 3000
    p = 0.0001
    wstep = 0.001
    wmin = 2150  # cm-1
    wmax = 2152  # cm-1
    broadening_max_width = 10  # cm-1

    # %% HITRAN calculation
    # -----------

    # Generate HAPI database locally
    db_begin(join(dirname(__file__), __file__.replace('.py', '_HAPIdata')))
    if not 'CO' in tableList():   # only if data not downloaded already
        fetch('CO', 5, 1, wmin-broadening_max_width /
              2, wmax+broadening_max_width/2)
        # HAPI doesnt correct for side effects

    # Calculate with HAPI
    nu, coef = absorptionCoefficient_Voigt(SourceTables='CO',
                                           Environment={'T': T,  # K
                                                        'p': p/1.01325,  # atm
                                                        },
                                           WavenumberStep=wstep,
                                           HITRAN_units=False)

    s_hapi = Spectrum.from_array(nu, coef, 'abscoeff', 'cm-1', 'cm_1',
                                 conditions={'Tgas': T},
                                 name='HAPI')

    # %% Calculate with RADIS
    # ----------
    sf = SpectrumFactory(
        wavenum_min=wmin,
        wavenum_max=wmax,
        mole_fraction=1,
        path_length=1,   # doesnt change anything
        wstep=wstep,
        pressure=p,
        broadening_max_width=broadening_max_width,
        isotope=[1],
        warnings={'MissingSelfBroadeningWarning':'ignore',
                  'NegativeEnergiesWarning':'ignore',
                  'HighTemperatureWarning':'ignore',
                  'GaussianBroadeningWarning':'ignore'}
        )  # 0.2)
    sf.load_databank('HITRAN-CO-TEST')
#    s = pl.non_eq_spectrum(Tvib=T, Trot=T, Ttrans=T)
    s = sf.eq_spectrum(Tgas=T, name='RADIS')
    
    
    if plot:  # plot broadening of line of largest linestrength
        sf.plot_broadening(i=sf.df1.S.idxmax())

    # Plot and compare
    res = abs(get_residual_integral(s, s_hapi, 'abscoeff'))
    if plot:
        plot_diff(s, s_hapi, var='abscoeff',
                  title='{0} bar, {1} K (residual {2:.2g}%)'.format(p, T, res*100), show_points=False)
        plt.xlim((wmin, wmax))
    if verbose:
        printm('residual:', res)
    assert res < rtol


@pytest.mark.fast
def test_broadening_methods_different_conditions(verbose=True, plot=False, *args, **kwargs):
    '''
    Test direct Voigt broadening vs convolution of Gaussian x Lorentzian
    for different spectral grid resolution
    
    Notes
    ----- 
    
    Reference broadening calculated manually with the HWHM formula of 
    `HITRAN.org <https://hitran.org/docs/definitions-and-units/>`_
    '''

    if plot:  # Make sure matplotlib is interactive so that test are not stuck in pytest
        plt.ion()

    setup_test_line_databases()  # add HITRAN-CO-TEST in ~/.radis if not there

    # Conditions
    wstep = 0.005
    wmin = 2150.4  # cm-1
    wmax = 2151.4  # cm-1
    broadening_max_width = 2  # cm-1
    
    
    for (T, p, fwhm_lorentz, fwhm_gauss) in [
         # K, bar, expected FWHM for Lotentz, gauss (cm-1)
         (3000, 1,      0.02849411,      0.01594728),
         (300,  1,      0.16023415,      0.00504297),
         (3000, 0.01,   0.00028494,      0.01594728),
         ]:

        # %% Calculate with RADIS
        # ----------
        sf = SpectrumFactory(
            wavenum_min=wmin,
            wavenum_max=wmax,
            mole_fraction=1,
            path_length=1,   # doesnt change anything
            wstep=wstep,
            pressure=p,
            broadening_max_width=broadening_max_width,
            isotope='1',
            verbose=False,
            warnings={'MissingSelfBroadeningWarning':'ignore',
                      'NegativeEnergiesWarning':'ignore',
                      'HighTemperatureWarning':'ignore',
                      'OutOfRangeLinesWarning':'ignore',
                      'GaussianBroadeningWarning':'ignore',
                      'CollisionalBroadeningWarning':'ignore'}
            )
        sf.load_databank('HITRAN-CO-TEST')
        # Manually filter line database, keep one line only:
        sf.df0.drop(sf.df0[sf.df0.vu!=1].index, inplace=True)
        assert isclose(sf.df0.wav, 2150.856008)
        
        # Calculate spectra (different broadening methods)
        sf._broadening_method = 'voigt'
        s_voigt = sf.eq_spectrum(Tgas=T, name='direct')
        
        # assert broadening FWHM are correct
        assert isclose(2*float(sf.df1.hwhm_gauss), fwhm_gauss)
        assert isclose(2*float(sf.df1.hwhm_lorentz), fwhm_lorentz)
        
        
        sf._broadening_method = 'convolve'
        s_convolve = sf.eq_spectrum(Tgas=T, name='convolve')
    
        # assert broadening FWHM are correct
        assert isclose(2*float(sf.df1.hwhm_gauss), fwhm_gauss)
        assert isclose(2*float(sf.df1.hwhm_lorentz), fwhm_lorentz)
        
        
        res = get_residual(s_voigt, s_convolve, 'abscoeff')
        
        if verbose:
            print('{0} K, {1} bar: FWHM lorentz = {2:.3f} cm-1, FWHM gauss = {3:.3f} cm-1'.format(
                    T, p, 2*float(sf.df1.hwhm_lorentz), 2*float(sf.df1.hwhm_gauss)))

        if plot:
            plot_diff(s_voigt, s_convolve, 'abscoeff', 
                      title=r'T {0} K, p {1} bar: w$_\mathrm{{L}}$ {2:.3f}, w$_\mathrm{{G}}$ {3:.3f} cm$^{{-1}}$'.format(
                      T, p, 2*float(sf.df1.hwhm_lorentz), float(sf.df1.hwhm_gauss)))
        
        # assert all broadening methods match
        assert res < 2e-4
        
        
@pytest.mark.fast
def test_broadening_methods_different_wstep(verbose=True, plot=False, *args, **kwargs):
    '''
    Test direct Voigt broadening vs convolution of Gaussian x Lorentzian
    for different spectral grid resolution
    '''

    if plot:  # Make sure matplotlib is interactive so that test are not stuck in pytest
        plt.ion()

    setup_test_line_databases()  # add HITRAN-CO-TEST in ~/.radis if not there

    # Conditions
    T = 3000
    p = 1
    wmin = 2150  # cm-1
    wmax = 2152  # cm-1
    broadening_max_width = 10  # cm-1
    
    for i, wstep in enumerate([0.01, 0.1, 0.5]):

        # %% Calculate with RADIS
        # ----------
        sf = SpectrumFactory(
            wavenum_min=wmin,
            wavenum_max=wmax,
            mole_fraction=1,
            path_length=1,   # doesnt change anything
            wstep=wstep,
            pressure=p,
            broadening_max_width=broadening_max_width,
            isotope='1',
            verbose=False,
            warnings={'MissingSelfBroadeningWarning':'ignore',
                      'NegativeEnergiesWarning':'ignore',
                      'HighTemperatureWarning':'ignore',
                      'GaussianBroadeningWarning':'ignore'}
            )  # 0.2)
        sf.load_databank('HITRAN-CO-TEST')
    #    s = pl.non_eq_spectrum(Tvib=T, Trot=T, Ttrans=T)
        sf._broadening_method = 'voigt'
        s_voigt = sf.eq_spectrum(Tgas=T, name='direct')
        
        sf._broadening_method = 'convolve'
        s_convolve = sf.eq_spectrum(Tgas=T, name='convolve')
    
        res = get_residual(s_voigt, s_convolve, 'abscoeff')
        
        if verbose:
            print('Residual:', res)

        # plot the last one    
        if plot:
            plot_diff(s_voigt, s_convolve, 'abscoeff', nfig='test_voigt_broadening_methods'+str(i),
                      title='P {0} bar, T {1} K, wstep {2} cm-1'.format(p, T, wstep))
        
        assert res < 2e-4
        
        
@pytest.mark.fast
def test_broadening_DLM(verbose=True, plot=False, *args, **kwargs):
    '''
    Test use of lineshape template for broadening calculation. 
    
    Ensures that results are the same with and without DLM.
    '''

    if plot:  # Make sure matplotlib is interactive so that test are not stuck in pytest
        plt.ion()

    setup_test_line_databases()  # add HITRAN-CO-TEST in ~/.radis if not there

    # Conditions
    T = 3000
    p = 1
    wstep = 0.002
    wmin = 2150  # cm-1
    wmax = 2152  # cm-1
    broadening_max_width = 10  # cm-1

    # %% Calculate with RADIS
    # ----------
    sf = SpectrumFactory(
        wavenum_min=wmin,
        wavenum_max=wmax,
        mole_fraction=1,
        path_length=1,   # doesnt change anything
        wstep=wstep,
        pressure=p,
        broadening_max_width=broadening_max_width,
        isotope='1',
        verbose=False,
        warnings={'MissingSelfBroadeningWarning':'ignore',
                  'NegativeEnergiesWarning':'ignore',
                  'HighTemperatureWarning':'ignore',
                  'GaussianBroadeningWarning':'ignore'}
        )  # 0.2)
    sf.load_databank('HITRAN-CO-TEST')
    
    # Reference: calculate without DLM 
    assert sf.misc['chunksize'] is None
    s_ref = sf.eq_spectrum(Tgas=T)
    s_ref.name = 'Reference ({0:.2f}s)'.format(s_ref.conditions['calculation_time'])
    
    # DLM:
    sf.misc['chunksize'] = 'DLM'
    sf._broadening_method = 'convolve'
    s_dlm = sf.eq_spectrum(Tgas=T)
    s_dlm.name = 'DLM ({0:.2f}s)'.format(s_dlm.conditions['calculation_time'])
    # DLM Voigt with Whiting approximation:
    sf._broadening_method = 'voigt'
    s_dlm_voigt = sf.eq_spectrum(Tgas=T)
    s_dlm_voigt.name = 'DLM Whiting ({0:.2f}s)'.format(s_dlm_voigt.conditions['calculation_time'])
    
    # Compare
    res = get_residual(s_ref, s_dlm, 'abscoeff')
    res_voigt = get_residual(s_dlm, s_dlm_voigt, 'abscoeff')
    
    if verbose:
        print('Residual:', res)

    # plot the last one    
    if plot:
        plot_diff(s_ref, s_dlm, 'abscoeff')
        plt.legend()
        plot_diff(s_dlm, s_dlm_voigt, 'abscoeff')
    
    assert res < 1.2e-5
    assert res_voigt < 1e-5

@pytest.mark.fast
def test_broadening_DLM_FT(verbose=True, plot=False, *args, **kwargs):
    '''
    Test use of DLM with and without Fourier Transform
    
    Ensures that results are the same, and compare calculation times.
    '''

    if plot:  # Make sure matplotlib is interactive so that test are not stuck in pytest
        plt.ion()

    setup_test_line_databases()  # add HITRAN-CO-TEST in ~/.radis if not there

    # Conditions
    T = 3000
    p = 1
    wstep = 0.002
    wmin = 2000  # cm-1
    wmax = 2300  # cm-1
    broadening_max_width = 10  # cm-1

    # %% Calculate with RADIS
    # ----------
    sf = SpectrumFactory(
        wavenum_min=wmin,
        wavenum_max=wmax,
        mole_fraction=1,
        path_length=1,   # doesnt change anything
        wstep=wstep,
        pressure=p,
        broadening_max_width=broadening_max_width,
        isotope='1',
        verbose=verbose,
        chunksize='DLM',
        warnings={'MissingSelfBroadeningWarning':'ignore',
                  'NegativeEnergiesWarning':'ignore',
                  'HighTemperatureWarning':'ignore',
                  'GaussianBroadeningWarning':'ignore'}
        )  # 0.2)
    sf.load_databank('HITRAN-CO-TEST')
    
    # DLM, real space
    if verbose: print('\nConvolve version \n')
    s_dlm = sf.eq_spectrum(Tgas=T)
    s_dlm.name = 'DLM ({0:.2f}s)'.format(s_dlm.conditions['calculation_time'])
    
    # DLM , with Fourier
    if verbose: print('\nFFT version \n')
    sf._broadening_method = 'fft'
    s_dlm_fft = sf.eq_spectrum(Tgas=T)
    s_dlm_fft.name = 'DLM FFT ({0:.2f}s)'.format(s_dlm_fft.conditions['calculation_time'])
    
    # Compare
    res = get_residual(s_dlm, s_dlm_fft, 'abscoeff')
    
    if verbose:
        print('Residual:', res)

    # plot    
    if plot:
        plot_diff(s_dlm, s_dlm_fft, 'abscoeff')
        plt.legend()
        
    assert res < 5e-6

@pytest.mark.fast
def test_broadening_DLM_noneq(verbose=True, plot=False, *args, **kwargs):
    '''
    Test Noneq version of DLM and makes sure it gives the same results as the eq
    one when used with Tvib=Trot 
    
    '''

    if plot:  # Make sure matplotlib is interactive so that test are not stuck in pytest
        plt.ion()

    setup_test_line_databases()  # add HITRAN-CO2-TEST in ~/.radis if not there

    # Conditions
    p = 1
    wstep = 0.002
    wmin = 2380  # cm-1
    wmax = 2400  # cm-1
    broadening_max_width = 10  # cm-1

    # %% Calculate with RADIS
    # ----------
    sf = SpectrumFactory(
        wavenum_min=wmin,
        wavenum_max=wmax,
        mole_fraction=1,
        path_length=1,   # doesnt change anything
        wstep=wstep,
        pressure=p,
        broadening_max_width=broadening_max_width,
        isotope='1',
        verbose=3,
        warnings={'MissingSelfBroadeningWarning':'ignore',
                  'NegativeEnergiesWarning':'ignore',
                  'HighTemperatureWarning':'ignore',
                  'GaussianBroadeningWarning':'ignore'}
        )  # 0.2)
    sf.load_databank('HITRAN-CO2-TEST')
    
    # DLM:
    sf.misc['chunksize'] = 'DLM'
    s_dlm_eq = sf.eq_spectrum(Tgas=3000)
    s_dlm_eq.name = 'DLM eq ({0:.2f}s)'.format(s_dlm_eq.conditions['calculation_time'])
    
    s_dlm_noneq = sf.non_eq_spectrum(Tvib=3000, Trot=3000)
    s_dlm_noneq.name = 'DLM noneq ({0:.2f}s)'.format(s_dlm_noneq.conditions['calculation_time'])
    
    # Compare
    res = get_residual(s_dlm_eq, s_dlm_noneq, 'radiance_noslit')
    
    if verbose:
        print('Residual:', res)

    # plot
    if plot:
        plot_diff(s_dlm_eq, s_dlm_noneq)
    
    assert res <= 4e-5 



#@pytest.mark.needs_config_file
#@pytest.mark.needs_db_HITEMP_CO2_DUNHAM
@pytest.mark.needs_connection
def test_abscoeff_continuum(plot=False, verbose=2, warnings=True, threshold=0.1, *args, **kwargs):
    ''' 
    Test calculation with pseudo-continuum

    Assert results on abscoeff dont change
    
    
    Notes
    -----
    
    Uses HITRAN so it can deployed and tested on [Travis]_, but we should switch
    to HITEMP if some HITEMP files can be downloaded automatically at the 
    execution time. 

    '''

    if plot:  # Make sure matplotlib is interactive so that test are not stuck in pytest
        plt.ion()

    try:
        if verbose:
            printm('>>> test_abscoeff_continuum')

        sf = SpectrumFactory(wavelength_min=4200, 
                             wavelength_max=4500,
                             parallel=False, bplot=False, cutoff=1e-23,
                             molecule='CO2',
                             isotope='1,2', 
                             db_use_cached=True,
                             broadening_max_width=10,
                             path_length=0.1,
                             mole_fraction=1e-3,
                             medium='vacuum',
                             verbose=verbose)
        sf.warnings.update({
            'MissingSelfBroadeningWarning': 'ignore',
            'NegativeEnergiesWarning': 'ignore',
            'LinestrengthCutoffWarning': 'ignore',
            'HighTemperatureWarning':'ignore'})
        sf.fetch_databank()   # uses HITRAN: not really valid at this temperature, but runs on all machines without install
#        sf.load_databank('HITEMP-CO2-DUNHAM')       # to take a real advantage of abscoeff continuum, should calculate with HITEMP
        sf._export_continuum = True   # activate it 

        # Calculate one without pseudo-continuum
        sf.params.pseudo_continuum_threshold = 0
        s1 = sf.eq_spectrum(Tgas=2000)
        s1.name = 'All lines resolved ({0}) ({1:.1f}s)'.format(s1.conditions['lines_calculated'],
                                                               s1.conditions['calculation_time'])
        assert s1.conditions['pseudo_continuum_threshold'] == 0

        # Calculate one with pseudo-continuum
        sf.params.pseudo_continuum_threshold = threshold
        s2 = sf.eq_spectrum(Tgas=2000)
        s2.name = 'Semi-continuum + {0} lines ({1:.1f}s)'.format(s2.conditions['lines_calculated'],
                                                                 s2.conditions['calculation_time'])
        assert s2.conditions['pseudo_continuum_threshold'] == threshold
        assert 'abscoeff_continuum' in s2.get_vars()

        # Plot
        if plot:
            plot_diff(s1, s2, 'radiance_noslit', Iunit='µW/cm2/sr/nm',
                      nfig='test_abscoeff_continuum: diff')

            s2.plot('abscoeff', label='Full spectrum',
                    nfig='test_abscoeff_continuum: show continuum', force=True)
            s2.plot('abscoeff_continuum', nfig='same', label='Pseudo-continuum'.format(
                    s2.conditions['lines_in_continuum']),
                    force=True)
            plt.legend()

        # Compare
        res = get_residual(s1, s2, 'abscoeff')
        if verbose:
            printm('residual:', res)

        globals().update(locals())

        assert res < 1.32e-6

    except DatabankNotFound as err:
        assert IgnoreMissingDatabase(err, __file__, warnings)

#@pytest.mark.needs_config_file
#@pytest.mark.needs_db_HITEMP_CO2_DUNHAM
@pytest.mark.needs_connection
def test_noneq_continuum(plot=False, verbose=2, warnings=True, *args, **kwargs):
    ''' 
    Test calculation with pseudo-continuum under nonequilibrium

    Assert results on emisscoeff dont change

    
    Notes
    -----
    
    Uses HITRAN so it can deployed and tested on [Travis]_, but we should switch
    to HITEMP if some HITEMP files can be downloaded automatically at the 
    execution time. 

    '''
   
    if plot:  # Make sure matplotlib is interactive so that test are not stuck in pytest
        plt.ion()

    try:
        if verbose:
            printm('>>> test_noneq_continuum')

        sf = SpectrumFactory(wavelength_min=4200, wavelength_max=4500,
                             parallel=False, bplot=False, cutoff=1e-23,
                             molecule='CO2',
                             isotope='1,2', db_use_cached=True,
                             broadening_max_width=10,
                             path_length=0.1,
                             mole_fraction=1e-3,
                             medium='vacuum',
                             verbose=verbose)
        sf.warnings.update({
            'MissingSelfBroadeningWarning': 'ignore',
            'NegativeEnergiesWarning': 'ignore',
            'LinestrengthCutoffWarning': 'ignore',
            'HighTemperatureWarning':'ignore'})
        sf.fetch_databank()   # uses HITRAN: not really valid at this temperature, but runs on all machines without install
#        sf.load_databank('HITEMP-CO2-DUNHAM')       # to take a real advantage of abscoeff continuum, should calculate with HITEMP
        sf._export_continuum = True   # activate it 

        # Calculate one without pseudo-continuum
        sf.params.pseudo_continuum_threshold = 0
        s1 = sf.non_eq_spectrum(Tvib=2000, Trot=1000)
        s1.name = 'All lines resolved ({0}) ({1:.1f}s)'.format(s1.conditions['lines_calculated'],
                                                               s1.conditions['calculation_time'])
        assert s1.conditions['pseudo_continuum_threshold'] == 0

        # Calculate one with pseudo-continuum
        sf.params.pseudo_continuum_threshold = 0.05
        s2 = sf.non_eq_spectrum(Tvib=2000, Trot=1000)
        s2.name = 'Semi-continuum + {0} lines ({1:.1f}s)'.format(s2.conditions['lines_calculated'],
                                                                 s2.conditions['calculation_time'])
        assert s2.conditions['pseudo_continuum_threshold'] == 0.05
        assert 'abscoeff_continuum' in s2.get_vars()
        assert 'emisscoeff_continuum' in s2.get_vars()

        # Plot
        if plot:
            plot_diff(s1, s2, 'radiance_noslit', Iunit='µW/cm2/sr/nm',
                      nfig='test_noneq_continuum: diff')

            plt.figure('test_noneq_continuum: show continuum').clear()
            s2.plot('emisscoeff', label=s2.name,
                    nfig='test_noneq_continuum: show continuum')
            s2.plot('emisscoeff_continuum', nfig='same', label='Pseudo-continuum (aggreg. {0:g} lines)'.format(
                    s2.conditions['lines_in_continuum']), force=True)

        # Compare
        res = get_residual(s1, s2, 'abscoeff') + get_residual(s1, s2, 'emisscoeff')
        if verbose:
            printm('residual:', res)

        assert res < 5.2e-6

    except DatabankNotFound as err:
        assert IgnoreMissingDatabase(err, __file__, warnings)



def _run_testcases(plot=False, verbose=True, *args, **kwargs):

   # Test broadening
    test_broadening_vs_hapi(plot=plot, verbose=verbose, *args, **kwargs)
    test_broadening_methods_different_conditions(plot=plot, verbose=verbose, *args, **kwargs)
    test_broadening_methods_different_wstep(plot=plot, verbose=verbose, *args, **kwargs)
    test_broadening_DLM(plot=plot, verbose=verbose, *args, **kwargs)
    test_broadening_DLM_FT(plot=plot, verbose=3, *args, **kwargs)
    test_broadening_DLM_noneq(plot=plot, verbose=verbose, *args, **kwargs)

    # Test pseudo-continuum
    test_abscoeff_continuum(plot=plot, verbose=verbose, *args, **kwargs)
    test_noneq_continuum(plot=plot, verbose=verbose, *args, **kwargs)
    
    return True


if __name__ == '__main__':
    printm('test_broadening: ', _run_testcases(plot=True, verbose=True, debug=False))
