import numpy as np
import scipy.constants as const
from astropy.io import fits
import scipy.interpolate as si
import glob
from ppxf import ppxf_util as P
import glob
import os

class Spectrum(object):
    """

    UPDATE THE DOCSTRING! A stripped down version of specTools in order to package this all together nicely
    Original Author: Ryan Houghton (20/4/11)
    Updated: Sam Vaughan (Early 2018)

    Purpose: A spectrum object to aid manipulation of wavelength/frequency and the associated
             flux values

    Inputs:
    This class MUST be initalised with EITHER
       lamspec - a NLAMxNSPEC numpy array with [:,0] being the wavelength array (AA) and [:,1:] being
                 the associated flux values (erg/s/cm**2/AA)
       muspec  - a NLAMxNSPEC numpy array with [:,0] being the frequency array (GHz) and [:,1:] being
                 the associated flux values (GJy = 1e-14 erg/s/cm**2/Hz)

    NOTE: 2xN seems a funny way to order the indicies but in numpy, this IS the efficient
    way to order the memory elements

    Definitions:
    Once you have initalised the class with either lamspec or muspec, you will find available
    the following variables (no matter if you initalised with lamspec or muspec):
       lam  - the wavelength array (AA)
       flam - the associated flux (erg/s/cm**2/AA)
       mu   - the frequency array (GHz)
       fmu  - the associated flux (GJy = 1e-14 erg/s/cm**2/Hz)

    Functions:
       calcABmag  - given a filter (spectrum class), return a magnitude on the AB system
       calcSTmag  - as above but for the ST system
       calcVEGAmag- as above but for the VEGA system; you must also supply a vega *spectrum*

    Notes:
       - You can define more than one flam/fmu for each lam/mu
       - Filters used in the functions should also be *spectrum* classes
       - Filters don't need to be sorted in mu or lam
    """

    global c, pc, d_sun

    def __init__(self, lam=None, lamspec=None, errlamspec=None, age=None, Z=None, wavesyst=None, userdict=None):

        """
        Inputs:
        -------
           lam     - the 1D array of wavelengths (becomes self.lam)
           lamspec - the (1D, 2D or 3D) array of fluxes on the specified lambda grid (becomes self.flam)
           errlamspec - standard deviation for pixels in lamspec (becomes self.eflam)
           age     - the 1D array of ages (Gyr) for the spectra
           Z       - the 1D array of metallicities for the spectra
           wavesyst - you MUST specify if the wavelengths are in the AIR or VAC system.
           userdict- add extra user-defined tags to the spectrum, such as airmass, water vapour etc.

        Notes:
        ------
        There are 3 ways to inialise a spectrum class:
            1. With a 1D lam/mu and a 1D spectrum (single spec mode).
                 - Magnitudes are returned as single values
            2. With a 1D lam/mu and a 2D spectrum array (NSPECxN{LAM/MU}) (multispec mode)
                 - In this case, you may specify AGE where len(age)=NSPEC.
                 - Magnitudes are returned as 1D arrays with NSPEC=len(age) elements
            #3. With a 1D lam/mu and a 3D spectrum array (NZxNAGExN{LAM/MU}) (multispec mode)
            #     - In this case, you can specify AGE and Z.
            #     - Magnitudes will be returned as 2D arrays with NZxNAGE elements
        """

        # start defining spectrum parts
        if lamspec is not None:
            # check that lam has been given
            if lam is None:
                raise Exception("If you give lamspec, you must also give lam")

            # make sure 2d
            flam = np.atleast_2d(lamspec)
            # get array size
            loc = flam.shape

            # check for bigger arrays
            # if len(loc)> 2: raise "lamspec not understood"

            # get sizes
            nlam = loc[1]
            nspec = loc[0]

            self.lam = lam
            self.flam = np.atleast_2d(singleOrList2Array(flam))

            if errlamspec is not None:
                eflam = np.atleast_2d(errlamspec)
                eloc = eflam.shape
                self.eflam = eflam
                # sanity check
                assert np.all(loc == eloc), "Flux and error arrays appear different sizes..."
            else:
                self.eflam = None

        # add age info
        if age is not None:
            self.age = singleOrList2Array(age)

            checkDims(self.age, "Age", self.flam.shape[:-1])
            self.logage = np.log10(self.age)
        else:
            self.age = None

        # add metallicitiy
        if Z is not None:

            self.Z = singleOrList2Array(Z)
            checkDims(self.Z, "Z", self.flam.shape[:-1])
        else:
            self.Z = None

        # add VAC or AIR
        if wavesyst is not None:
            if (wavesyst == "vac" or wavesyst == "VAC" or wavesyst == "Vac"):
                self.wavesyst = "vac"
            elif (wavesyst == "air" or wavesyst == "AIR" or wavesyst == "Air"):
                self.wavesyst = "air"
            else:
                raise ValueError("wavesyst not understood. Should be air or vac.")
        else:
            warn.warn("You failed to specify if the wavelength is defined in AIR or VAC units.")
            self.wavesyst = None

        # add user dictionary for extra info
        if userdict is not None:
            self.__userdict__ = userdict
            keys = list(userdict.keys())
            for key in keys:
                setattr(self, key, singleOrList2Array(userdict[key]))
        else:
            self.__userdict__ = None

def checkDims(var, varname, parentShape):
    assert (np.isscalar(var)) or (np.all(np.equal(np.array(var).shape, parentShape))), varname + " dimensions not understood."

def singleOrList2Array(invar):
    """
    If a single value, leave. If a list, convert to array. If array, leave as array.
    But for size-1 arrays/lits, convert back to scalar.
    """

    if isinstance(invar, list):
        # convert to array unless size-1
        rval = np.squeeze(np.array(invar))
        if rval.size == 1:
            rval = np.asscalar(rval)
    elif isinstance(invar, np.ndarray):
        # leave except if size-1
        rval = np.squeeze(invar)
        if rval.size == 1:
            rval = np.asscalar(rval)
    else:
        # leave
        rval = invar
    # return
    return rval

def load_varelem_CvD16ssps(varelem_template_location, imf='kroupa', verbose=True):

    '''
    Load the CvD16 spectra (response functions) with variable elemental abundances.

    Arguments:
        dirname (string): Base location of the stellar models directory
        folder (string): The folder inside the base directory which contains the response functions
        imf (string): The IMF of the response functions. Choices are 'kroupa' or 'salpeter'
        verbose (bool): Print information to the console
    Returns:
        (dict): A dictionary of spectra. See below for example

    Example:
        The spectra are returned in a dictionary. The dictionary is indexed by a string referring to the name of the respone function,
        e.g. 'Na+'. At each index, we have a Spectrum class which contains the response functions at various ages and metallicities

        >>> spectra=load_varelem_CvD16ssps(dirname='/path/to/dir', folder='folder', imf='kroupa')
        >>> spectra['Na+'] #Gives all RFs with [Na/H]=+0.3

    '''

    varelem_template_location = os.path.expanduser(varelem_template_location)

    if imf in ['kroupa', 'krpa', 'Kroupa', 'Krpa']:
        model_spectra = sorted(glob.glob('{}/atlas_ssp_*.krpa.s100'.format(varelem_template_location)))
        imf_name = 'krpa'
    elif imf in ['Salpeter', 'salpeter', 'salp', 'Salp']:
        model_spectra = sorted(glob.glob('{}/atlas_ssp_*.salp.s100'.format(varelem_template_location)))
        imf_name = 'salp'
    else:
        raise NameError('IMF type not understood')

    data = np.genfromtxt(model_spectra[0])
    lams = data[:, 0]

    model_Zs_names = ['m1.5', 'm1.0', 'm0.5', 'p0.0', 'p0.2']
    model_age_names = ['01', '03', '05', '09', '13']

    model_elem_order = ['Solar', 'Na+', 'Na-', 'Ca+', 'Ca-', 'Fe+', 'Fe-', 'C+', 'C-', 'a/Fe+', 'N+', 'N-', 'as/Fe+', 'Ti+', 'Ti-', 'Mg+', 'Mg-', 'Si+', 'Si-', 'T+', 'T-', 'Cr+', 'Mn+', 'Ba+', 'Ba-', 'Ni+', 'Co+', 'Eu+', 'Sr+', 'K+', 'V+', 'Cu+', 'Na+0.6', 'Na+0.9']

    Zs = [-1.5, -1.0, -0.5, 0.0, 0.2]
    ages = [float(a) for a in model_age_names]

    n_ages = len(model_age_names)
    n_zs = len(model_Zs_names)
    n_elems = len(model_elem_order)

    templates = np.empty((n_elems, n_ages, n_zs, len(lams)))

    for a, Z in enumerate(model_Zs_names):
        for b, age in enumerate(model_age_names):

            model = glob.glob('{}/atlas_ssp*t{}*{}*{}.s100'.format(varelem_template_location, age, Z, imf_name))[0]
            if verbose:
                print('Loading {}'.format(model))
            data = np.genfromtxt(model)

            for i, elem in enumerate(model_elem_order):
                templates[i, b, a, :] = data[:, i + 1]

    spectra = {}
    for i, elem in enumerate(model_elem_order):

        age_values = np.repeat(ages, n_zs).reshape(n_ages, n_zs)
        Z_values = np.repeat(Zs, n_ages).reshape(n_zs, n_ages).T

        spectra[elem] = Spectrum(lam=lams, lamspec=templates[i, :, :, :], age=age_values, Z=Z_values, wavesyst='vac', userdict={'elem': elem})

    return spectra

################################################################################################################################################################
def prepare_CvD_interpolator_twopartIMF(base_template_location, templates_lam_range, velscale, verbose=True, instrumental_resolution=None):
    """
    Set up the interpolator for the base SSP spectra using the log-rebinned templates we get from `prepare_CvD2_templates_twopartIMF`.

    Arguments:
        templates_lam_range (array-like): A 2 component vector with the start and stop wavelengths we want for the templates
        velscale (float): The velocity difference between two adjacent pixels in the log-rebinned spectrum. Use the same value
        for the templates as you *measure* from the spectrum!
        verbose (bool): Print information to the consolse

    Returns:
        (tuple): A two component tuple containing:
            * interp: the interpolate object. Axes are wavelength, age, Z, imf_x1 and imf_x2
            * logLam_template: the log-rebinned wavelength array of the templates
    """
    templates, logLam_template = prepare_CvD2_templates_twopartIMF(base_template_location, templates_lam_range, velscale, verbose=verbose, instrumental_resolution=instrumental_resolution)

    nimfs = 16
    ages = [1., 3., 5., 7., 9., 11.0, 13.5]
    Zs = [-1.5, -1.0, -0.5, 0.0, 0.2]
    n_imfs = 16
    imfs_X1 = 0.5 + np.arange(n_imfs) / 5.0
    imfs_X2 = 0.5 + np.arange(n_imfs) / 5.0

    linear_interp = si.RegularGridInterpolator(((logLam_template, ages, Zs, imfs_X1, imfs_X2)), templates, bounds_error=False, fill_value=None)

    return linear_interp, logLam_template

################################################################################################################################################################

def prepare_CvD2_templates_twopartIMF(template_location, templates_lam_range, velscale, verbose=True, instrumental_resolution=None):

    '''
    Load the CvD16 base spectra, those that vary age, [Z/H] and the IMF. The templates are log-rebinned to have a uniform wavelength
    spacing in **log** wavelength, to ensure the difference in velocity between two adjacent pixels is the same, regardless of wavelength.

    They are returned in an array of shape (n_ages, n_Zs, n_pixels).

    Arguments:
        templates_lam_range (array-like): A 2 component vector with the start and stop wavelengths we want for the templates
        velscale (float): The velocity difference between two adjacent pixels in the log-rebinned spectrum. Use the same value
        for the templates as you *measure* from the spectrum!
        verbose (bool): Print information to the consolse
    Returns:
        (tuple): A two component tuple:
            * templates (array): An array of templates with shape (N_wavelength, N_ages, N_metallicities, N_imf_x1, N_imf_x2)
            * logLam_template (array): The log-rebinned wavelength array of the templates
    '''

    template_glob=os.path.expanduser(f'{template_location}/VCJ_v8_mcut0.08_t*')

    vcj_models=sorted(glob.glob(template_glob))
    models=np.genfromtxt(vcj_models[-1])

    temp_lamdas=models[:, 0]

    n_ages=7
    n_zs=5
    n_imfs=16

    Zs=['m1.5', 'm1.0', 'm0.5', 'p0.0', 'p0.2']
    ages=['01.0', '03.0', '05.0', '07.0', '09.0', '11.0', '13.5']
    imfs_X1=0.5+np.arange(n_imfs)/5.0
    imfs_X2=0.5+np.arange(n_imfs)/5.0

    t_mask = ((temp_lamdas > templates_lam_range[0]) & (temp_lamdas <templates_lam_range[1]))

    y=models[t_mask, 1]
    x=temp_lamdas[t_mask]
    #Make a new lamda array, carrying on the delta lamdas of high resolution bit
    new_x=temp_lamdas[t_mask][0]+0.9*(np.arange(np.ceil((temp_lamdas[t_mask][-1]-temp_lamdas[t_mask][0])/0.9))+1)
    interp=si.interp1d(x, y, fill_value='extrapolate')
    out=interp(new_x)

    sspNew, logLam_template, template_velscale = P.log_rebin(templates_lam_range, out, velscale=velscale)
    templates=np.empty((len(sspNew), n_ages, n_zs, len(imfs_X1), len(imfs_X2)))

    #Resolution of the templates in km/s

    for a, Z in enumerate(Zs):
        for b, age in enumerate(ages):
            model=glob.glob(os.path.expanduser(f'{template_location}/VCJ_v8_mcut0.08_t{age}*{Z}.ssp.imf_varydoublex.s100'))[0]
            print('Loading {}'.format(model))
            data=np.genfromtxt(model)

            for c, counter1 in enumerate(imfs_X1):
                for d, counter2 in enumerate(imfs_X2):

                    #Interpolate templates onto a uniform wavelength grid and then log-rebin
                    y=data[:, c*n_imfs+d+1][t_mask]
                    x=temp_lamdas[t_mask]

                    #Make a new lamda array, carrying on the delta lamdas of high resolution bit
                    new_x=temp_lamdas[t_mask][0]+0.9*(np.arange(np.ceil((temp_lamdas[t_mask][-1]-temp_lamdas[t_mask][0])/0.9))+1)

                    interp=si.interp1d(x, y, fill_value='extrapolate')
                    out=interp(new_x)

                    #log rebin them
                    sspNew, logLam_template, template_velscale = P.log_rebin(templates_lam_range, out, velscale=velscale)
                    if instrumental_resolution is not None:
                        sspNew = P.gaussian_filter1d(sspNew, instrumental_resolution/velscale)

                    templates[:, b, a, c, d]=sspNew#/np.median(sspNew)

    return templates, logLam_template

################################################################################################################################################################

################################################################################################################################################################
def prepare_CvD_correction_interpolators(varelem_template_location, templates_lam_range, velscale, elements, verbose=True, element_imf='kroupa', instrumental_resolution=None):

    """
    Set up the interpolator for the response functions using the log-rebinned templates we get from `prepare_CvD2_element_templates`.
    The response functions allow for variation in a general set of elements (Fe, Ca, Mg, etc), a set of elements which can only be
    positive (Mn, K, V, etc), and Sodium, which can be extended from -0.3 dex to +1.0 dex. We have to treat each of these cases
    separately, so we end up with three different interpolation objects. You can also vary the effective temperature of the isochrone,
    but this is now deprecated in V2.

    Arguments:
        templates_lam_range (array-like): A 2 component vector with the start and stop wavelengths we want for the templates
        velscale (float): The velocity difference between two adjacent pixels in the log-rebinned spectrum. Use the same value
        for the templates as you *measure* from the spectrum!
        elements (array): This is an array of integers corresponding to the elements we want to vary in the interpolators. The
        numbers come from the headings in the SSP RF file. Na is element 1, Ca is element 2, etc. TODO: Improve this! Make it clearer

        verbose (bool): Print information to the consolse

    Returns:
        (tuple): A two component tuple containing:
            * correction_interps: A tuple of the interpolators. The order is [general_interp, na_interp, positive_only_interp, T_interp). TODO explain this more
            * logLam_template: The log-rebinned wavelength array of the templates
    """

    all_corrections, logLam_template=prepare_CvD2_element_templates(varelem_template_location, templates_lam_range, velscale, elements, verbose=verbose, element_imf=element_imf, instrumental_resolution=instrumental_resolution)

    general_templates, na_templates, carbon_templates, positive_only_templates, T_templates=all_corrections

    positive_only_elems, Na_elem, Carbon_elem, normal_elems=elements

    #It's not clear if the last value here should be 13.5 or 13
    ages=np.array([  1.,   3.,   5.,   9.,  13.0])
    Zs=[-1.5, -1.0, -0.5, 0.0, 0.2]

    elem_steps=[-0.45, -0.3, -0.2, -0.1, 0.0, 0.1, 0.2, 0.3, 0.45]
    Na_elem_steps=[-0.45, -0.3, -0.2, -0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
    C_elem_steps = [-0.2, -0.15, -0.1, -0.05, 0.0, 0.05, 0.1, 0.15, 0.2]
    positive_only_elem_steps=[0.0, 0.1, 0.2, 0.3, 0.45]
    T_steps=[-50.0, -40.0, -30.0, -20.0, -10.0, 0.0, 10.0, 20.0, 30.0, 40.0, 50.0]

    na_interp=si.RegularGridInterpolator(((Na_elem_steps, ages, Zs, logLam_template)), na_templates, bounds_error=False, fill_value=None, method='linear')

    carbon_interp = si.RegularGridInterpolator(((C_elem_steps, ages, Zs, logLam_template)), carbon_templates, bounds_error=False, fill_value=None, method='linear')

    T_interp=si.RegularGridInterpolator(((T_steps, ages, Zs, logLam_template)), T_templates, bounds_error=False, fill_value=None, method='linear')

    #If we only have one positive element to check, we need to do something different- can't have a dimension with only one element in RegularGridInterpolator apparently.
    if len(positive_only_elems)>1:
        positive_only_interp=si.RegularGridInterpolator(((np.arange(len(positive_only_elems)), positive_only_elem_steps, ages, Zs, logLam_template)), positive_only_templates, bounds_error=False, fill_value=None, method='linear')
    else:
        positive_only_interp=si.RegularGridInterpolator((positive_only_elem_steps, ages, Zs, logLam_template), positive_only_templates[0, :], bounds_error=False, fill_value=None, method='linear')

    if len(normal_elems)>1:
        general_interp=si.RegularGridInterpolator(((np.arange(len(normal_elems)), elem_steps, ages, Zs, logLam_template)), general_templates, bounds_error=False, fill_value=None, method='linear')
    else:
        general_interp=si.RegularGridInterpolator((elem_steps, ages, Zs, logLam_template), general_templates[0, :], bounds_error=False, fill_value=None, method='linear')

    correction_interps=[general_interp, na_interp, carbon_interp, positive_only_interp, T_interp]

    return correction_interps, logLam_template

##########

def new_wavelength_array_high_resolution(wavelengths):
    """
    The templates have different spacings in wavelength between the blue and red ends. This function makes a new set of wavelength values which ensures the same wavelength spacing throughout
    """

    new_wavelength_array = wavelengths[0]+0.9*(np.arange(np.ceil((wavelengths[-1]-wavelengths[0])/0.9))+1)

    return new_wavelength_array

def prepare_CvD2_element_templates(varelem_template_location, templates_lam_range, velscale, elements, verbose=True, element_imf='kroupa', instrumental_resolution=None):

    var_elem_spectra=load_varelem_CvD16ssps(varelem_template_location, imf=element_imf)

    ages=var_elem_spectra['Solar'].age[:, 0]
    Zs=var_elem_spectra['Solar'].Z[0, :]
    n_ages=len(ages)
    n_Zs=len(Zs)

    temp_lamdas=var_elem_spectra['Solar'].lam

    t_mask = ((temp_lamdas > templates_lam_range[0]) & (temp_lamdas <templates_lam_range[1]))

    positive_only_elems, Na_elem, Carbon_elem, normal_elems=elements

    elem_steps=[-0.45, -0.3, -0.2, -0.1, 0.0, 0.1, 0.2, 0.3, 0.45]
    Na_elem_steps=[-0.45, -0.3, -0.2, -0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
    C_elem_steps = [-0.2, -0.15, -0.1, -0.05, 0.0, 0.05, 0.1, 0.15, 0.2]
    positive_only_elem_steps=[0.0, 0.1, 0.2, 0.3, 0.45]
    T_steps=[-50.0, -40.0, -30.0, -20.0, -10.0, 0.0, 10.0, 20.0, 30.0, 40.0, 50.0]

    x=var_elem_spectra['Solar'].lam[t_mask]
    y=var_elem_spectra['Solar'].flam[-1, -1, t_mask]
    #Make a new lamda array, carrying on the delta lamdas of high resolution bit
    new_x = new_wavelength_array_high_resolution(var_elem_spectra['Solar'].lam[t_mask])
    interp=si.interp1d(x, y, fill_value='extrapolate')
    data=interp(new_x)

    sspNew, logLam_template, template_velscale = P.log_rebin(templates_lam_range, data, velscale=velscale)

    positive_only_templates = np.empty((len(positive_only_elems), len(positive_only_elem_steps), n_ages, n_Zs, len(sspNew)))
    general_templates = np.empty((len(normal_elems), len(elem_steps), n_ages, n_Zs, len(sspNew)))

    na_templates = np.empty((len(Na_elem_steps), n_ages, n_Zs, len(sspNew)))
    carbon_templates = np.empty((len(C_elem_steps), n_ages, n_Zs, len(sspNew)))
    T_templates = np.empty((len(T_steps), n_ages, n_Zs, len(sspNew)))

    print('Making the Positive-Only Correction templates')
    #Do the positve only correction templates:
    for a, elem in enumerate(positive_only_elems):
        print('\t{}'.format(elem))
        for b, step in enumerate(positive_only_elem_steps):
            for c, _ in enumerate(ages):
                for d, _ in enumerate(Zs):

                    if step !=0.0:
                        y=(var_elem_spectra[elem].flam[c, d, t_mask]/var_elem_spectra['Solar'].flam[c, d, t_mask] - 1.0)*((10**(step)-1.0)/(10**(0.3)-1.0))

                    else:
                        y=np.zeros_like(var_elem_spectra['Solar'].flam[c, d, t_mask])

                    x=var_elem_spectra[elem].lam[t_mask]
                    #Make a new lamda array, carrying on the delta lamdas of high resolution bit
                    new_x=new_wavelength_array_high_resolution(x)
                    interp=si.interp1d(x, y, fill_value='extrapolate')
                    data=interp(new_x)

                    sspNew, logLam_template, template_velscale = P.log_rebin(templates_lam_range, data, velscale=velscale)
                    if instrumental_resolution is not None:
                        sspNew = P.gaussian_filter1d(sspNew, instrumental_resolution/velscale)

                    positive_only_templates[a, b, c, d, :]=sspNew

    print('Making the General Correction templates')
    #Do the general templates
    for a, elem in enumerate(normal_elems):

        print('\t{}'.format(elem))
        for b, step in enumerate(elem_steps):
            for c, _ in enumerate(ages):
                for d, _ in enumerate(Zs):

                    if step>0.0:
                        e='{}+'.format(elem)
                        gen_step=step
                    elif step<0.0:
                        e='{}-'.format(elem)
                        gen_step=np.abs(step)

                    if step !=0.0:
                        y=(var_elem_spectra[e].flam[c, d, t_mask]/var_elem_spectra['Solar'].flam[c, d, t_mask]-1)*((10**(gen_step)-1.0)/(10**(0.3)-1.0))
                    else:
                        y=np.zeros_like(var_elem_spectra['Solar'].flam[c, d, t_mask])

                    x=var_elem_spectra[e].lam[t_mask]
                    #Make a new lamda array, carrying on the delta lamdas of high resolution bit
                    new_x=new_wavelength_array_high_resolution(x)
                    interp=si.interp1d(x, y, fill_value='extrapolate')
                    data=interp(new_x)

                    #data=util.gaussian_filter1d(data, sigs/velscale)

                    sspNew, logLam_template, template_velscale = P.log_rebin(templates_lam_range, data, velscale=velscale)
                    if instrumental_resolution is not None:
                        sspNew = P.gaussian_filter1d(sspNew, instrumental_resolution/velscale)

                    general_templates[a, b, c, d, :]=sspNew

    # Make the Carbon templates
    # Treat Carbon differently because it's templates are at +0.15 and -0.15 dex rather than +0.3 and -0.3 dex
    print('Making the Carbon Correction template')
    for b, step in enumerate(C_elem_steps):
        for c, _ in enumerate(ages):
            for d, _ in enumerate(Zs):

                if step>0.0:
                    e='{}+'.format(elem)
                    gen_step=step
                elif step<0.0:
                    e='{}-'.format(elem)
                    gen_step=np.abs(step)

                if step !=0.0:
                    y=(var_elem_spectra[e].flam[c, d, t_mask]/var_elem_spectra['Solar'].flam[c, d, t_mask]-1)*((10**(gen_step)-1.0)/(10**(0.15)-1.0))
                else:
                    y=np.zeros_like(var_elem_spectra['Solar'].flam[c, d, t_mask])

                x=var_elem_spectra[e].lam[t_mask]
                #Make a new lamda array, carrying on the delta lamdas of high resolution bit
                new_x=new_wavelength_array_high_resolution(x)
                interp=si.interp1d(x, y, fill_value='extrapolate')
                data=interp(new_x)

                #data=util.gaussian_filter1d(data, sigs/velscale)

                sspNew, logLam_template, template_velscale = P.log_rebin(templates_lam_range, data, velscale=velscale)
                if instrumental_resolution is not None:
                        sspNew = P.gaussian_filter1d(sspNew, instrumental_resolution/velscale)

                carbon_templates[b, c, d, :]=sspNew

    #Do the Na templates:
    print('Making the Na Correction template')
    for a, step in enumerate(Na_elem_steps):
        for b, _ in enumerate(ages):
            for c, _ in enumerate(Zs):

                if step <0.0:
                    e='Na-'
                    base_enhancement=0.3
                    Na_step=np.abs(step)
                elif 0.0<=step<0.45:
                    e='Na+'
                    base_enhancement=0.3
                    Na_step=step
                elif 0.45<=step<0.75:
                    e='Na+0.6'
                    base_enhancement=0.6
                    Na_step=step
                elif 0.75<=step<1.0:
                    e='Na+0.9'
                    base_enhancement=0.9
                    Na_step=step

                if step !=0.0:
                    y=(var_elem_spectra[e].flam[b, c, t_mask]/var_elem_spectra['Solar'].flam[b, c, t_mask]-1)*((10**(Na_step)-1.0)/(10**(base_enhancement)-1.0))

                else:

                    y=np.zeros_like(var_elem_spectra['Solar'].flam[b, c, t_mask])

                x=var_elem_spectra[e].lam[t_mask]
                #Make a new lamda array, carrying on the delta lamdas of high resolution bit
                new_x=new_wavelength_array_high_resolution(x)
                interp=si.interp1d(x, y, fill_value='extrapolate')
                data=interp(new_x)

                sspNew, logLam_template, template_velscale = P.log_rebin(templates_lam_range, data, velscale=velscale)
                if instrumental_resolution is not None:
                        sspNew = P.gaussian_filter1d(sspNew, instrumental_resolution/velscale)

                na_templates[a, b, c, :]=sspNew

    print('Making the Temperature Correction template')
    for a, step in enumerate(T_steps):
        for b, _ in enumerate(ages):
            for c, _ in enumerate(Zs):

                if step>0.0:
                    e='T+'
                    T_step=step
                elif step<0.0:
                    e='T-'
                    T_step=np.abs(step)

                if step !=0.0:
                    y=(var_elem_spectra[e].flam[b, c, t_mask]/var_elem_spectra['Solar'].flam[b, c, t_mask]-1)*(T_step/50.0)

                else:

                    y=np.zeros_like(var_elem_spectra['Solar'].flam[b, c, t_mask])

                x=var_elem_spectra[e].lam[t_mask]
                #Make a new lamda array, carrying on the delta lamdas of high resolution bit
                new_x=new_wavelength_array_high_resolution(x)
                interp=si.interp1d(x, y, fill_value='extrapolate')
                data=interp(new_x)

                sspNew, logLam_template, template_velscale = P.log_rebin(templates_lam_range, data, velscale=velscale)
                if instrumental_resolution is not None:
                        sspNew = P.gaussian_filter1d(sspNew, instrumental_resolution/velscale)

                T_templates[a, b, c, :]=sspNew

    return [general_templates, na_templates, carbon_templates, positive_only_templates, T_templates], logLam_template


"""BLPose Train: Generator

Author: Bo Lin (@linbo0518)
Date: 2020-09-11
"""

import json
import math
import numpy as np

__all__ = ["dict2json", "json2dict", "TargetGenerator"]

def dict2json(dict_obj: dict, filename: str) -> None:
    with open(filename, "w") as f:
        json.dump(dict_obj, f, indent=2)

def json2dict(filename: str) -> dict:
    dict_obj: dict
    with open(filename, "r") as f:
        dict_obj = json.load(f)
    return dict_obj

def _gaussian_kernel(height, width, x, y, sigma=1.75):
    grid_y, grid_x = np.mgrid[0:height, 0:width]
    return np.exp(-((grid_x - x) ** 2 + (grid_y - y) ** 2) / (2.0 * sigma ** 2))

def _part_affinity_field(height, width, x1, y1, x2, y2, thickness=1):
    pafmap = np.zeros((2, height, width), dtype=np.float32)
    countmap = np.zeros_like(pafmap, dtype=np.uint)

    limb_vec_x = x2 - x1
    limb_vec_y = y2 - y1
    limb_len = math.sqrt(limb_vec_x ** 2 + limb_vec_y ** 2)

    if limb_len < 1e-7:
        return pafmap, countmap

    limb_unit_x = limb_vec_x / limb_len
    limb_unit_y = limb_vec_y / limb_len

    min_x = max(min(x1, x2) - thickness, 0)
    max_x = min(max(x1, x2) + thickness, width)
    min_y = max(min(y1, y2) - thickness, 0)
    max_y = min(max(y1, y2) + thickness, height)

    grid_y, grid_x = np.mgrid[min_y:max_y, min_x:max_x]
    x_12 = grid_x - x1
    y_12 = grid_y - y1

    limb_width = np.abs(x_12 * limb_unit_y - y_12 * limb_unit_x)
    limb_mask = limb_width < thickness

    pafmap[:, grid_y, grid_x] = np.repeat(
        limb_mask[np.newaxis, :, :], repeats=2, axis=0
    )
    pafmap[0, grid_y, grid_x] *= limb_unit_x
    pafmap[1, grid_y, grid_x] *= limb_unit_y

    limb_mask = (pafmap[0] != 0) | (pafmap[1] != 0)
    countmap[:, limb_mask] += 1
    return pafmap, countmap

class TargetGenerator:
    def __init__(
        self, n_keypoints, limbs, image_shape, stride, sigma=1.75, thickness=1
    ):
        self.n_keypoints = n_keypoints
        self.n_limbs = len(limbs)
        self.limbs = limbs
        height, width, _ = image_shape
        self.height = height // stride  # or round
        self.width = width // stride  # or round
        self.stride = stride
        self.sigma = sigma
        self.thickness = thickness

    def __call__(self, annotation):
        heatmap = np.zeros(
            (self.n_keypoints + 1, self.height, self.width), dtype=np.float32
        )
        pafmap = np.zeros((self.n_limbs * 2, self.height, self.width), dtype=np.float32)
        countmap = np.zeros_like(pafmap, dtype=np.uint)
        for keypoints in annotation:
            heat = self.gen_heatmap(keypoints)
            heatmap = np.maximum(heatmap, heat)

            paf, count = self.gen_pafmap(keypoints)
            pafmap += paf
            countmap += count

        countmap[countmap == 0] = 1
        pafmap /= countmap
        heatmap[-1] = 1.0 - np.max(heatmap[:-1], axis=0)
        return heatmap, pafmap

    def gen_heatmap(self, keypoints):
        heatmap = np.zeros(
            (self.n_keypoints + 1, self.height, self.width), dtype=np.float32
        )
        for idx, keypoint in enumerate(keypoints):
            x, y, v = keypoint
            if v == 0:
                continue
            x, y = round(x / self.stride), round(y / self.stride)

            heatmap[idx] = _gaussian_kernel(self.height, self.width, x, y, self.sigma)
            heatmap[idx][heatmap[idx] > 1.0] = 1
            heatmap[idx][heatmap[idx] < 0.01] = 0
        return heatmap

    def gen_pafmap(self, keypoints):
        pafmap = np.zeros((self.n_limbs * 2, self.height, self.width), dtype=np.float32)
        countmap = np.zeros_like(pafmap, dtype=np.uint)
        for idx, limb in zip(range(0, self.n_limbs * 2, 2), self.limbs):
            x1, y1, v1 = keypoints[limb[0]]
            x2, y2, v2 = keypoints[limb[1]]
            if v1 == 0 or v2 == 0:
                continue
            x1, y1 = round(x1 / self.stride), round(y1 / self.stride)
            x2, y2 = round(x2 / self.stride), round(y2 / self.stride)

            paf, count = _part_affinity_field(
                self.height, self.width, x1, y1, x2, y2, self.thickness
            )
            pafmap[idx : idx + 2] = paf
            countmap[idx : idx + 2] = count
        return pafmap, countmap


"""A stochastic policy"""

import numpy as np
import tensorflow as tf

from context import flags
import env_info
from utils import build_mlp, trainable_vars

class StochasticPolicy:
    """
    Provides tensorflow interface to sampling-with-density and acting
    according to it.
    """

    def tf_sample_action_with_log_prob(self, obs_ns):
        """Return a sample from the policy, along with its log probability."""
        raise NotImplementedError  # return sample, log prob

    def tf_greedy_action(self, obs_ns):
        """Return a sample from the greedy version of the policy."""
        raise NotImplementedError

    def act(self, obs_ns):
        """Eager version of sampling an action"""
        raise NotImplementedError

    def greedy_act(self, obs_ns):
        """Eager version of sampling a greedy action"""
        raise NotImplementedError

    def expectation(self, obs_ns, fn):
        """
        Given a function fn : (acs_na, log pi(acs_na)) -> reals
        which maps a TF batch of actions and their likelihoods (taken according
        to this stochastic policy), this returns a tensor which, when
        differentiated wrt the policy parameters,
        gives an estimator for the gradient of E[fn(A)], where A follows
        this stochastic policy. The expectation is wrt the state distribution
        an empirical sample of which should be given as the obs_ns argument
        here.
        """
        pass

    def tf_report(self, reporter, obs_ns):
        """
        Update a TF reporter with various statistics according to the given
        sample of states.
        """
        pass

class SquashedGaussianPolicy(StochasticPolicy):
    """
    This stochastic policy implements a unimodal multivariate
    Gaussian distribution with a diagonal covariance, with the both the mean
    and covariance conditional on the state. This generates
    a latent distribution, which then is then clipped for stability
    and squashed into a compact [-1, 1] interval.
    """

    def __init__(self, scope='sac/policy'):
        """
        Generates a diagonal gaussian policy which keeps its log standard
        deviations between the lower and upper bounds provided as arguments.
        """
        self._scope = scope
        self._obs_ph_ns = tf.placeholder(tf.float32, [None, env_info.ob_dim()])
        self._acs_na, _ = self.tf_sample_action_with_log_prob(self._obs_ph_ns)
        self._greedy_acs_na = self.tf_greedy_action(self._obs_ph_ns)
        self.variables = trainable_vars(self._scope)

    def _tf_mu_logstd(self, obs_ns):
        unclipped_mu_logstd_n2a = build_mlp(
            obs_ns, scope=self._scope,
            output_size=(env_info.ac_dim() * 2),
            n_layers=flags().sac.learner_depth,
            size=flags().sac.learner_width,
            activation=tf.nn.relu,
            reuse=tf.AUTO_REUSE)
        unclipped_mu_na = unclipped_mu_logstd_n2a[:, :env_info.ac_dim()]
        unclipped_logstd_na = unclipped_mu_logstd_n2a[:, env_info.ac_dim():]

        # What follows is black magic clipping for stability. These constants
        # are luckily not specific to any enviornment, since we are going to
        # squash outputs with tanh anyway. In turn, they're mostly determined
        # by the range of tanh and the machine accuracy of floats. However,
        # the logstd upper bound can be higher in this sense, but this was a
        # magic constant taken from Tuomas' code that I do not have the
        # confidence to mess with.
        #
        # TODO:
        # Consider adding a policy regularization term to encourage staying in
        # the active interval, e.g., add loss max(x - ub, 0) ** 2 or something.
        logstd_na = tf.clip_by_value(unclipped_logstd_na, -20, 2)
        mu_na = tf.clip_by_value(unclipped_mu_na, -5, 5)
        return mu_na, logstd_na

    @staticmethod
    def _tf_sample(mu_na, logstd_na):
        return tf.random_normal(tf.shape(mu_na)) * tf.exp(logstd_na) + mu_na

    @staticmethod
    def _tf_log_prob(x_na, mu_na, logstd_na):
        diffs_na = x_na - mu_na
        quadratic_n = tf.reduce_sum(
            (diffs_na * tf.exp(-logstd_na)) ** 2, axis=1)

        norm_factor_n = 2 * tf.reduce_sum(logstd_na, axis=1)
        norm_factor_n += tf.to_float(env_info.ac_dim()) * np.log(2 * np.pi)
        logprob_n = -0.5 * (quadratic_n + norm_factor_n)
        return logprob_n

    def tf_sample_action_with_log_prob(self, obs_ns):
        return self._tf_sample_action_with_log_prob(
            obs_ns, stop_gradient=True)

    def _tf_sample_action_with_log_prob(self, obs_ns, stop_gradient=True):
        # stop_gradient stops the contribution of the parameters to the
        # sampling directly.
        # when using the reparameterization trick, we *do* want the
        # Jacobian effect of the policy parameters on the sampling, so
        # we would have stop_gradient set to false. However for vanilla
        # policy gradient we don't want this on because it would not
        # correspond to performing a valid reinforce gradient estimation.

        # latent mean, logstd, and log prob
        mu_na, logstd_na = self._tf_mu_logstd(obs_ns)

        # stop the gradient for computing the sample log prob, the
        # parameters of the policy should only affect the log prob
        # gradient by their influence on the policy distribution.
        unbounded_sample_na = self._tf_sample(mu_na, logstd_na)
        if stop_gradient:
            unbounded_sample_na = tf.stop_gradient(unbounded_sample_na)

        unbounded_logprob_n = self._tf_log_prob(
            unbounded_sample_na, mu_na, logstd_na)
        # squash with tanh and scale
        logprob_n = self._scale_correction(
            unbounded_sample_na, unbounded_logprob_n)
        sample_na = self._scale(unbounded_sample_na)
        return sample_na, logprob_n

    def _scale(self, unbounded_acs_na):
        return tf.tanh(unbounded_acs_na)

    @staticmethod
    def _scale_correction(unbounded_acs_na, unbounded_log_prob_n):
        # chain rule to correct for density after modifications
        log_prob_n = unbounded_log_prob_n
        # original paper has 1-tanh**2, we use sech**2 for stability
        # log(sech**2) == 2 * (log 2 - log (e^x + e^-x))
        #              == 2 * (log 2 - |x| - log (1 + e^(-2|x|)))
        abs_na = tf.abs(unbounded_acs_na)
        log_prob_n -= 2 * tf.reduce_sum(
            -1 * abs_na - tf.log1p(tf.exp(-2 * abs_na)) + np.log(2), axis=1)
        # scale to unit action
        log_prob_n -= np.log(0.5) * env_info.ac_dim()
        space = env_info.ac_space()
        # scale to box of action space
        log_prob_n -= np.sum(np.log(space.high - space.low))
        return log_prob_n

    def tf_greedy_action(self, obs_ns):
        mu_na, _ = self._tf_mu_logstd(obs_ns)
        return self._scale(mu_na)

    def act(self, obs_ns):
        return tf.get_default_session().run(self._acs_na, feed_dict={
            self._obs_ph_ns: obs_ns})

    def greedy_act(self, obs_ns):
        return tf.get_default_session().run(self._greedy_acs_na, feed_dict={
            self._obs_ph_ns: obs_ns})

    def tf_report(self, reporter, obs_ns):
        mu_na, logstd_na = self._tf_mu_logstd(obs_ns)
        reporter.stats('logstd', logstd_na)
        reporter.stats('abs(mu)', tf.abs(mu_na))
        acs_na, log_pi_n = self.tf_sample_action_with_log_prob(obs_ns)
        reporter.stats('log pi', log_pi_n)
        reporter.stats('abs(scaled sample acs)', tf.abs(acs_na))

    def expectation(self, obs_ns, fn):
        if flags().sac.reparameterization_trick:
            # for the reparameterization trick, we do not stop the gradient
            # when sampling from the latent distribution, intentionally.
            # This way we can use gradient information of fn to improve the
            # parameters
            onpol_act_na, log_prob_acs_n = (
                self._tf_sample_action_with_log_prob(
                    obs_ns, stop_gradient=False))
            return tf.reduce_mean(fn(onpol_act_na, log_prob_acs_n))
        # else use the reinforce trick
        onpol_act_na, log_prob_acs_n = self._tf_sample_action_with_log_prob(
            obs_ns, stop_gradient=True)
        return tf.reduce_mean(
            log_prob_acs_n * tf.stop_gradient(
                fn(onpol_act_na, log_prob_acs_n)))


#!/usr/bin/env python

"""Convert astronomical positions between various coordinate systems.

The CnvObj code is getting ugly enough (I had to add ICRS<->ICRS2000 cases
to handle the proper motion) that I'm wondering if it wouldn't be better
to just hand code the converter as I did in the TCC. I hate all those
if/then/else statements but it's probably more obvious what's going on
then this mess of magic stuff.

I could also hand code all apparent->ICRS2000 and back routines
and then ditch the entire catalog of methods. That might be cleanest.

History:
2002-08-07 ROwen    Beta release. Not yet well tested and I hope to clean up the code somewhat.
2002-08-23 ROwen    Renamed from Convert; added coordConv method and stopped exporting CnvObj
    (I'm giving up on CnvObj being useful for now, but may revive it later).
    Added velocity handling and default date handling. Still beta; still not well tested.
2002-12-23 ROwen    Bug fix: was not importing ICRSFromFK4 (thanks to pychecker);
    Reorganized the code that computates method dictionaries to make it clearer;
    Reduced the number of globals and made them self-hiding (initial underscore).
2003-05-28 ROwen    Modified to handle opscore.RO.CoordSys 2003-05-09 (defaultDate->currDefaultDate).
2005-01-21 ROwen    Bug fix: was miscomputing proper motion, due to using ICRS
                    instead of ICRS2000 as the intermediate coordinate system.
2005-04-26 ROwen    Bug fix: conversions requiring a App Topo<->Observed step were broken
                    (thanks to Emmanouil Angelakis for the report).
2007-04-24 ROwen    Converted from Numeric to numpy.
2015-09-24 ROwen    Replace "== None" with "is None" to modernize the code.
"""
__all__ = ["coordConv"]

import numpy
import opscore.RO.CoordSys
from opscore.RO.Astro import Tm
from .AppGeoData import AppGeoData
from .FK4FromICRS import fk4FromICRS
from .FK5Prec import fk5Prec
from .GalFromICRS import galFromICRS
from .GeoFromICRS import geoFromICRS
from .GeoFromTopo import geoFromTopo
from .ICRSFromFK4 import icrsFromFK4
from .ICRSFromFixedFK4 import icrsFromFixedFK4
from .ICRSFromGal import icrsFromGal
from .ICRSFromGeo import icrsFromGeo
from .ObsFromTopo import obsFromTopo
from .TopoFromGeo import topoFromGeo
from .TopoFromObs import topoFromObs

_CSysList = (opscore.RO.CoordSys.ICRS, opscore.RO.CoordSys.FK5, opscore.RO.CoordSys.FK4, opscore.RO.CoordSys.Galactic,
    opscore.RO.CoordSys.Geocentric, opscore.RO.CoordSys.Topocentric, opscore.RO.CoordSys.Observed,
    "ICRS2000")

class _CnvObj (object):
    """
    An object that can convert between coordinate systems. To use,
    create one of these objects and then call its "coordConv" method.

    This code is a slightly awkward mix of functional and object-oriented programming.
    Some data is stored as instance variables purely for ease of passing to functions,
    but at present the caching serves no other purpose. In the future I hope to figure out an
    interface whereby the user can save some data and repeatedly ask for conversions.
    For instance one could create an object and then repeatedly ask for its current
    observed position. Then CnvObj would start to feel more like a normal object.

    Also, the method dictionary really should be saved as a class variable, but these
    are relatively new to Python and clumsy, and I'm not yet sure I want to bother
    (especially as it makes the code incompatible with older versions of Python).
    """
    ZeroV = (0.0, 0.0, 0.0)

    def __init__(self):
        # Compute self.methDictDict, a dictionary of _CnvObj method dictionaries
        # to convert between any two coordinate systems:
        # - key: the "from" coordinate system
        # - value: the method dictionary to use if starting from that coordinate system:
        #   - key: a "to" coordinate system
        #   - value: a tuple consisting of:
        #     - the _CnvObj method to convert to "to" from "from" coordinate system
        #     - the "from" coordinate system
        # Each method dictionary gives a chain of methods that stretches from a given
        # starting "from" coordinate system to all other coordinate systems

        def addMethods(methDict, fromSys):
            """Recursively add all methods to a method dictionary
            starting from a given "from" coordinate system.

            For the first call, set methDict = {}

            Inputs:
            - methDict: the method dictionary
            - a "from" coordinate system
            """
            # if first execution, set entry for "from" = "to":
            if methDict == {}:
                methDict[fromSys] = (None, None)

            # add an entry for each "to" coordinate system
            for toSys in _CSysList:
                if toSys == fromSys:
                    continue
                if toSys in methDict:
                    continue
                funcName = "%sFrom%s" % (toSys, fromSys)
                if hasattr(_CnvObj, funcName):
                    methDict[toSys] = (getattr(_CnvObj, funcName), fromSys)
                    #print "methDict[%s] = %r" % (toSys, methDict[toSys])
                    addMethods(methDict, toSys)

        # pre-compute conversion dictionaries for every possible starting coordinate system
        # (which so far includes all in opscore.RO.CoordSys except Physical and Mount)
        methDictDict = {}
        for _whatGiven in _CSysList:
            # define method dictionary for coordsys _whatGiven given
            methDict = {}
            addMethods(methDict, _whatGiven)
            assert len(methDict) == len(_CSysList), "incomplete function list; check your conversion functions"
            methDictDict[_whatGiven] = methDict
            #print "methDictDict[%s] = %r" % (_whatGiven, methDictDict[_whatGiven])

        self.methDictDict = methDictDict

    def coordConv(self, fromP, fromV, fromSys, fromDate, toSys, toDate, obsData=None, refCo=None):
        """Converts a position from one coordinate system to another.
        See notes for the coordConv function below.
        """
        # handle default dates
        if fromDate is None:
            fromDate = opscore.RO.CoordSys.getSysConst(fromSys).currDefaultDate()
        if toDate is None:
            toDate = opscore.RO.CoordSys.getSysConst(toSys).currDefaultDate()

        self.fromDate = fromDate
        self.toDate = toDate
        self.obsData = obsData
        if refCo:
            self.refCo = refCo[:]
        else:
            refCo = None

        # convert to ICRS from fromSys
        icrsP, icrsV = self._getItem("ICRS2000", self.methDictDict[fromSys], fromP, fromV)
        #print "Cnv.coordConv: icrsP=%s, icrsV=%s" % (icrsP, icrsV)

        # convert to toSys from ICRS
        toP, toV = self._getItem(toSys, self.methDictDict["ICRS2000"], icrsP, icrsV)
        #print "Cnv.coordConv: toP=%s, toV=%s" % (toP, toV)
        return (toP, toV)

    def _getItem(self, csys, methDict, initP, initV):
        """Internal routine that returns the position initP and velocity initV
        converted to the requested coordinate system.

        To perform a full conversion, first convert fromP, fromV, fromSys, fromDate to ICRS
        and then convert ICRS to toP, toV, toSys, toDate.

        Before calling, set self.fromDate, self.toDate, self.obsData and self.refCo.

        Inputs:
        - csys      desired coordinate system
        - methDict  _MethodDict[initSys], where initSys is the initial coordinate system
                    (a dictionary of unbound coordinate conversion methods)
        - initP     position in initSys
        - initV     velocity in initSys
        """
        #print "Cnv._CnvObj._getItem: csys=%s; fromP=%s; fromV=%s" % (csys, initP, initV)
        func, nextSys = methDict[csys]
        #print "Cnv._CnvObj._getItem: nextSys=%s; func=%s" % (nextSys, func)
        if func:
            toP, toV = func(self, *self._getItem(nextSys, methDict, initP, initV))
            #print "Cnv._CnvObj._getItem: toP=%s, toV=%s" % (toP, toV)
            return (toP, toV)
        else:
            #print "Cnv._CnvObj._getItem: null conversion; toP, V = fromP, V"
            return (initP, initV)

    # conversion functions
    def ICRSFromICRS2000(self, fromP, fromV):
        return (fromP + (numpy.asarray(fromV, dtype=float) * (self.toDate - 2000.0)), fromV)

    def ICRS2000FromICRS(self, fromP, fromV):
        return (fromP + (numpy.asarray(fromV, dtype=float) * (2000.0 - self.fromDate)), fromV)

    def FK5FromICRS2000(self, fromP, fromV):
        return fk5Prec(fromP, fromV, 2000.0, self.toDate)

    def ICRS2000FromFK5(self, fromP, fromV):
        return fk5Prec(fromP, fromV, self.fromDate, 2000.0)

    def FK4FromICRS2000(self, fromP, fromV):
        return fk4FromICRS (fromP, fromV, self.toDate)

    def ICRS2000FromFK4(self, fromP, fromV):
        if tuple(fromV) == (0.0, 0.0, 0.0):
            return (icrsFromFixedFK4 (fromP, self.fromDate), _CnvObj.ZeroV)
        else:
            return icrsFromFK4(fromP, fromV, self.fromDate)

    def GalacticFromICRS2000(self, fromP, fromV):
        return galFromICRS (fromP, fromV, self.toDate)

    def ICRS2000FromGalactic(self, fromP, fromV):
        return icrsFromGal(fromP, fromV, self.fromDate)

    def GeocentricFromICRS2000(self, fromP, fromV):
        agData = AppGeoData(Tm.epJFromMJD(self.toDate))
        return (geoFromICRS(fromP, fromV, agData), _CnvObj.ZeroV)

    def ICRS2000FromGeocentric(self, fromP, dumV):
        agData = AppGeoData(Tm.epJFromMJD(self.fromDate))
        return (icrsFromGeo(fromP, agData), _CnvObj.ZeroV)

    def TopocentricFromGeocentric(self, fromP, dumV):
        if self.obsData is None:
            raise ValueError("must specify obsData to cnvert to Topocentric from Geocentric")
        return (
            topoFromGeo(fromP, Tm.lastFromUT1(self.toDate, self.obsData.longitude), self.obsData),
            _CnvObj.ZeroV,
        )

    def GeocentricFromTopocentric(self, fromP, dumV):
        if self.obsData is None:
            raise ValueError("must specify obsData to convert to Geocentric from Topocentric")
        return (
            geoFromTopo(fromP, Tm.lastFromUT1(self.fromDate, self.obsData.longitude), self.obsData),
            _CnvObj.ZeroV,
        )

    def ObservedFromTopocentric(self, fromP, dumV):
        if self.refCo is None:
            raise ValueError("must specify refCo to convert to Observed from Topocentric")
        pos, tooLow = obsFromTopo(fromP, self.refCo)
        return (pos, _CnvObj.ZeroV)

    def TopocentricFromObserved(self, fromP, dumV):
        if self.refCo is None:
            raise ValueError("must specify refCo to convert to Topocentric from Observed")
        pos, tooLow = topoFromObs(fromP, self.refCo)
        return (pos, _CnvObj.ZeroV)

# create a singleton _CnvObj for the coordConv method to use
_TheCnvObj = _CnvObj()

# this is the only public function defined in this file
def coordConv(fromP, fromV, fromSys, fromDate, toSys, toDate, obsData=None, refCo=None):
        """Converts a position from one coordinate system to another.

        Inputs:
        - fromP(3)  cartesian position (au)
        - fromV(3)  cartesian velocity (au/year); ignored if fromSys
                    is Geocentric, Topocentric or Observed
        - fromSys   coordinate system from which to convert;
                    any of the entries in the table below; use opscore.RO.CoordSys constants.
        - fromDate  date*
        - toSys     coordinate system to which to convert (see fromSys)
        - toDate    date*
        - obsData   an opscore.RO.Astro.Cnv.ObserverData object; required if fromSys or toSys
                    is Topocentric or Observed; ignored otherwise.
        - refCo(2)  refraction coefficients; required if fromSys or toSys is Observed;
                    ignored otherwise.

        Returns:
        - toP(3)    converted cartesian position (au)
        - toV(3)    converted cartesian velocity (au/year)

        *the units of date depend on the associated coordinate system:
        coord sys   def date    date
        ICRS        2000.0      Julian epoch of observation
        FK5         2000.0      Julian epoch of equinox and observation
        FK4         1950.0      Besselian epoch of equinox and observation
        Galactic     now        Julian epoch of observation
        Geocentric   now        UT1 (MJD)
        Topocentric  now        UT1 (MJD)
        Observed     now        UT1 (MJD)

        **Setting fromV all zero means the object is fixed. This slighly affects
        conversion to or from FK4, which has fictitious proper motion.

        Error Conditions:
        - If obsData or refCo are absent and are required, raises ValueError.

        Details:
        The conversion is performed in two stages:
        - fromP/fromSys/fromDate -> ICRS
        - ICRS -> toP/toSys/toDate

        Each of these two stages is performed using the following graph:

        FK5 ------\
        FK4 ------ ICRS --- Geocentric -*- Topocentric -**- Observed
        Galactic--/

        * obsData required
        ** refCo required
        """
        return _TheCnvObj.coordConv(fromP, fromV, fromSys, fromDate, toSys, toDate, obsData, refCo)


# -*- coding: utf-8 -*-
"""
File bounding_box.py
@author:ZhengYuwei
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

class BoundingBox(object):
    """ 定义BoundingBox类，描述bounding box的坐标及预定义一些操作 """

    def __init__(self, w, h):
        """ box参数初始化
        :param w: box的宽
        :param h: box的高
        """
        self.w = w
        self.h = h
        self.area = w * h

    @staticmethod
    def distance(box, boxes):
        """
        计算一个box与一列表boxes的 1 - IoUs 的静态方法 （作为距离）
        :param box: box对象
        :param boxes：box对象列表
        :return box对象与box对象列表中每一个元素的 1-IOU
        """
        boxes_w, boxes_h, boxes_area = map(np.array, zip(*[(box.w, box.h, box.area) for box in boxes]))
        intersections = np.minimum(box.w, boxes_w) * np.minimum(box.h, boxes_h)
        dist_iou = 1 - intersections / (box.area + boxes_area - intersections)
        return dist_iou

    @staticmethod
    def mean(boxes):
        """
        计算一个box对象列表的均值中心点
        :param boxes: box对象列表
        :return box对象列表的均值中心点box对象
        """
        boxes_w, boxes_h = map(np.array, zip(*[(box.w, box.h) for box in boxes]))
        return BoundingBox(boxes_w.mean(), boxes_h.mean())

    @staticmethod
    def plot(centroids, boxes, group, name='GT bounding prediction cluster'):
        """
        绘制boxes散点图：聚类中心点为黑色，其他每一类一种颜色
        :param centroids: 各个group中心点
        :param boxes: boxes列表
        :param group: boxes对应的归属的中心点
        :param name: 散点图名称
        :return:
        """
        col = ['silver', 'red', 'peru', 'yellow', 'green', 'cyan', 'blue', 'fuchsia', 'pink', 'black']
        # 中心点和boxes的 width 和 height 解析
        center_w, center_h = map(np.array, zip(*[(box.w, box.h) for box in centroids]))
        center_colors = col[:len(centroids)]  # 'k'
        boxes_w, boxes_h = map(np.array, zip(*[(box.w, box.h) for box in boxes]))
        boxes_colors = [col[i] for i in group]
        # 绘制散点图
        plt.grid(ls='--')
        plt.scatter(boxes_w, boxes_h, c=boxes_colors, s=36, alpha=0.3)
        plt.scatter(center_w, center_h, c=center_colors, marker='p', s=48, alpha=1)
        shift = np.max(center_h) * 0.1
        for box in centroids:
            plt.text(box.w, box.h - shift, '({:.3f}, {:.3f})'.format(box.w, box.h))
        plt.title(name)
        plt.xlabel('Width')
        plt.ylabel('Height')
        plt.show()

    @staticmethod
    def plot3D(centroids, boxes, group, cls_names, name='GT bounding prediction cluster'):
        """
        绘制boxes的3维散点图
        :param centroids: 各个group中心点
        :param boxes: boxes列表
        :param group: boxes对应的归属的中心点
        :param cls_names: 物体类别名称
        :param name: 散点图名称
        :return:
        """
        fig = plt.figure()
        ax = Axes3D(fig)
        col = ['black', 'silver', 'red', 'peru', 'yellow', 'green', 'cyan', 'blue', 'fuchsia', 'pink']
        # 中心点和boxes的 width 和 height 解析
        center_w, center_h = map(np.array, zip(*[(box.w, box.h) for box in centroids]))
        center_colors = 'k'  # col[:len(centroids)]
        boxes_w, boxes_h = map(np.array, zip(*[(box.w, box.h) for box in boxes]))
        boxes_colors = [col[i] for i in group]
        # 绘制散点图
        plt.grid(ls='--')
        ax.scatter(boxes_w, boxes_h, [0] * len(boxes_h), c=boxes_colors, alpha=0.3)
        ax.scatter(center_w, center_h, [0] * len(center_w), c=center_colors, marker='p', s=48, alpha=1)
        for index, cls_name in enumerate(cls_names):
            cls_indices = np.where([name == cls_name for name in cls_names])[0]
            ax.scatter(boxes_w[cls_indices], boxes_h[cls_indices], [(index + 1) * 10] * len(cls_indices),
                       c=[boxes_colors[i] for i in cls_indices], alpha=0.3)
            ax.scatter(center_w, center_h, [(index + 1) * 10] * len(center_w),
                       c=center_colors, marker='p', s=48, alpha=1)
        plt.title(name)
        ax.set_xlabel('Width')
        ax.set_ylabel('Height')
        ax.set_zlabel('Classes')
        ax.set_zlim(0, 10 * len(cls_names) + 10)
        plt.savefig("scatter3D.png")
        plt.show()

    @staticmethod
    def plot_pareto(centroids, boxes, group, title='IOU-Ratio Curve'):
        """
        绘制 IOU-probability曲线
        :param centroids: 聚类中心
        :param boxes: box列表
        :param group: boxes对应的归属的中心点
        :param title: 标题
        :return:
        """
        plt.grid(ls='--')
        plt.title('{} Pareto'.format(title))
        col = ['black', 'silver', 'red', 'peru', 'yellow', 'green', 'cyan', 'blue', 'fuchsia', 'pink']
        for i, box in enumerate(centroids):
            pos = np.where(group == i)[0]
            x = 1 - np.sort(BoundingBox.distance(box, [boxes[j] for j in pos]))  # 因为距离是 1 - IOU
            y = np.arange(len(x)) / len(x)
            plt.plot(x, y, color=col[i], label='cluster {}'.format(i))

        plt.legend()
        plt.xlabel('IoU')
        plt.ylabel('Sample ratio')
        plt.show()


# -*- coding: utf-8 -*-

"""
This module does calculations compliant to ISO 3382-1 in order to obtain room
acoustic paramters.

It has an implementation of Lundeby et al. [1] algorithm to estimate the
correction factor for the cumulative integral, as suggested by the ISO 3382-1.

"""

import numpy as np
import matplotlib.pyplot as plt
from numba import njit
from pytta import SignalObj, OctFilter, Analysis, ImpulsiveResponse
from pytta.utils import fractional_octave_frequencies as FOF
import traceback
import copy as cp

def _filter(signal,
            order: int = 4,
            nthOct: int = 3,
            minFreq: float = 20,
            maxFreq: float = 20000,
            refFreq: float = 1000,
            base: int = 10):
    of = OctFilter(order=order,
                   nthOct=nthOct,
                   samplingRate=signal.samplingRate,
                   minFreq=minFreq,
                   maxFreq=maxFreq,
                   refFreq=refFreq,
                   base=base)
    result = of.filter(signal)
    return result[0]

@njit
def _level_profile(timeSignal, samplingRate,
                    numSamples, numChannels, blockSamples=None):
    """
    Gets h(t) in octave bands and do the local time averaging in nblocks.
    Returns h^2_averaged(block).
    """
    def mean_squared(x):
        return np.mean(x**2)

    if blockSamples is None:
        blockSamples = 100
    nblocks = int(numSamples // blockSamples)
    profile = np.zeros((nblocks, numChannels), dtype=np.float32)
    timeStamp = np.zeros((nblocks, 1))

    for ch in range(numChannels):
        # if numChannels == 1:
        #     tmp = timeSignal
        # else:
        tmp = timeSignal[:, ch]
        for idx in range(nblocks):
            profile[idx, ch] = mean_squared(tmp[:blockSamples])
            timeStamp[idx, 0] = idx*blockSamples/samplingRate
            tmp = tmp[blockSamples:]
    return profile, timeStamp

@njit
def _start_sample_ISO3382(timeSignal, threshold) -> np.ndarray:
    squaredIR = timeSignal**2
    # assume the last 10% of the IR is noise, and calculate its noise level
    last10Idx = -int(len(squaredIR)//10)
    noiseLevel = np.mean(squaredIR[last10Idx:])
    # get the maximum of the signal, that is the assumed IR peak
    max_val = np.max(squaredIR)
    max_idx = np.argmax(squaredIR)
    # check if the SNR is enough to assume that the signal is an IR. If not,
    # the signal is probably not an IR, so it starts at sample 1
    idxNoShift = np.asarray([max_val < 100*noiseLevel or
                             max_idx > int(0.9*squaredIR.shape[0])])
    # less than 20dB SNR or in the "noisy" part
    if idxNoShift.any():
        print("noiseLevelCheck: The SNR too bad or this is not an " +
              "impulse response.")
        return 0
    # find the first sample that lies under the given threshold
    threshold = abs(threshold)
    startSample = 1
#    # TODO - envelope mar/pdi - check!
#    if idxNoShift:
#        print("Something wrong!")
#        return
    # if maximum lies on the first point, then there is no point in searching
    # for the beginning of the IR. Just return this position.
    if max_idx > 0:
        abs_dat = 10*np.log10(squaredIR[:max_idx]) \
                  - 10.*np.log10(max_val)
        thresholdNotOk = True
        thresholdShift = 0
        while thresholdNotOk:
            if len(np.where(abs_dat < (-threshold+thresholdShift))[0]) > 0:
                lastBelowThreshold = \
                    np.where(abs_dat < (-threshold+thresholdShift))[0][-1]
                thresholdNotOk = False
            else:
                thresholdShift += 1
        if thresholdShift > 0:
            print("_start_sample_ISO3382: 20 dB threshold too high. " +
                  "Decreasing it.")
        if lastBelowThreshold > 0:
            startSample = lastBelowThreshold
        else:
            startSample = 1
    return startSample

@njit
def _circular_time_shift(timeSignal, threshold=20):
    # find the first sample where inputSignal level > 20 dB or > bgNoise level
    startSample = _start_sample_ISO3382(timeSignal, threshold)
    newTimeSignal = timeSignal[startSample:]
    return (newTimeSignal, startSample)

@njit
def _Lundeby_correction(band, timeSignal, samplingRate, numSamples,
                         numChannels, timeLength):
    returnTuple = (np.float32(0), np.float32(0), np.int32(0), np.float32(0))
    timeSignal, sampleShift = _circular_time_shift(timeSignal)
    if sampleShift is None:
        return returnTuple
    winTimeLength = 0.03  # 30 ms window
    numSamples -= sampleShift  # discount shifted samples
    numParts = 5  # number of parts per 10 dB decay. N = any([3, 10])
    dBtoNoise = 7  # stop point 10 dB above first estimated background noise
    useDynRange = 15  # dynamic range

    # 1) local time average:
    blockSamples = int(winTimeLength * samplingRate)
    timeWinData, timeVecWin = _level_profile(timeSignal, samplingRate,
                                              numSamples, numChannels,
                                              blockSamples)

    # 2) estimate noise from h^2_averaged(block):
    bgNoiseLevel = 10 * \
                   np.log10(
                            np.mean(timeWinData[-int(timeWinData.size/10):]))

    # 3) Calculate premilinar slope
    startIdx = np.argmax(np.abs(timeWinData/np.max(np.abs(timeWinData))))
    stopIdx = startIdx + np.where(10*np.log10(timeWinData[startIdx+1:])
                                  >= bgNoiseLevel + dBtoNoise)[0][-1]
    dynRange = 10*np.log10(timeWinData[stopIdx]) \
        - 10*np.log10(timeWinData[startIdx])
    if (stopIdx == startIdx) or (dynRange > -5)[0]:
        print(band, "[Hz] band: SNR too low for the preliminar slope",
              "calculation.")
        return returnTuple

    # X*c = EDC (energy decaying curve)
    X = np.ones((stopIdx-startIdx, 2), dtype=np.float32)
    X[:, 1] = timeVecWin[startIdx:stopIdx, 0]
    c = np.linalg.lstsq(X, 10*np.log10(timeWinData[startIdx:stopIdx]),
                        rcond=-1)[0]

    if (c[1] == 0)[0] or np.isnan(c).any():
        print(band, "[Hz] band: regression failed. T would be inf.")
        return returnTuple

    # 4) preliminary intersection
    crossingPoint = (bgNoiseLevel - c[0]) / c[1]  # [s]
    if (crossingPoint > 2*(timeLength + sampleShift/samplingRate))[0]:
        print(band, "[Hz] band: preliminary intersection point between",
              "bgNoiseLevel and the decay slope greater than signal length.")
        return returnTuple

    # 5) new local time interval length
    nBlocksInDecay = numParts * dynRange[0] / -10

    dynRangeTime = timeVecWin[stopIdx] - timeVecWin[startIdx]
    blockSamples = int(samplingRate * dynRangeTime[0] / nBlocksInDecay)

    # 6) average
    timeWinData, timeVecWin = _level_profile(timeSignal, samplingRate,
                                              numSamples, numChannels,
                                              blockSamples)

    oldCrossingPoint = 11+crossingPoint  # arbitrary higher value to enter loop
    loopCounter = 0

    while (np.abs(oldCrossingPoint - crossingPoint) > 0.001)[0]:
        # 7) estimate background noise level (BGL)
        bgNoiseMargin = 7
        idxLast10Percent = int(len(timeWinData)-(len(timeWinData)//10))
        bgStartTime = crossingPoint - bgNoiseMargin/c[1]
        if (bgStartTime > timeVecWin[-1:][0])[0]:
            idx10dBDecayBelowCrossPoint = len(timeVecWin)-1
        else:
            idx10dBDecayBelowCrossPoint = \
                np.where(timeVecWin >= bgStartTime)[0][0]
        BGL = np.mean(timeWinData[np.min(
                np.array([idxLast10Percent,
                          idx10dBDecayBelowCrossPoint])):])
        bgNoiseLevel = 10*np.log10(BGL)

        # 8) estimate late decay slope
        stopTime = (bgNoiseLevel + dBtoNoise - c[0])/c[1]
        if (stopTime > timeVecWin[-1])[0]:
            stopIdx = 0
        else:
            stopIdx = int(np.where(timeVecWin >= stopTime)[0][0])

        startTime = (bgNoiseLevel + dBtoNoise + useDynRange - c[0])/c[1]
        if (startTime < timeVecWin[0])[0]:
            startIdx = 0
        else:
            startIdx = int(np.where(timeVecWin <= startTime)[0][0])

        lateDynRange = np.abs(10*np.log10(timeWinData[stopIdx]) \
            - 10*np.log10(timeWinData[startIdx]))

        # where returns empty
        if stopIdx == startIdx or (lateDynRange < useDynRange)[0]:
            print(band, "[Hz] band: SNR for the Lundeby late decay slope too",
                "low. Skipping!")
            # c[1] = np.inf
            c[1] = 0
            break

        X = np.ones((stopIdx-startIdx, 2), dtype=np.float32)
        X[:, 1] = timeVecWin[startIdx:stopIdx, 0]
        c = np.linalg.lstsq(X, 10*np.log10(timeWinData[startIdx:stopIdx]),
                            rcond=-1)[0]

        if (c[1] >= 0)[0]:
            print(band, "[Hz] band: regression did not work, T -> inf.",
                "Setting slope to 0!")
            # c[1] = np.inf
            c[1] = 0
            break

        # 9) find crosspoint
        oldCrossingPoint = crossingPoint
        crossingPoint = (bgNoiseLevel - c[0]) / c[1]

        loopCounter += 1
        if loopCounter > 30:
            print(band, "[Hz] band: more than 30 iterations on regression.",
                "Canceling!")
            break

    interIdx = crossingPoint * samplingRate # [sample]

    return c[0][0], c[1][0], np.int32(interIdx[0]), BGL

@njit
def energy_decay_calculation(band, timeSignal, timeVector, samplingRate,
                             numSamples, numChannels, timeLength, bypassLundeby):
    """Calculate the Energy Decay Curve."""
    if not bypassLundeby:
        lundebyParams = \
            _Lundeby_correction(band,
                                timeSignal,
                                samplingRate,
                                numSamples,
                                numChannels,
                                timeLength)
        _, c1, interIdx, BGL = lundebyParams
        lateRT = -60/c1 if c1 != 0 else 0
    else:
        interIdx = 0
        lateRT = 1

    if interIdx == 0:
        interIdx = -1

    truncatedTimeSignal = timeSignal[:interIdx, 0]
    truncatedTimeVector = timeVector[:interIdx]

    if lateRT != 0.0:
        if not bypassLundeby:
            C = samplingRate*BGL*lateRT/(6*np.log(10))
        else:
            C = 0
        sqrInv = truncatedTimeSignal[::-1]**2
        energyDecayFull = np.cumsum(sqrInv)[::-1] + C
        energyDecay = energyDecayFull/energyDecayFull[0]
    else:
        print(band, "[Hz] band: could not estimate C factor")
        C = 0
        energyDecay = np.zeros(truncatedTimeVector.size)
    return (energyDecay, truncatedTimeVector, lundebyParams)

def cumulative_integration(inputSignal,
                           bypassLundeby,
                           plotLundebyResults,
                           **kwargs):
    """Cumulative integration with proper corrections."""

    def plot_lundeby():
        c0, c1, interIdx, BGL = lundebyParams
        fig = plt.figure(figsize=(10, 5))
        ax = fig.add_axes([0.08, 0.15, 0.75, 0.8], polar=False,
                            projection='rectilinear', xscale='linear')
        line = c1*timeVector + c0
        ax.plot(timeVector, 10*np.log10(timeSignal**2),label='IR')
        ax.axhline(y=10*np.log10(BGL), color='#1f77b4', label='BG Noise')
        ax.plot(timeVector, line,label='Late slope')
        ax.axvline(x=interIdx/samplingRate, label='Truncation point')
        plt.title('{0:.0f} [Hz]'.format(band))
        ax.legend(loc='upper center', shadow=True, fontsize='x-large')

    timeSignal = inputSignal.timeSignal[:]
    # Substituted by SignalObj.crop in analyse function
    # timeSignal, sampleShift = _circular_time_shift(timeSignal)
    # del sampleShift
    hSignal = SignalObj(timeSignal,
                        inputSignal.lengthDomain,
                        inputSignal.samplingRate)
    hSignal = _filter(hSignal, **kwargs)
    bands = FOF(nthOct=kwargs['nthOct'],
                freqRange=[kwargs['minFreq'],kwargs['maxFreq']])[:,1]
    listEDC = []
    for ch in range(hSignal.numChannels):
        signal = hSignal[ch]
        band = bands[ch]
        timeSignal = cp.copy(signal.timeSignal[:])
        timeVector = signal.timeVector[:]
        samplingRate = signal.samplingRate
        numSamples = signal.numSamples
        numChannels = signal.numChannels
        timeLength = signal.timeLength
        energyDecay, energyVector, lundebyParams = \
            energy_decay_calculation(band,
                                     timeSignal,
                                     timeVector,
                                     samplingRate,
                                     numSamples,
                                     numChannels,
                                     timeLength,
                                     bypassLundeby)
        listEDC.append((energyDecay, energyVector))
        if plotLundebyResults:  # Placed here because Numba can't handle plots.
            # plot_lundeby(band, timeVector, timeSignal,  samplingRate,
            #             lundebyParams)
            plot_lundeby()
    return listEDC

@njit
def reverb_time_regression(energyDecay, energyVector, upperLim, lowerLim):
    """Interpolate the EDT to get the reverberation time."""
    if not np.any(energyDecay):
        return 0
    first = np.where(10*np.log10(energyDecay) >= upperLim)[0][-1]
    last = np.where(10*np.log10(energyDecay) >= lowerLim)[0][-1]
    if last <= first:
        # return np.nan
        return 0
    X = np.ones((last-first, 2))
    X[:, 1] = energyVector[first:last]
    c = np.linalg.lstsq(X, 10*np.log10(energyDecay[first:last]), rcond=-1)[0]
    return -60/c[1]

def reverberation_time(decay, nthOct, samplingRate, listEDC):
    """Call the reverberation time regression."""
    try:
        decay = int(decay)
        y1 = -5
        y2 = y1 - decay
    except ValueError:
        if decay in ['EDT', 'edt']:
            y1 = 0
            y2 = -10
        else:
            raise ValueError("Decay must be either 'EDT' or an integer \
                             corresponding to the amount of energy decayed to \
                             evaluate, e.g. (decay='20' | 20).")
    RT = []
    for ED in listEDC:
        edc, edv = ED
        RT.append(reverb_time_regression(edc, edv, y1, y2))
    return RT

def G_Lpe(IR, nthOct, minFreq, maxFreq, IREndManualCut=None):
    """
    Calculate the energy level from the room impulsive response.

    Reference:
        Christensen, C. L.; Rindel, J. H. APPLYING IN-SITU RECALIBRATION FOR
        SOUND STRENGTH MEASUREMENTS IN AUDITORIA.

    :param IR: one channel impulsive response
    :type IR: ImpulsiveResponse

    :param nthOct: number of fractions per octave
    :type nthOct: int

    :param minFreq: analysis inferior frequency limit
    :type minFreq: float

    :param maxFreq: analysis superior frequency limit
    :type maxFreq: float

    :return: Analysis object with the calculated parameter
    :rtype: Analysis
    """
    # Code snippet to guarantee that generated object name is
    # the declared at global scope
    # for frame, line in traceback.walk_stack(None):
    for framenline in traceback.walk_stack(None):
        # varnames = frame.f_code.co_varnames
        varnames = framenline[0].f_code.co_varnames
        if varnames is ():
            break
    # creation_file, creation_line, creation_function, \
    #     creation_text = \
    extracted_text = \
        traceback.extract_stack(framenline[0], 1)[0]
        # traceback.extract_stack(frame, 1)[0]
    # creation_name = creation_text.split("=")[0].strip()
    creation_name = extracted_text[3].split("=")[0].strip()

    # firstChNum = IR.systemSignal.channels.mapping[0]
    # if not IR.systemSignal.channels[firstChNum].calibCheck:
    #     raise ValueError("'IR' must be a calibrated ImpulsiveResponse")
    if isinstance(IR, SignalObj):
        SigObj = cp.copy(IR)
    elif isinstance(IR, ImpulsiveResponse):
        SigObj = cp.copy(IR.systemSignal)
    else:
        raise TypeError("'IR' must be an ImpulsiveResponse or SignalObj.")
    # Cutting the IR
    if IREndManualCut is not None:
        SigObj.crop(0, IREndManualCut)
    timeSignal, _ = _circular_time_shift(SigObj.timeSignal[:,0])
    # Bands filtering
    # hSignal = SignalObj(SigObj.timeSignal[:,0],
    hSignal = SignalObj(timeSignal,
                        SigObj.lengthDomain,
                        SigObj.samplingRate)
    hSignal = _filter(signal=hSignal, nthOct=nthOct, minFreq=minFreq,
                      maxFreq=maxFreq)
    bands = FOF(nthOct=nthOct,
                freqRange=[minFreq,maxFreq])[:,1]
    Lpe = []
    for chIndex in range(hSignal.numChannels):
        Lpe.append(
            10*np.log10(np.trapz(y=hSignal.timeSignal[:,chIndex]**2/(2e-5**2),
                                 x=hSignal.timeVector)))
    LpeAnal = Analysis(anType='mixed', nthOct=nthOct, minBand=float(bands[0]),
                       maxBand=float(bands[-1]), data=Lpe,
                       comment='h**2 energy level')
    LpeAnal.creation_name = creation_name
    return LpeAnal

def G_Lps(IR, nthOct, minFreq, maxFreq):
    """G_Lps

    Calculates the recalibration level, for both in-situ and
    reverberation chamber. Lps is applied for G calculation.

    During the recalibration: source height and mic heigth must be >= 1 [m],
    while the distance between source and mic must be <= 1 [m]. The distances
    must be the same for in-situ and reverberation chamber measurements.

    Reference:
        Christensen, C. L.; Rindel, J. H. APPLYING IN-SITU RECALIBRATION FOR
        SOUND STRENGTH MEASUREMENTS IN AUDITORIA.

    :param IR: one channel impulsive response
    :type IR: ImpulsiveResponse

    :param nthOct: number of fractions per octave
    :type nthOct: int

    :param minFreq: analysis inferior frequency limit
    :type minFreq: float

    :param maxFreq: analysis superior frequency limit
    :type maxFreq: float

    :return: Analysis object with the calculated parameter
    :rtype: Analysis
    """
    # Code snippet to guarantee that generated object name is
    # the declared at global scope
    # for frame, line in traceback.walk_stack(None):
    for framenline in traceback.walk_stack(None):
        # varnames = frame.f_code.co_varnames
        varnames = framenline[0].f_code.co_varnames
        if varnames is ():
            break
    # creation_file, creation_line, creation_function, \
    #     creation_text = \
    extracted_text = \
        traceback.extract_stack(framenline[0], 1)[0]
        # traceback.extract_stack(frame, 1)[0]
    # creation_name = creation_text.split("=")[0].strip()
    creation_name = extracted_text[3].split("=")[0].strip()

    # firstChNum = IR.systemSignal.channels.mapping[0]
    # if not IR.systemSignal.channels[firstChNum].calibCheck:
    #     raise ValueError("'IR' must be a calibrated ImpulsiveResponse")
    if isinstance(IR, SignalObj):
        SigObj = IR
    elif isinstance(IR, ImpulsiveResponse):
        SigObj = IR.systemSignal
    else:
        raise TypeError("'IR' must be an ImpulsiveResponse or SignalObj.")
    # Windowing the IR
    # dBtoOnSet = 20
    # dBIR = 10*np.log10((SigObj.timeSignal[:,0]**2)/((2e-5)**2))
    # windowStart = np.where(dBIR > (max(dBIR) - dBtoOnSet))[0][0]

    broadBandTimeSignal = cp.copy(SigObj.timeSignal[:,0])
    broadBandTimeSignalNoStart, sampleShift = \
        _circular_time_shift(broadBandTimeSignal)
    windowLength = 0.0032 # [s]
    windowEnd = int(windowLength*SigObj.samplingRate)

    hSignal = SignalObj(broadBandTimeSignalNoStart[:windowEnd],
    # hSignal = SignalObj(timeSignal,
                        SigObj.lengthDomain,
                        SigObj.samplingRate)
    hSignal = _filter(signal=hSignal, nthOct=nthOct, minFreq=minFreq,
                      maxFreq=maxFreq)
    bands = FOF(nthOct=nthOct,
                freqRange=[minFreq,maxFreq])[:,1]
    Lps = []
    for chIndex in range(hSignal.numChannels):
        timeSignal = cp.copy(hSignal.timeSignal[:,chIndex])
        # timeSignalNoStart, sampleShift = _circular_time_shift(timeSignal)
        # windowLength = 0.0032 # [s]
        # windowEnd = int(windowLength*SigObj.samplingRate)

        Lps.append(
            # 10*np.log10(np.trapz(y=timeSignalNoStart[:windowEnd]**2/(2e-5**2),
            10*np.log10(np.trapz(y=timeSignal**2/(2e-5**2),
                                #  x=hSignal.timeVector[sampleShift:sampleShift+windowEnd])))
                                 x=hSignal.timeVector)))
    LpsAnal = Analysis(anType='mixed', nthOct=nthOct, minBand=float(bands[0]),
                       maxBand=float(bands[-1]), data=Lps,
                       comment='Source recalibration method IR')
    LpsAnal.creation_name = creation_name
    LpsAnal.windowLimits = ((sampleShift)/SigObj.samplingRate,
                            (sampleShift+windowEnd)/SigObj.samplingRate)
    # Plot IR cutting
    # fig = plt.figure(figsize=(10, 5))
    # ax = fig.add_axes([0.08, 0.15, 0.75, 0.8], polar=False,
    #                         projection='rectilinear', xscale='linear')
    # ax.plot(SigObj.timeVector, 10*np.log10(SigObj.timeSignal**2/2e-5**2))
    # ax.axvline(x=(sampleShift)/SigObj.samplingRate, linewidth=4, color='k')
    # ax.axvline(x=(sampleShift+windowEnd)/SigObj.samplingRate, linewidth=4, color='k')
    # ax.set_xlim([(sampleShift-100)/SigObj.samplingRate, (sampleShift+windowEnd+100)/SigObj.samplingRate])
    return LpsAnal

def strength_factor(Lpe, Lpe_revCh, V_revCh, T_revCh, Lps_revCh, Lps_inSitu):
    S0 = 1 # [m2]

    bands = T_revCh.bands
    nthOct = T_revCh.nthOct
    terms = []
    for bandData in T_revCh.data:
        if bandData == 0:
            terms.append(0)
        else:
            term = (V_revCh * 0.16) / (bandData * S0)
            terms.append(term)
    terms = [10*np.log10(term) if term != 0 else 0 for term in terms]

    revChTerm = Analysis(anType='mixed', nthOct=nthOct, minBand=float(bands[0]),
                         maxBand=float(bands[-1]), data=terms)
    Lpe.anType = 'mixed'
    Lpe_revCh.anType = 'mixed'
    Lps_revCh.anType = 'mixed'
    Lps_inSitu.anType = 'mixed'
    G = Lpe - Lpe_revCh - revChTerm + 37 \
        + Lps_revCh - Lps_inSitu
    G.anType = 'G'
    return G

def _clarity(temp, signalObj, nthOct, **kwargs):  # TODO
    """

    """
#    try:
#        temp = int(temp)*signalObj.samplingRate//1000
#    except ValueError:
#        raise ValueError("The temp parameter must be an integer or a string \
#                         of integers, e.g. (temp='80' | 80).")
#    output = []
#    for ch in range(signalObj.num_channels()):
#        filtResp = filtered_response(signalObj[ch], nthOct, **kwargs)
#        C = []
#        for bd in range(len(filtResp)):
#            C.append(round(np.sum(filtResp[bd][:temp], axis=0)
#                           / np.sum(filtResp[bd][temp:], axis=0)[0], 2))
#        output.append(C)
#    return output
    pass

def _definition(temp, signalObj, nthOct, **kwargs):  # TODO
    """

    """
#    try:
#        temp = int(temp)*signalObj.samplingRate//1000
#    except ValueError:
#        raise ValueError("The temp parameter must be an integer or a string \
#                         of integers, e.g. (temp='50' | 50).")
#    output = []
#    for ch in range(signalObj.num_channels()):
#        filtResp = filtered_response(signalObj[ch], nthOct, **kwargs)
#        D = []
#        for bd in range(len(filtResp)):
#            D.append(round(10*np.log10(
#                        np.sum(filtResp[bd][:temp], axis=0)
#                        / np.sum(filtResp[bd][:], axis=0))[0], 2))
#        output.append(D)
#    return output
    pass

def crop_IR(SigObj, IREndManualCut):
    """Cut the impulse response at background noise level."""
    timeSignal = cp.copy(SigObj.timeSignal)
    timeVector = SigObj.timeVector
    samplingRate = SigObj.samplingRate
    numSamples = SigObj.numSamples
    # numChannels = SigObj.numChannels
    if SigObj.numChannels > 1:
        print('crop_IR: The provided impulsive response has more than one ' +
              'channel. Cropping based on channel 1.')
    numChannels = 1
    # Cut the end automatically or manual
    if IREndManualCut is None:
        winTimeLength = 0.1  # [s]
        meanSize = 5  # [blocks]
        dBtoReplica = 6  # [dB]
        blockSamples = int(winTimeLength * samplingRate)
        timeWinData, timeVecWin = _level_profile(timeSignal, samplingRate,
                                                numSamples, numChannels,
                                                blockSamples)
        endTimeCut = timeVector[-1]
        for blockIdx, blockAmplitude in enumerate(timeWinData):
            if blockIdx >= meanSize:
                anteriorMean = 10*np.log10( \
                    np.sum(timeWinData[blockIdx-meanSize:blockIdx])/meanSize)
                if 10*np.log10(blockAmplitude) > anteriorMean+dBtoReplica:
                    endTimeCut = timeVecWin[blockIdx-meanSize//2]
                    break
    else:
        endTimeCut = IREndManualCut
    endTimeCutIdx = np.where(timeVector >= endTimeCut)[0][0]
    timeSignal = timeSignal[:endTimeCutIdx]
    # Cut the start automatically
    timeSignal, _ = _circular_time_shift(timeSignal)
    result = SignalObj(timeSignal,
                       'time',
                       samplingRate,
                       signalType='energy')
    return result

def analyse(obj, *params,
            bypassLundeby=False,
            plotLundebyResults=False,
            IREndManualCut=None, **kwargs):
    """
    Receives an one channel SignalObj or ImpulsiveResponse and calculate the
    room acoustic parameters especified in the positional input arguments.

    :param obj: one channel impulsive response
    :type obj: SignalObj or ImpulsiveResponse

    Input parameters for reverberation time, 'RT':
        :param RTdecay: decay interval for RT calculation. e.g. 20
        :type RTdecay: int

    Input parameters for clarity, 'C':
        TODO

    Input parameters for definition, 'D':
        TODO

    Input parameters for strength factor, 'G':
        TODO

    :param nthOct: number of fractions per octave
    :type nthOct: int

    :param minFreq: analysis inferior frequency limit
    :type minFreq: float

    :param maxFreq: analysis superior frequency limit
    :type maxFreq: float

    :param bypassLundeby: bypass lundeby correction
    to False
    :type bypassLundeby: bool, optional

    :param plotLundebyResults: plot the Lundeby correction parameters, defaults to False
    :type plotLundebyResults: bool, optional

    :return: Analysis object with the calculated parameter
    :rtype: Analysis

    """
    # Code snippet to guarantee that generated object name is
    # the declared at global scope
    # for frame, line in traceback.walk_stack(None):
    for framenline in traceback.walk_stack(None):
        # varnames = frame.f_code.co_varnames
        varnames = framenline[0].f_code.co_varnames
        if varnames is ():
            break
    # creation_file, creation_line, creation_function, \
    #     creation_text = \
    extracted_text = \
        traceback.extract_stack(framenline[0], 1)[0]
    # traceback.extract_stack(frame, 1)[0]
    # creation_name = creation_text.split("=")[0].strip()
    creation_name = extracted_text[3].split("=")[0].strip()

    if not isinstance(obj, SignalObj) and not isinstance(obj, ImpulsiveResponse):
        raise TypeError("'obj' must be an one channel SignalObj or" +
                        " ImpulsiveResponse.")
    if isinstance(obj, ImpulsiveResponse):
        SigObj = obj.systemSignal
    else:
        SigObj = obj

    if SigObj.numChannels > 1:
        raise TypeError("'obj' can't contain more than one channel.")
    samplingRate = SigObj.samplingRate

    SigObj = crop_IR(SigObj, IREndManualCut)

    listEDC = cumulative_integration(SigObj,
                                     bypassLundeby,
                                     plotLundebyResults,
                                     **kwargs)
    for _ in params:
        if 'RT' in params:
            RTdecay = params[params.index('RT')+1]
            nthOct = kwargs['nthOct']
            RT = reverberation_time(RTdecay, nthOct, samplingRate, listEDC)
            result = Analysis(anType='RT', nthOct=nthOct,
                              minBand=kwargs['minFreq'],
                              maxBand=kwargs['maxFreq'],
                              data=RT)
        # if 'C' in prm:
        #     Ctemp = prm[1]
        # if 'D' in prm:
        #     Dtemp = prm[1]
    result.creation_name = creation_name
    return result