#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
Reddening laws.
"""
import pylab as py
import numpy as np
from scipy import interpolate
import pysynphot
from scipy.linalg import solve_banded
import pdb

def get_red_law(str):
    """
    Given a reddening law name, return the reddening
    law object.

    Parameters:
    ----------
    str: str
        Reddening law name and additional params (comma-separated).
        Name must match
    """
    # Parse the string, extracting redlaw name and other params
    tmp = str.split(',')
    name = tmp[0]
    params = ()
    if len(tmp) > 1:
        for ii in range(len(tmp) - 1):
            params = params + (float(tmp[ii+1]),)

    # Define dictionary connecting redlaw names to the redlaw classes
    name_dict = {'N09':RedLawNishiyama09,
                     'C89': RedLawCardelli,
                     'RZ07': RedLawRomanZuniga07,
                     'RL85': RedLawRiekeLebofsky,
                     'D16': RedLawDamineli16,
                     'DM16': RedLawDeMarchi16,
                     'F09': RedLawFitzpatrick09,
                     'S16': RedLawSchlafly16,
                     'pl': RedLawPowerLaw,
                     'F11': RedLawFritz11,
                     'H18': RedLawHosek18,
                     'H18b': RedLawHosek18b,
                     'NL18': RedLawNoguerasLara18}

    # Make reddening law object, including params if necessary.
    # This is not great coding, but I really strugged to generalize this...
    if len(params) == 0:
        red_law = name_dict[name]()
    elif len(params) == 1:
        red_law = name_dict[name](params[0])
    elif len(params) == 2:
        red_law = name_dict[name](params[0], params[1])
    elif len(params) == 3:
        red_law = name_dict[name](params[0], params[1], params[2])
    elif len(params) == 4:
        red_law = name_dict[name](params[0], params[1], params[2], params[3])
    else:
        mes = 'Redlaw contains more params than reddening.get_red_law currently supports'
        raise ValueError(mes)

    return red_law

class RedLawNishiyama09(pysynphot.reddening.CustomRedLaw):
    """
    Defines extinction law from `Nishiyama et al. 2009
    <https://ui.adsabs.harvard.edu/abs/2009ApJ...696.1407N/abstract>`_
    toward the Galactic Center. This is the default extinction law.
    The law is defined between 0.5 -- 8 microns.
    """
    def __init__(self):
        # Fetch the extinction curve, pre-interpolate across 3-8 microns
        wave = np.arange(0.5, 8.0, 0.001)

        # This will eventually be scaled by AKs when you
        # call reddening(). Right now, calc for AKs=1
        wave_vals, Alambda_scaled = RedLawNishiyama09._derive_nishiyama09(wave)

        # Convert wavelength to angstrom
        wave_vals *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave_vals,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Nishiyama09',
                                                  litref='Nishiyama+ 2009')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave_vals)
        self.high_lim = max(wave_vals)
        self.name = 'N09'

    @staticmethod
    def _derive_nishiyama09(wavelength):
        """
        Calculate the N09 extinction law as defined in the paper:
        a A_lambda/AKs = power law of exponent -2.0 between JHK. Then
        use a *linear* interpolation in 1/lambda space to go from J to the V-band observation,
        in order to avoid imposing more structure. A cublic spline interpolation
        across the wavelength points is used longward of K-band

        Parameters
        ----------
        wavelength : float
            in microns
        AKs : float
            in magnitudes
        """
        #-----Define power law extinction law between JHK----#
        jhk_idx = np.where( (wavelength >= 1.25) & (wavelength <= 2.14) )

        alpha = 2.0
        wave_jhk = wavelength[jhk_idx]

        A_jhk = wave_jhk**(-1.0*alpha)
        A_Ks_jhk = A_jhk / A_jhk[-1]

        #----Now do a linear interpolation (in log(1/lambda) vs log(A/AKs) space) between 1.25 microns and 0.551 microns---#
        jv_idx = np.where( (wavelength < 1.25) & (wavelength > 0.551) )
        Av = 16.13
        func = interpolate.interp1d(np.log10(np.array([1.0/1.25, 1.0/0.551])), np.log10(np.array([A_Ks_jhk[0], Av])),
                                        kind='linear')
        A_Ks_jv = func(np.log10(1.0 / wavelength[jv_idx]))

        # Convert back to linear space
        A_Ks_jv = 10**A_Ks_jv

        #---Do a spline interpolation for the rest of the (long-wavelength) law---#
        # We do this since no other function form is given
        long_idx = np.where(wavelength > 2.14)
        wave = np.array([0.551, 1.25, 1.63, 2.14, 3.545, 4.442, 5.675, 7.760])
        A_AKs = np.array([16.13, 3.02, 1.73, 1.00, 0.500, 0.390, 0.360, 0.430])

        spline_interp = interpolate.splrep(wave, A_AKs, k=3, s=0)
        A_AKs_long = interpolate.splev(wavelength[long_idx], spline_interp)

        # Stitch together sections for the final law
        wave_vals = np.concatenate((wavelength[jv_idx[0]], wavelength[jhk_idx[0]]))
        A_AKs_vjhk = np.concatenate((A_Ks_jv, A_Ks_jhk))

        # Now add the long-wavelength law
        wave_vals = np.concatenate((wave_vals, wavelength[long_idx[0]]))
        A_AKs_final = np.concatenate((A_AKs_vjhk, A_AKs_long))

        return wave_vals, A_AKs_final

    def Nishiyama09(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawCardelli(pysynphot.reddening.CustomRedLaw):
    """
    Defines the extinction law from
    `Cardelli et al. 1989 <https://ui.adsabs.harvard.edu/abs/1989ApJ...345..245C/abstract>`_.
    The law is defined from 0.3 - 3 microns, and in terms
    of :math:`A_{\lambda} / A_{Ks}`, where Ks is 2.174 microns.

    Parameters
    ----------
    Rv : float
        Ratio of absolute to selective extinction, :math:`A(V) / E(B-V)`.
        The standard value for the diffuse ISM is 3.1.
    """
    def __init__(self, Rv):
        # Fetch the extinction curve, pre-interpolate across 0.3-3 microns
        wave = np.arange(0.3, 3.0, 0.001)

        # This will eventually be scaled by AKs when you
        # call reddening(). Produces A_lambda for AKs = 1, which will be
        # scaled later. Expects wavelength in microns
        Alambda_scaled = RedLawCardelli._derive_cardelli(wave, Rv)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Cardelli89',
                                                  litref='Cardelli+ 2009')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'C89,{0}'.format(Rv)

    @staticmethod
    def _derive_cardelli(wavelength, Rv):
        """
        Cardelli extinction law. This produces extinction values expected
        for AKs = 1
        """
        x = 1.0 / np.array(wavelength)

        # check for applicability
        if (np.min(x) < 0.3):
            print( 'wavelength is longer than applicable range for Cardelli law')
            return None

        if (np.max(x) > 8.0):
            print( 'wavelength is shorter than applicable range for Cardelli law')
            return None

        # Set up some arrays for coefficients that we will need
        a = np.zeros(len(x), dtype=float)
        b = np.zeros(len(x), dtype=float)

        y = x - 1.82

        # Calculate coefficients for long wavelengths (low wavenumber)
        # Wavenumger <= 1.1 (Eq. 2a, 2b)
        idx = np.where(x <= 1.1)[0]
        a[idx] =  0.574 * x[idx] ** 1.61
        b[idx] = -0.527 * x[idx] ** 1.61

        # Calculate coefficients for intermediate wavelengths
        # 1.1 < wavenumber <= 3.3 (Eq. 3a, 3b)
        idx = np.where((x > 1.1) & (x <= 3.3))[0]
        yy = y[idx]
        a[idx] = 1 + (0.17699 * yy) - (0.50447 * yy ** 2) - \
            (0.02427 * yy ** 3) + (0.72085 * yy ** 4) + \
            (0.01979 * yy ** 5) - (0.77530 * yy ** 6) + \
            (0.32999 * yy ** 7)
        b[idx] = (1.41338 * yy) + (2.28305 * yy ** 2) + \
            (1.07233 * yy ** 3) - (5.38434 * yy ** 4) - \
            (0.62251 * yy ** 5) + (5.30260 * yy ** 6) - \
            (2.09002 * yy ** 7)

        # Calculate the long wavelength
        # 3.3 < wavenumber < 5.9 (Eq. 4a, 4b)
        idx = np.where((x > 3.3) & (x < 5.9))[0]
        xx = x[idx]
        a[idx] = 1.752 - (0.316 * xx) - (0.104/((xx - 4.67) ** 2 + 0.341))
        b[idx] = -3.090 + (1.825 * xx) + (1.206/((xx - 4.62) ** 2 + 0.263))

        # Calculate the longest wavelength
        # 5.9 <= wavenumber (Eq. 4a, 4b)
        idx = np.where(x >= 5.9)[0]
        xx = x[idx]
        a[idx] = 1.752 - (0.316 * xx) - (0.104/((xx - 4.67) ** 2 + 0.341)) + \
            (-0.04473 * (xx - 5.9) ** 2) - (0.009779 * (xx - 5.9) ** 3)
        b[idx] = -3.090 + (1.825 * xx) + (1.206/((xx - 4.62) ** 2 + 0.263)) + \
            (0.2130 * (xx - 5.9) ** 2) + (0.1207 * (xx - 5.9) ** 3)

        # A(lam) / A(V), from Eq. 1
        extinction = a + b/Rv

        # Now, want to produce A_lambda / AKs, to match other laws
        k_ind = np.where(abs(x-0.46) == min(abs(x-0.46)))
        Aks_Av = a[k_ind] + b[k_ind]/Rv # Aks / Av
        Av_Aks = 1.0 / Aks_Av # Av / Aks

        output = extinction * Av_Aks # (A(lamb) / Av) * (Av / Aks) = (A(lamb) / Aks)

        return output

    def Cardelli89(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawRomanZuniga07(pysynphot.reddening.CustomRedLaw):
    """
    Defines extinction law from `Roman-Zuniga et al. 2007
    <https://ui.adsabs.harvard.edu/abs/2007ApJ...664..357R/abstract>`_
    for the dense cloud core Barnard 59. It is defined between 1.0 - 8.0
    microns.
    """
    def __init__(self):
        # Fetch the extinction curve, pre-interpolate across 1-8 microns
        wave = np.arange(1.0, 8.0, 0.01)

        # This will eventually be scaled by AKs when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawRomanZuniga07._derive_romanzuniga07(wave)

        # Convert wavelength to angstrom
        wave *= 10**4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='RomanZuniga07',
                                                  litref='Roman-Zuniga+ 2007')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'RZ07'

    @staticmethod
    def _derive_romanzuniga07(wavelength):
        filters = ['J', 'H', 'Ks', '[3.6]', '[4.5]', '[5.8]', '[8.0]']
        wave =      np.array([1.240, 1.664, 2.164, 3.545, 4.442, 5.675, 7.760])
        A_AKs =     np.array([2.299, 1.550, 1.000, 0.618, 0.525, 0.462, 0.455])
        A_AKs_err = np.array([0.530, 0.080, 0.000, 0.077, 0.063, 0.055, 0.059])

        # Interpolate over the curve
        spline_interp = interpolate.splrep(wave, A_AKs, k=3, s=0)
        A_AKs_at_wave = interpolate.splev(wavelength, spline_interp)

        return A_AKs_at_wave

    def RomanZuniga07(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawRiekeLebofsky(pysynphot.reddening.CustomRedLaw):
    """
    Defines the extinction law from `Rieke & Lebofsky 1985
    <https://ui.adsabs.harvard.edu/abs/1985ApJ...288..618R/abstract>`_
    for the Galactic Center. The law is defined between 1.0 - 13 microns.
    """
    def __init__(self):
        # Fetch the extinction curve, pre-interpolate across 0.365-13 microns
        wave = np.arange(0.365, 13.0, 0.001)

        # This will eventually be scaled by AKs when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawRiekeLebofsky._derive_RiekeLebofsky(wave)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='RiekeLebofsky',
                                                  litref='Rieke+Lebovsky 1985')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'RL85'

    @staticmethod
    def _derive_RiekeLebofsky(wavelength):
        """
        Calculate the resulting extinction for an array of wavelengths.
        The extinction is normalized with A_Ks.

        Data pulled from Rieke+Lebofsky 1985, Table 3
        """
        filters = ['U', 'B', 'V', 'R', 'I', 'J', 'H', 'K', 'L', 'M',
                   '[8.0]', '[8.5]', '[9.0]', '[9.5]', '[10.0]', '[10.5]',
                   '[11.0]', '[11.5]', '[12.0]', '[12.5]', '[13.0]']
        #wave = np.array([0.365, 0.445, 0.551, 0.658, 0.806, 1.25, 1.635, 2.2,
        #                 3.77, 4.68, 4.75, 8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0,
        #                11.5, 12.0, 12.5, 13.0])

        # Wavelengths from Nishiyama+09 plot of RL+85 law...slightly different than standard,
        # drop N filter
        wave = np.array([0.365, 0.445, 0.551, 0.658, 0.806, 1.17, 1.57, 2.12,
                         3.40, 4.75, 8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0,
                        11.5, 12.0, 12.5, 13.0])
        A_Av = np.array([1.531, 1.324, 1.00, 0.748, 0.482, 0.282, 0.175, 0.112,
                         0.058, 0.023, 0.02, 0.043, 0.074, 0.087, 0.083,
                         0.074, 0.060, 0.047, 0.037, 0.030, 0.027])
        # Want to change this from A/Av to A/AK
        k_ind = np.where(np.array(filters) == 'K')
        Ak_Av = A_Av[k_ind]
        Av_Ak = 1.0 / Ak_Av

        A_Ak = A_Av * Av_Ak

        # Interpolate over the curve
        spline_interp = interpolate.splrep(wave, A_Ak, k=3, s=0)
        A_Ak_at_wave = interpolate.splev(wavelength, spline_interp)

        return A_Ak_at_wave

    def RiekeLebofsky85(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawDamineli16(pysynphot.reddening.CustomRedLaw):
    """
    Defines the extinction law of `Damineli et al. 2016
    <https://ui.adsabs.harvard.edu/abs/2016MNRAS.463.2653D/abstract>`_,
    derived for the Wd1 cluster. The law is derived between
    0.5 - 8.0 microns.
    """
    def __init__(self):
        # Fetch the extinction curve, pre-interpolate across 1-8 microns
        wave = np.arange(0.3, 8.0, 0.001)

        # This will eventually be scaled by AKs when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawDamineli16._derive_Damineli16(wave)
        #Alambda_scaled = RedLawDamineli16.derive_Damineli16_old(wave, 1.0)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Damineli16',
                                                  litref='Damineli+ 2016')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'D16'

    @staticmethod
    def _derive_Damineli16(wavelength):
        """
        Calculate the Damineli+16 extinction law using their equation 19

        Parameters
        ----------
        wavelength : float
            in microns
        AKs : float
            in magnitudes
        """
        # From their eq 19
        x = np.log10(2.159 / wavelength)
        log_A_AKs = -0.015 + 2.33*x + 0.522*x**2. - 3.001*x**3. + 2.034*x**4.

        # Now to convert this back to linear space
        A_AKs_at_wave = 10**log_A_AKs

        return A_AKs_at_wave

    def Damineli16(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawDeMarchi16(pysynphot.reddening.CustomRedLaw):
    """
    Defines extinction law from `De Marchi et al. 2016
    <https://ui.adsabs.harvard.edu/abs/2016MNRAS.455.4373D/abstract>`_
    derived for 30 Dorodus. The law is defined between 0.3 - 8.0 microns.
    """
    def __init__(self):
        # Fetch the extinction curve, pre-interpolate across 1-8 microns
        wave = np.arange(0.3, 8.0, 0.001)

        # This will eventually be scaled by AK when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawDeMarchi16._derive_DeMarchi16(wave)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='DeMarchi16',
                                                  litref='DeMarchi+ 2016')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'DM16'

    @staticmethod
    def _derive_DeMarchi16(wavelength):
        """
        Calculate the resulting extinction for an array of wavelengths.
        The extinction is normalized with A_Ks.

        Data pulled from DeMarchi+16, Table 3

        Note: Authors measure R_VI (V) = 3.09 +/- 0.15,
        so we use this to calculate the absolute extinctions in all
        of the other bands. This corresponds to A_I/ A_V = 0.676

        Note that they extrapolate their curve to get to K-band

        Parameters
        ----------
        wavelength : float
            in microns
        AKs : float
            in magnitudes
        """
        AI_AV = 0.676

        # Extracting the values from the paper
        filters = ['U', 'B', 'V', 'R', 'I', 'J', 'H', 'K']
        wave = np.array([0.365, 0.445, 0.551, 0.658, 0.806, 1.22, 1.63, 2.19])
        R_VI = np.array([4.41, 3.78, 3.09, 2.58, 2.09, 1.26, 0.84, 0.52])
        R_VI_err = np.array([0.18, 0.15, 0.15, 0.13, 0.17, 0.18, 0.12, 0.08])

        # We'll calculate A_AKs from R_VI
        A_Av = R_VI * (1. - AI_AV)
        AK_Av = A_Av[-1]
        A_AK = A_Av / AK_Av

        # Interpolate over the curve
        spline_interp = interpolate.splrep(wave, A_AK, k=3, s=0)
        A_AK_at_wave = interpolate.splev(wavelength, spline_interp)

        return A_AK_at_wave

    def DeMarchi16(self, wavelength, AK):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AK (since law assumes AK = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AK

        return A_at_wave

class RedLawFitzpatrick09(pysynphot.reddening.CustomRedLaw):
    """
    Defines the extinction law from
    `Fitzpatrick et al. 2009 <https://ui.adsabs.harvard.edu/abs/2009ApJ...699.1209F/abstract>`_.
    The law is defined between 0.3 -- 3 microns.

    The extinction law is as defined in their equation 5, and has two
    free parameters: :math:`\alpha` and R(V). Averaged over 14 sight-lines,
    the authors generally find either :math:`alpha` ~ 2.5, R(V) ~ 3, or
    :math:`alpha` ~ 1.8, R(V) ~ 5 (their Figure 6).

    Parameters
    ----------
    alpha : float
         alpha parameter for extinction law.

    RV : float
        R(V) parameter for extinction law.
    """
    def __init__(self, alpha, RV):
        # Fetch the extinction curve, pre-interpolate across 1-8 microns
        wave = np.arange(0.7, 3.0, 0.001)

        # This will eventually be scaled by AK when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawFitzpatrick09._derive_Fitzpatrick09(wave, alpha, RV)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Fitzpatrick09',
                                                  litref='Fitzpatrick+ 2009')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'F09,{0},{1}'.format(alpha, RV)

    @staticmethod
    def _derive_Fitzpatrick09(wavelength, alpha, RV):
        """
        Calculate the resulting extinction for an array of wavelengths.
        The extinction is normalized with A_Ks.

        Data pulled from Fitzpactrick09, equation 5

        Parameters
        ----------
        wavelength : float
            in microns

        alpha: float
            Free parameter alpha

        RV: float
            Free parameter RV
        """
        alpha = float(alpha)
        RV = float(RV)

        # First we'll calculate k(lambda - V) = E(lambda - V) / E(B - V),
        # directly from equation 5
        k = (0.349 + 2.087*RV) * (1.0 / (1.0 + (wavelength / 0.507)**alpha)) - RV

        # We'll calculate Alam/Av from K + Rv
        Alam_Av = (k / RV) + 1.

        # Finally, to get A_lambda/Aks we need to divide Alam_Av by AKs_Av.
        # We'll assume central wavelength of 2.14 for Ks
        idx = np.where(abs(wavelength - 2.14) == min(abs(wavelength - 2.14)))

        A_AKs_at_wave = Alam_Av / Alam_Av[idx]

        return A_AKs_at_wave

    def Fitzpatrick09(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawSchlafly16(pysynphot.reddening.CustomRedLaw):
    """
    Defines the extinction law from `Schlafly et al. 2016
    <https://ui.adsabs.harvard.edu/abs/2016ApJ...821...78S/abstract>`_.
    The law is defined between 0.5 - 8 microns.

    Parameters
    ----------
    AH_AKs : float
        Ratio of A_H / A_Ks, which sets the normalization of the law (see Schlafly+16)
    x : float
        Free parameter in extinction law (see Schlafly+16, Eqn 6)
    """
    def __init__(self, AH_AKs, x):
        # Fetch the extinction curve, pre-interpolate across 1-8 microns
        wave = np.arange(0.5, 4.8, 0.001)

        # This will eventually be scaled by AK when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawSchlafly16._derive_Schlafly16(wave, AH_AKs, x)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Schlafly16',
                                                  litref='Schlafly+ 2016')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'S16,{0},{1}'.format(AH_AKs,x)

    @staticmethod
    def _derive_Schlafly16(wavelength, AH_AKs, x):
        """
        Calculate Schalfly+16 extinction law according to
        code provided in appendix of the paper. AH_AKs sets the
        gray component while x sets the shape of the law in an
        Rv-like way
        """
        # Use the function from the Schlafly+16 appendix to get the extinciton law
        # for given AH_AKs and x value. This is given in terms of A_lambda / A(5420)
        law_func = RedLawSchlafly16._Schlafly_appendix(x, AH_AKs)

        # Evaluate function for desired wavelengths (in angstroms)
        law = law_func(wavelength*10**4)

        # Now normalize to A_lambda/AKs, rather than A_lambda/A(5420)
        idx = np.where( abs(wavelength - 2.14) == min(abs(wavelength - 2.14)) )
        law_out = law / law[idx]

        return law_out

    @staticmethod
    def _Schlafly_appendix(x, rhk):
        """
        Schlafly+16 extinction law as defined in paper appendix. We've modified
        the wrapper slightly so that the user has control of rhk and x. Here is
        the comments from that code:

        Returns the extinction curve, A(lambda)/A(5420 A), according to
        Schlafly+2016, for the parameter "x," which controls the overall shape of
        the extinction curve in an R(V)-like way.  The extinction curve returned
        is a callable function, which is then invoked with the wavelength, in
        angstroms, of interest.

        The extinction curve is based on broad band photometry between the PS1 g
        band and the WISE W2 band, which have effective wavelengths between 5000
        and 45000 A.  The extinction curve is blindly extrapolated outside that
        range.  The gray component of the extinction curve is fixed by enforcing
        A(H)/A(K) = 1.55 (Indebetouw+2005).  The gray component is relatively
        uncertain, and its variation with x is largely made up.

        Args:
            x: some number controlling the shape of the extinction curve
            ra: extinction vector at anchor wavelengths, default to Schlafly+2016
            dra: derivative of extinction vector at anchor wavelengths, default to
             Schlafly+2016
            lam: anchor wavelengths (angstroms), default to Schlafly+2016

        Returns: the extinction curve E, so the extinction alam = A(lam)/A(5420 A)
        is given by:
            A = extcurve(x)
            alam = A(lam)
        """
        # Schlafly+2016
        ra = np.array([ 0.65373283,  0.39063843,  0.20197893,  0.07871701, -0.00476316,
                   -0.14213929, -0.23660605, -0.28522577, -0.321301  , -0.33503192])
        dra = np.array([-0.54278669,  0.03404903,  0.36841725,  0.42265873,  0.38247769,
                     0.14148814, -0.04020524, -0.13457319, -0.26883343, -0.36269229])

        # "isoreddening wavelengths" for extinction curve, at E(g-r) = 0.65 reddening
        # T_eff = 4500, Fe/H = 0, log g = 2.5
        lam = np.array([  5032.36441067,   6280.53335141,   7571.85928312,   8690.89321059,
                      9635.52560909,  12377.04268274,  16381.78146718,  21510.20523237,
                     32949.54009328,  44809.4919175 ])

        anchors = ra + x*dra
        # fix gray component so that A(H)/A(K) = 1.55
        anchors += (-anchors[6] + rhk*anchors[7])/(1 - rhk)
        cs0 = CubicSpline(lam, anchors, yp='3d=0')
        # normalize at 5420 angstroms
        return CubicSpline(lam, anchors/cs0(5420.), yp='3d=0')

    def Schlafly16(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawPowerLaw(pysynphot.reddening.CustomRedLaw):
    """
    Extinction object that is a power-law extinction law:
    :math:`A_{\lambda} \propto \lambda^{\alpha}`.

    For example, to create an extinction law between
    0.8 and 3 microns where :math:`\alpha = 2.21`,
    where :math:`A_{\lambda} / A_{Ks} = 1` at 2.12 microns:

    >>> red_law = reddening.RedLawPowerLaw(2.21, 2.12, wave_min=0.8, wave_max=3.0)

    Parameters
    ----------
    alpha : float
        Exponent of the extinction power-law.

    K_wave : float
        Extinction law is normalized such that AKs = 1 at `K_wave`.

    wave_min : float; optional
        Minimum wavelength of the extinction law, in microns.
        Default is 0.5 microns.

    wave_max : float; optional
        Maximum wavelength of the extinction law, in microns.
        Default is 5.0 microns
    """
    def __init__(self, alpha, K_wave, wave_min=0.5, wave_max=5.0):
        # Fetch the extinction curve, pre-interpolate across wave_min to wave_max
        wave = np.arange(wave_min, wave_max, 0.001)

        # This will eventually be scaled by AK when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawPowerLaw._derive_powerlaw(wave, alpha, K_wave)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Power law')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'pl,{0},{1},{2},{3}'.format(alpha,K_wave,wave_min,wave_max)

    @staticmethod
    def _derive_powerlaw(wavelength, alpha, K_wave):
        """
        Calculate the resulting extinction for an array of wavelengths.
        The extinction is normalized with A_Ks.

        Parameters
        ----------
        wavelength : float
            in microns

        alpha: float
            -1.0 * (power law exponent)

        K_wave: float
            Desired K-band wavelength, in microns
        """
        # Create extinction law
        law = wavelength**(-1.0 * alpha)

        # We'll identify K-band as 2.14 microns
        idx = np.where(abs(wavelength - K_wave) == min(abs(wavelength - K_wave)))
        A_AKs_at_wave = law / law[idx]

        return A_AKs_at_wave

    def powerlaw(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawFritz11(pysynphot.reddening.CustomRedLaw):
    """
    Defines extinction law from `Fritz et al. 2011
    <https://ui.adsabs.harvard.edu/abs/2011ApJ...737...73F/abstract>`_
    for the Galactic Center. The law is defined from 1.0 -- 19 microns.
    """
    def __init__(self):
        # Fetch the extinction curve, pre-interpolate across 3-8 microns
        wave = np.arange(1.0, 19, 0.001)

        # This will eventually be scaled by AKs when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawFritz11._derive_Fritz11(wave)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Fritz09',
                                                  litref='Fritz+2011')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'F11'

    @staticmethod
    def _derive_Fritz11(wavelength):
        """
        Calculate the resulting extinction for an array of wavelengths.
        The extinction is normalized to A_Ks = 1

        Data pulled from Fritz+11, Table 2

        Parameters
        ----------
        wavelength : float
            Wavelength range to derive extinction law over, in microns
        """
        # Extinction law definition
        wave = np.array([1.282, 1.736, 2.166, 2.625, 2.758, 2.873, 3.039, 3.297, 3.74, 3.819, 3.907, 4.052,
                             4.376, 5.128, 5.908, 6.772, 7.459, 7.502, 8.76, 12.371, 19.062])
        A_AKs = np.array([7.91, 4.30, 2.49, 1.83, 1.51, 1.84, 2.07, 1.66, 1.19, 1.19, 1.09, 1.01, 1.09, 0.99,
                              1.04, 0.84, 0.81, 0.79, 2.04, 1.34, 1.34])

        # Interpolate over the curve
        spline_interp = interpolate.splrep(wave, A_AKs, k=3, s=0)
        A_at_wave = interpolate.splev(wavelength, spline_interp)

        # We'll call 2.14 microns the K-band
        idx = np.where( abs(wavelength - 2.14) == min(abs(wavelength - 2.14)) )
        A_AKs_at_wave = A_at_wave / A_at_wave[idx]

        return A_AKs_at_wave

    def Fritz11(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawHosek18(pysynphot.reddening.CustomRedLaw):
    """
    Defines extinction law from `Hosek et al. 2018
    <https://ui.adsabs.harvard.edu/abs/2018ApJ...855...13H/abstract>`_
    for the Arches Cluster and Wd1. The law is defined between
    0.7 - 3.54 microns.

    WARNING: DEPRECATED! This law has revised to RedLawHosek18b, which
    should be used instead
    """
    def __init__(self):
        # Fetch the extinction curve, pre-interpolate across 3-8 microns
        wave = np.arange(0.7, 3.545, 0.001)

        # This will eventually be scaled by AKs when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawHosek18._derive_Hosek18(wave)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Hosek+18',
                                                  litref='Hosek+ 2018')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'H18'

    @staticmethod
    def _derive_Hosek18(wavelength):
        """
        Derive the Hosek+18 extinction law, using the data from Table 4.

        Calculate the resulting extinction for an array of wavelengths.
        The extinction is normalized with A_Ks.

        Data pulled from Hosek+18, Table 4

        Parameters
        ----------
        wavelength : float
            Wavelength range to define extinction law over, in microns
        """
        # Extinction law definition
        wave = np.array([0.8059, 0.962, 1.25, 1.53, 2.14, 3.545])
        A_AKs = np.array([9.66, 6.29, 3.56, 2.33, 1.0, 0.50])

        # Following Hosek+18, Interpolate over the curve with cubic spline interpolation
        spline_interp = interpolate.splrep(wave, A_AKs, k=3, s=0)
        A_AKs_at_wave = interpolate.splev(wavelength, spline_interp)

        # This curve already assumes A_Ks = 1.0, so we can go straight to
        # output
        return A_AKs_at_wave

    def Hosek18(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawHosek18b(pysynphot.reddening.CustomRedLaw):
    """
    Defines extinction law from `Hosek et al. 2019
    <https://ui.adsabs.harvard.edu/abs/2019ApJ...870...44H/abstract>`_
    for the Arches cluster and Wd1. This should be used over RedLawHosek18b.
    The law is derived between 0.7 - 3.54 microns
    """
    def __init__(self):
        # Fetch the extinction curve, pre-interpolate across 3-8 microns
        wave = np.arange(0.7, 3.545, 0.001)

        # This will eventually be scaled by AKs when you
        # call reddening(). Right now, calc for AKs=1
        Alambda_scaled = RedLawHosek18b._derive_Hosek18b(wave)

        # Convert wavelength to angstrom
        wave *= 10 ** 4

        pysynphot.reddening.CustomRedLaw.__init__(self, wave=wave,
                                                  waveunits='angstrom',
                                                  Avscaled=Alambda_scaled,
                                                  name='Hosek+18b',
                                                  litref='Hosek+ 2018b')

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = min(wave)
        self.high_lim = max(wave)
        self.name = 'H18b'

    @staticmethod
    def _derive_Hosek18b(wavelength):
        """
        Derive the Hosek+18 extinction law, using the data from Table 4.

        Calculate the resulting extinction for an array of wavelengths.
        The extinction is normalized with A_Ks.

        Data pulled from Hosek+18, Table 4

        Parameters
        ----------
        wavelength : float
            Wavelength range to define extinction law over, in microns
        """
        # Extinction law definition
        wave = np.array([0.8059, 0.962, 1.25, 1.53, 2.14, 3.545])
        A_AKs = np.array([7.943, 5.715, 3.142, 2.04, 1.0, 0.50])

        # Following Hosek+18, Interpolate over the curve with cubic spline interpolation
        spline_interp = interpolate.splrep(wave, A_AKs, k=3, s=0)
        A_AKs_at_wave = interpolate.splev(wavelength, spline_interp)

        # This curve already assumes A_Ks = 1.0, so we can go straight to
        # output
        return A_AKs_at_wave

    def Hosek18b(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

class RedLawNoguerasLara18(RedLawPowerLaw):
    """
    Defines extinction law from `Nogueras-Lara et al. 2018
    <https://ui.adsabs.harvard.edu/abs/2018A%26A...610A..83N/abstract>`_
    for the Galactic Center. It is defined between 0.8 - 2.5 microns.
    """
    def __init__(self):
        wave_min = 0.8
        wave_max = 2.8
        RedLawPowerLaw.__init__(self, 2.30, 2.15, wave_min=wave_min, wave_max=wave_max)

        # Set the upper/lower wavelength limits of law (in angstroms)
        self.low_lim = wave_min*10**4
        self.high_lim = wave_max*10**4
        self.name = 'NL18'

    def NoguerasLara18(self, wavelength, AKs):
        """
        Return the extinction at a given wavelength assuming the
        extinction law and an overall `AKs` value.

        Parameters
        ----------
        wavelength : float or array
            Wavelength to return extinction for, in microns
        AKs : float
            Total extinction in AKs, in mags
        """
        # If input entry is a single float, turn it into an array
        try:
            len(wavelength)
        except:
            wavelength = [wavelength]

        # Return error if any wavelength is beyond interpolation range of
        # extinction law
        if ((min(wavelength) < (self.low_lim*10**-4)) | (max(wavelength) > (self.high_lim*10**-4))):
            return ValueError('{0}: wavelength values beyond interpolation range'.format(self))

        # Extract wave and A/AKs from law, turning wave into micron units
        wave = self.wave * (10**-4)
        law = self.obscuration

        # Find the value of the law at the closest points
        # to wavelength
        A_AKs_at_wave = []
        for ii in wavelength:
            idx = np.where( abs(wave - ii) == min(abs(wave - ii)) )
            A_AKs_at_wave.append(law[idx][0])

        # Now multiply by AKs (since law assumes AKs = 1)
        A_at_wave = np.array(A_AKs_at_wave) * AKs

        return A_at_wave

#---------------------------#
# Cubic spline function from Schalfly+16 appendix
#---------------------------#
def splint(spl, x):
    npts = len(spl.x)
    lo = np.searchsorted(spl.x, x)-1
    lo = np.clip(lo, 0, npts-2)
    hi = lo + 1
    dx = spl.x[hi] - spl.x[lo]
    a = (spl.x[hi] - x)/dx
    b = (x-spl.x[lo])/dx
    y = (a*spl.y[lo]+b*spl.y[hi]+
         ((a**3-a)*spl.y2[lo]+(b**3-b)*spl.y2[hi])*dx**2./6.)
    return y

class CubicSpline:
    def __init__(self, x, y, yp=None):
        npts = len(x)
        mat = np.zeros((3, npts))
        # enforce continuity of 1st derivatives
        mat[1,1:-1] = (x[2:  ]-x[0:-2])/3.
        mat[2,0:-2] = (x[1:-1]-x[0:-2])/6.
        mat[0,2:  ] = (x[2:  ]-x[1:-1])/6.
        bb = np.zeros(npts)
        bb[1:-1] = ((y[2:  ]-y[1:-1])/(x[2:  ]-x[1:-1]) -
                    (y[1:-1]-y[0:-2])/(x[1:-1]-x[0:-2]))
        if yp is None: # natural cubic spline
            mat[1,0] = 1.
            mat[1,-1] = 1.
            bb[0] = 0.
            bb[-1] = 0.
        elif yp == '3d=0':
            mat[1, 0] = -1./(x[1]-x[0])
            mat[0, 1] =  1./(x[1]-x[0])
            mat[1,-1] =  1./(x[-2]-x[-1])
            mat[2,-2] = -1./(x[-2]-x[-1])
            bb[ 0] = 0.
            bb[-1] = 0.
        else:
            mat[1, 0] = -1./3.*(x[1]-x[0])
            mat[0, 1] = -1./6.*(x[1]-x[0])
            mat[2,-2] =  1./6.*(x[-1]-x[-2])
            mat[1,-1] =  1./3.*(x[-1]-x[-2])
            bb[ 0] = yp[0]-1.*(y[ 1]-y[ 0])/(x[ 1]-x[ 0])
            bb[-1] = yp[1]-1.*(y[-1]-y[-2])/(x[-1]-x[-2])
        y2 = solve_banded((1,1), mat, bb)
        self.x, self.y, self.y2 = (x, y, y2)
    def __call__(self, x):
        return splint(self, x)


#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
J-metric density fitting
'''

import time
import tempfile
import numpy
import h5py
from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.df import incore
from pyscf.df import outcore
from pyscf.df import r_incore
from pyscf.df import addons
from pyscf.df import df_jk
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo.incore import _conc_mos, iden_coeffs
from pyscf import __config__

class DF(lib.StreamObject):
    r'''
    Object to hold 3-index tensor

    Attributes:
        auxbasis : str or dict
            Same input format as :attr:`Mole.basis`

        auxmol : Mole object
            Read only Mole object to hold the auxiliary basis.  auxmol is
            generated automatically in the initialization step based on the
            given auxbasis.  It is used in the rest part of the code to
            determine the problem size, the integral batches etc.  This object
            should NOT be modified.
        _cderi_to_save : str
            If _cderi_to_save is specified, the DF integral tensor will be
            saved in this file.
        _cderi : str or numpy array
            If _cderi is specified, the DF integral tensor will be read from
            this HDF5 file (or numpy array). When the DF integral tensor is
            provided from the HDF5 file, it has to be stored under the dataset
            'j3c'.
            The DF integral tensor :math:`V_{x,ij}` should be a 2D array in C
            (row-major) convention, where x corresponds to index of auxiliary
            basis, and the combined index ij is the orbital pair index. The
            hermitian symmetry is assumed for the combined ij index, ie
            the elements of :math:`V_{x,i,j}` with :math:`i\geq j` are existed
            in the DF integral tensor.  Thus the shape of DF integral tensor
            is (M,N*(N+1)/2), where M is the number of auxbasis functions and
            N is the number of basis functions of the orbital basis.
        blockdim : int
            When reading DF integrals from disk the chunk size to load.  It is
            used to improve the IO performance.
    '''
    def __init__(self, mol):
        self.mol = mol
        self.stdout = mol.stdout
        self.verbose = mol.verbose
        self.max_memory = mol.max_memory
        self._auxbasis = None

##################################################
# Following are not input options
        self.auxmol = None
# If _cderi_to_save is specified, the 3C-integral tensor will be saved in this file.
        self._cderi_to_save = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
# If _cderi is specified, the 3C-integral tensor will be read from this file
        self._cderi = None
        self._call_count = getattr(__config__, 'df_df_DF_call_count', None)
        self.blockdim = getattr(__config__, 'df_df_DF_blockdim', 240)
        self._keys = set(self.__dict__.keys())

    @property
    def auxbasis(self):
        return self._auxbasis
    @auxbasis.setter
    def auxbasis(self, x):
        if self._auxbasis != x:
            self._auxbasis = x
            self.auxmol = None
            self._cderi = None

    def dump_flags(self):
        log = logger.Logger(self.stdout, self.verbose)
        log.info('******** %s flags ********', self.__class__)
        if self.auxmol is None:
            log.info('auxbasis = %s', self.auxbasis)
        else:
            log.info('auxbasis = auxmol.basis = %s', self.auxmol.basis)
        log.info('max_memory = %s', self.max_memory)
        if isinstance(self._cderi, str):
            log.info('_cderi = %s  where DF integrals are loaded (readonly).',
                     self._cderi)
        if isinstance(self._cderi_to_save, str):
            log.info('_cderi_to_save = %s', self._cderi_to_save)
        else:
            log.info('_cderi_to_save = %s', self._cderi_to_save.name)
        return self

    def build(self):
        t0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        self.check_sanity()
        self.dump_flags()

        mol = self.mol
        auxmol = self.auxmol = addons.make_auxmol(self.mol, self.auxbasis)
        nao = mol.nao_nr()
        naux = auxmol.nao_nr()
        nao_pair = nao*(nao+1)//2

        max_memory = (self.max_memory - lib.current_memory()[0]) * .8
        int3c = mol._add_suffix('int3c2e')
        int2c = mol._add_suffix('int2c2e')
        if (nao_pair*naux*3*8/1e6 < max_memory and
            not isinstance(self._cderi_to_save, str)):
            self._cderi = incore.cholesky_eri(mol, int3c=int3c, int2c=int2c,
                                              auxmol=auxmol, verbose=log)
        else:
            if isinstance(self._cderi_to_save, str):
                cderi = self._cderi_to_save
            else:
                cderi = self._cderi_to_save.name
            if isinstance(self._cderi, str):
                log.warn('Value of _cderi is ignored. DF integrals will be '
                         'saved in file %s .', cderi)
            outcore.cholesky_eri(mol, cderi, dataname='j3c',
                                 int3c=int3c, int2c=int2c, auxmol=auxmol,
                                 max_memory=max_memory, verbose=log)
            if nao_pair*naux*8/1e6 < max_memory:
                with addons.load(cderi, 'j3c') as feri:
                    cderi = numpy.asarray(feri)
            self._cderi = cderi
            log.timer_debug1('Generate density fitting integrals', *t0)
        return self

    def kernel(self, *args, **kwargs):
        return self.build(*args, **kwargs)

    def loop(self, blksize=None):
        if self._cderi is None:
            self.build()
        if blksize is None:
            blksize = self.blockdim
        with addons.load(self._cderi, 'j3c') as feri:
            naoaux = feri.shape[0]
            for b0, b1 in self.prange(0, naoaux, blksize):
                eri1 = numpy.asarray(feri[b0:b1], order='C')
                yield eri1

    def prange(self, start, end, step):
        if isinstance(self._call_count, int):
            self._call_count += 1
            if self._call_count % 2 == 1:
                for i in reversed(range(start, end, step)):
                    yield i, min(i+step, end)
            else:
                for i in range(start, end, step):
                    yield i, min(i+step, end)

        else:
            for i in range(start, end, step):
                yield i, min(i+step, end)

    def get_naoaux(self):
# determine naoaux with self._cderi, because DF object may be used as CD
# object when self._cderi is provided.
        if self._cderi is None:
            self.build()
        with addons.load(self._cderi, 'j3c') as feri:
            return feri.shape[0]

    def get_jk(self, dm, hermi=1, vhfopt=None, with_j=True, with_k=True):
        return df_jk.get_jk(self, dm, hermi, vhfopt, with_j, with_k)

    def get_eri(self):
        nao = self.mol.nao_nr()
        nao_pair = nao * (nao+1) // 2
        ao_eri = numpy.zeros((nao_pair,nao_pair))
        for eri1 in self.loop():
            lib.dot(eri1.T, eri1, 1, ao_eri, 1)
        return ao2mo.restore(8, ao_eri, nao)
    get_ao_eri = get_eri

    def ao2mo(self, mo_coeffs,
              compact=getattr(__config__, 'df_df_DF_ao2mo_compact', True)):
        if isinstance(mo_coeffs, numpy.ndarray) and mo_coeffs.ndim == 2:
            mo_coeffs = (mo_coeffs,) * 4
        ijmosym, nij_pair, moij, ijslice = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact)
        klmosym, nkl_pair, mokl, klslice = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact)
        mo_eri = numpy.zeros((nij_pair,nkl_pair))
        sym = (iden_coeffs(mo_coeffs[0], mo_coeffs[2]) and
               iden_coeffs(mo_coeffs[1], mo_coeffs[3]))
        Lij = Lkl = None
        for eri1 in self.loop():
            Lij = _ao2mo.nr_e2(eri1, moij, ijslice, aosym='s2', mosym=ijmosym, out=Lij)
            if sym:
                Lkl = Lij
            else:
                Lkl = _ao2mo.nr_e2(eri1, mokl, klslice, aosym='s2', mosym=klmosym, out=Lkl)
            lib.dot(Lij.T, Lkl, 1, mo_eri, 1)
        return mo_eri
    get_mo_eri = ao2mo

class DF4C(DF):
    '''Relativistic 4-component'''
    def build(self):
        log = logger.Logger(self.stdout, self.verbose)
        mol = self.mol
        auxmol = self.auxmol = addons.make_auxmol(self.mol, self.auxbasis)
        n2c = mol.nao_2c()
        naux = auxmol.nao_nr()
        nao_pair = n2c*(n2c+1)//2

        max_memory = (self.max_memory - lib.current_memory()[0]) * .8
        if nao_pair*naux*3*16/1e6*2 < max_memory:
            self._cderi =(r_incore.cholesky_eri(mol, auxmol=auxmol, aosym='s2',
                                                int3c='int3c2e_spinor', verbose=log),
                          r_incore.cholesky_eri(mol, auxmol=auxmol, aosym='s2',
                                                int3c='int3c2e_spsp1_spinor', verbose=log))
        else:
            raise NotImplementedError
        return self

    def loop(self):
        if self._cderi is None:
            self.build()
        with addons.load(self._cderi[0], 'j3c') as ferill:
            naoaux = ferill.shape[0]
            with addons.load(self._cderi[1], 'j3c') as feriss: # python2.6 not support multiple with
                for b0, b1 in self.prange(0, naoaux, self.blockdim):
                    erill = numpy.asarray(ferill[b0:b1], order='C')
                    eriss = numpy.asarray(feriss[b0:b1], order='C')
                    yield erill, eriss

    def get_jk(self, dm, hermi=1, vhfopt=None, with_j=True, with_k=True):
        return df_jk.r_get_jk(self, dm, hermi)

    def ao2mo(self, mo_coeffs):
        raise NotImplementedError


"""Simplified retina model."""

import logging
import numpy as np
import cv2
import cv2.cv as cv
from collections import OrderedDict

from lumos.context import Context
from lumos.input import Projector, run

from ..photoreceptor import Rod, Cone

class Retina(object):
  """
  A multi-layered surface for hosting different types of neurons that make up a retina, simplified version.

  [Deprecated] Use VisualSystem instead.

  """

  default_image_size = (480, 480)

  def __init__(self, imageSize=default_image_size, timeNow=0.0):
    # * Initialize members, parameters
    self.context = Context.getInstance()
    self.logger = logging.getLogger(__name__)
    self.logger.debug("Creating simplified Retina")  # to distinguish from other Retina versions
    self.imageSize = imageSize
    self.imageCenter = (self.imageSize[1] / 2, self.imageSize[0] / 2)
    self.timeNow = timeNow
    self.bounds = np.float32([[0.0, 0.0, 2.0], [self.imageSize[0] - 1, self.imageSize[1] - 1, 4.0]])
    self.center = (self.bounds[0] + self.bounds[1]) / 2
    self.logger.debug("Retina center: {}, image size: {}".format(self.center, self.imageSize))

    self.bipolarBlurSize = (5, 5)  # size of blurring kernel used when computing Bipolar cell response
    self.ganglionCenterSurroundKernel = np.float32(
      [ [ -1, -1, -1, -1, -1, -1, -1 ],
        [ -1, -1, -1, -1, -1, -1, -1 ],
        [ -1, -1,  7,  7,  7, -1, -1 ],
        [ -1, -1,  7,  9,  7, -1, -1 ],
        [ -1, -1,  7,  7,  7, -1, -1 ],
        [ -1, -1, -1, -1, -1, -1, -1 ],
        [ -1, -1, -1, -1, -1, -1, -1 ] ])
    self.ganglionCenterSurroundKernel /= np.sum(self.ganglionCenterSurroundKernel)  # normalize
    #self.logger.info("Ganglion center-surround kernel:\n{}".format(self.ganglionCenterSurroundKernel))  # [debug]
    self.ganglionKernelLevels = 4
    self.ganglionKernels = [None] * self.ganglionKernelLevels
    self.ganglionKernels[0] = self.ganglionCenterSurroundKernel
    for i in xrange(1, self.ganglionKernelLevels):
      self.ganglionKernels[i] = cv2.resize(self.ganglionKernels[i - 1], dsize=None, fx=2, fy=2)
      self.ganglionKernels[i] /= np.sum(self.ganglionKernels[i])  # normalize
    #self.logger.info("Ganglion center-surround kernel sizes ({} levels): {}".format(self.ganglionKernelLevels, ", ".join("{}".format(k.shape) for k in self.ganglionKernels)))  # [debug]

    # * Image and related members
    self.imageCenter = (self.imageSize[1] / 2, self.imageSize[0] / 2)
    self.imageShapeC3 = (self.imageSize[1], self.imageSize[0], 3)  # numpy shape for 3 channel images
    self.imageShapeC1 = (self.imageSize[1], self.imageSize[0])  # numpy shape for single channel images
    # NOTE Image shapes (h, w, 1) and (h, w) are not compatible unless we use keepdims=True for numpy operations
    self.imageTypeInt = np.uint8  # numpy dtype for integer-valued images
    self.imageTypeFloat = np.float32  # numpy dtype for real-valued images
    self.images = OrderedDict()

    # ** RGB and HSV images
    self.images['BGR'] = np.zeros(self.imageShapeC3, dtype=self.imageTypeInt)
    self.images['HSV'] = np.zeros(self.imageShapeC3, dtype=self.imageTypeInt)
    self.images['H'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeInt)
    self.images['S'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeInt)
    self.images['V'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeInt)

    # ** Freq/hue-dependent response images for rods and different cone types
    self.imageRod = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesCone = dict()  # NOTE dict keys must match names of Cone.cone_types
    self.imagesCone['S'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesCone['M'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesCone['L'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)

    # ** Bipolar and Ganglion cell response images
    # TODO Add more Ganglion cell types with different receptive field properties (color-opponent cells)
    #   'RG' +Red    -Green
    #   'GR' +Green  -Red
    #   'RB' +Red    -Blue
    #   'BR' +Blue   -Red
    #   'BY' +Blue   -Yellow
    #   'YB' +Yellow -Blue
    #   'WK' +White  -Black (currently 'ON')
    #   'KW' +Black  -White (currently 'OFF')
    # NOTE: R = L cones, G = M cones, B = S cones
    self.imagesBipolar = dict()
    self.imagesBipolar['ON'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesBipolar['OFF'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesGanglion = dict()
    self.imagesGanglion['ON'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesGanglion['OFF'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    # TODO Verify why image shapes (h, w, 1) and (h, w) are not compatible (use keepdims=True for numpy operations)
    self.imagesGanglion['RG'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesGanglion['GR'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesGanglion['RB'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesGanglion['BR'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesGanglion['BY'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    self.imagesGanglion['YB'] = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)

    # ** Combined response (salience) image
    self.imageSalience = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)

    # ** Spatial attention map with a central (covert) spotlight (currently unused; TODO move to VisualCortex? also, use np.ogrid?)
    self.imageAttention = np.zeros(self.imageShapeC1, dtype=self.imageTypeFloat)
    cv2.circle(self.imageAttention, (self.imageSize[1] / 2, self.imageSize[0] / 2), self.imageSize[0] / 3, 1.0, cv.CV_FILLED)
    self.imageAttention = cv2.blur(self.imageAttention, (self.imageSize[0] / 4, self.imageSize[0] / 4))  # coarse blur

    # ** Output image(s)
    if self.context.options.gui:
      self.imageOut = np.zeros(self.imageShapeC3, dtype=self.imageTypeInt)

  def initialize(self, imageIn, timeNow):
    pass  # to emulate FrameProcessor-like interface

  def process(self, imageIn, timeNow):
    self.timeNow = timeNow
    self.logger.debug("Retina update @ {}".format(self.timeNow))

    # * Get HSV
    self.images['BGR'][:] = imageIn
    self.images['HSV'] = cv2.cvtColor(self.images['BGR'], cv2.COLOR_BGR2HSV)
    self.images['H'], self.images['S'], self.images['V'] = cv2.split(self.images['HSV'])

    # * Compute Rod and Cone responses
    # TODO Need non-linear response to hue, sat, val (less dependent on sat, val for cones)
    self.imageRod = np.float32(180 - cv2.absdiff(self.images['H'], Rod.rod_type.hue) % 180) * 255 * self.images['V'] * Rod.rod_type.responseFactor  # hack: use constant sat = 200 to make response independent of saturation
    self.imagesCone['S'] = np.float32(180 - cv2.absdiff(self.images['H'], Cone.cone_types[0].hue) % 180) * self.images['S'] * self.images['V'] * Cone.cone_types[0].responseFactor
    self.imagesCone['M'] = np.float32(180 - cv2.absdiff(self.images['H'], Cone.cone_types[1].hue) % 180) * self.images['S'] * self.images['V'] * Cone.cone_types[1].responseFactor
    self.imagesCone['L'] = np.float32(180 - cv2.absdiff(self.images['H'], Cone.cone_types[2].hue) % 180) * self.images['S'] * self.images['V'] * Cone.cone_types[2].responseFactor

    # * Compute Bipolar and Ganglion cell responses
    # ** Blurring is a step that is effectively achieved in biology by horizontal cells
    imageRodBlurred = cv2.blur(self.imageRod, self.bipolarBlurSize)
    self.imagesBipolar['ON'] = np.clip(self.imageRod - 0.75 * imageRodBlurred, 0.0, 1.0)
    self.imagesBipolar['OFF'] = np.clip((1.0 - self.imageRod) - 0.75 * (1.0 - imageRodBlurred), 0.0, 1.0)  # same as (1 - ON response)?
    #imagesConeSBlurred = cv2.blur(self.imagesCone['S'], self.bipolarBlurSize)
    #imagesConeMBlurred = cv2.blur(self.imagesCone['M'], self.bipolarBlurSize)
    #imagesConeLBlurred = cv2.blur(self.imagesCone['L'], self.bipolarBlurSize)
    # ** Ganglion cells simply add up responses from a (bunch of) central bipolar cell(s) (ON/OFF) and surrounding antagonistic bipolar cells (OFF/ON)
    # *** Method 1: Center - Surround
    #imageGanglionCenterON = cv2.filter2D(self.imagesBipolar['ON'], -1, self.ganglionCenterKernel)
    #imageGanglionSurroundOFF = cv2.filter2D(self.imagesBipolar['OFF'], -1, self.ganglionSurroundKernel)
    #self.imagesGanglion['ON'] = 0.75 * imageGanglionCenterON + 0.25 * imageGanglionSurroundOFF
    # *** Method 2: Center-Surround kernel
    #self.imagesGanglion['ON'] = np.clip(cv2.filter2D(self.imagesBipolar['ON'], -1, self.ganglionCenterSurroundKernel), 0.0, 1.0)
    #self.imagesGanglion['OFF'] = np.clip(cv2.filter2D(self.imagesBipolar['OFF'], -1, self.ganglionCenterSurroundKernel), 0.0, 1.0)
    # *** Method 3: Multi-level Center-Surround kernels, taking maximum
    self.imagesGanglion['ON'].fill(0.0)
    self.imagesGanglion['OFF'].fill(0.0)
    self.imagesGanglion['RG'].fill(0.0)
    self.imagesGanglion['GR'].fill(0.0)
    self.imagesGanglion['RB'].fill(0.0)
    self.imagesGanglion['BR'].fill(0.0)
    self.imagesGanglion['BY'].fill(0.0)
    self.imagesGanglion['YB'].fill(0.0)

    for k in self.ganglionKernels:
      # Rod pathway
      self.imagesGanglion['ON'] = np.maximum(self.imagesGanglion['ON'], np.clip(cv2.filter2D(self.imagesBipolar['ON'], -1, k), 0.0, 1.0))
      self.imagesGanglion['OFF'] = np.maximum(self.imagesGanglion['OFF'], np.clip(cv2.filter2D(self.imagesBipolar['OFF'], -1, k), 0.0, 1.0))
      # Cone pathway
      imageRG = self.imagesCone['L'] - self.imagesCone['M']
      imageRB = self.imagesCone['L'] - self.imagesCone['S']
      imageBY = self.imagesCone['S'] - (self.imagesCone['L'] + self.imagesCone['M']) / 2
      self.imagesGanglion['RG'] = np.maximum(self.imagesGanglion['RG'], np.clip(cv2.filter2D(imageRG, -1, k), 0.0, 1.0))
      self.imagesGanglion['GR'] = np.maximum(self.imagesGanglion['GR'], np.clip(cv2.filter2D(-imageRG, -1, k), 0.0, 1.0))
      self.imagesGanglion['RB'] = np.maximum(self.imagesGanglion['RB'], np.clip(cv2.filter2D(imageRB, -1, k), 0.0, 1.0))
      self.imagesGanglion['BR'] = np.maximum(self.imagesGanglion['BR'], np.clip(cv2.filter2D(-imageRB, -1, k), 0.0, 1.0))
      self.imagesGanglion['BY'] = np.maximum(self.imagesGanglion['BY'], np.clip(cv2.filter2D(imageBY, -1, k), 0.0, 1.0))
      self.imagesGanglion['YB'] = np.maximum(self.imagesGanglion['YB'], np.clip(cv2.filter2D(-imageBY, -1, k), 0.0, 1.0))

    # * Compute combined (salience) image; TODO incorporate attention weighting (spatial, as well as by visual feature)
    # ** Method 1: Max of all Ganglion cell images
    self.imageSalience.fill(0.0)
    for ganglionType, ganglionImage in self.imagesGanglion.iteritems():
      self.imageSalience = np.maximum(self.imageSalience, ganglionImage)

    #self.imageSalience *= self.imageAttention  # TODO evaluate if this is necessary

    # * TODO Compute feature vector of attended region

    # * Show output images if in GUI mode
    if self.context.options.gui:
      #cv2.imshow("Hue", self.images['H'])
      #cv2.imshow("Saturation", self.images['S'])
      #cv2.imshow("Value", self.images['V'])
      cv2.imshow("Rod response", self.imageRod)
      cv2.imshow("S-cone response", self.imagesCone['S'])
      cv2.imshow("M-cone response", self.imagesCone['M'])
      cv2.imshow("L-cone response", self.imagesCone['L'])
      cv2.imshow("ON Bipolar cells", self.imagesBipolar['ON'])
      cv2.imshow("OFF Bipolar cells", self.imagesBipolar['OFF'])
      #cv2.imshow("ON Ganglion cells", self.imagesGanglion['ON'])
      #cv2.imshow("OFF Ganglion cells", self.imagesGanglion['OFF'])
      for ganglionType, ganglionImage in self.imagesGanglion.iteritems():
        cv2.imshow("{} Ganglion cells".format(ganglionType), ganglionImage)
      cv2.imshow("Salience", self.imageSalience)

      # Designate a representative output image
      self.imageOut = self.imageSalience
      #_, self.imageOut = cv2.threshold(self.imageOut, 0.15, 1.0, cv2.THRESH_TOZERO)  # apply threshold to remove low-response regions

    return True, self.imageOut

if __name__ == "__main__":
  Context.createInstance(description="Test application that uses a SimplifiedProjector to run image input through a (simplified) Retina.")
  run(Projector(Retina()))


import time
import multiprocessing as mp
from typing import Tuple, Dict, Optional

import numpy as np
import pandas as pd
from scipy.interpolate.interpolate import interp1d

import sha_calc as sha_calc
from gmhazard_calc import site
from gmhazard_calc import utils
from gmhazard_calc import shared
from gmhazard_calc import gm_data
from gmhazard_calc import site_source
from gmhazard_calc import constants as const
from gmhazard_calc.im import IM
from .HazardResult import BranchHazardResult, EnsembleHazardResult

DEFAULT_N_IM_VALUES = 200

def run_ensemble_hazard(
    ensemble: gm_data.Ensemble,
    site_info: site.SiteInfo,
    im: IM,
    branch_hazard: Optional[Dict[str, BranchHazardResult]] = None,
    im_values: Optional[np.ndarray] = None,
    calc_percentiles: bool = True,
) -> EnsembleHazardResult:
    """Computes the weighted hazard curve for all branches in
    the specified ensemble.

    Parameters
    ----------
    ensemble: Ensemble
        ensemble to use for calculation
    site_info: SiteInfo
        The site at which to calculate the hazard curve
    im: IM
        IM object for specifying the IM to use for calculations
    branch_hazard: Dictionary of str: HazardResult, optional
        Where the key is the branch name the hazard result is for.
        If specified then this saves re-computing the hazard
        results for the branches.
    im_values: np.ndarray, optional
        The range of IM values for which to calculate the
        hazard, not used if branches_hazard is passed in
    calc_percentiles: bool, optional
        True or False to calculate the 16th and 84th percentiles

    Returns
    -------
    HazardResult
    """

    def get_weighted_branch_hazard(hazard: BranchHazardResult):
        return (
            hazard.branch.weight * hazard.fault_hazard,
            hazard.branch.weight * hazard.ds_hazard,
        )

    ensemble.check_im(im)

    # Get the hazard per branch
    if branch_hazard is None:
        branch_hazard = run_branches_hazard(
            ensemble,
            site_info,
            im,
            im_values=im_values,
        )

    # Combine the branches according to their weights
    fault_hazard, ds_hazard = None, None
    for branch_name, cur_hazard in branch_hazard.items():
        cur_fault_h, cur_ds_h = get_weighted_branch_hazard(cur_hazard)
        if fault_hazard is None:
            fault_hazard, ds_hazard = cur_fault_h, cur_ds_h
        else:
            fault_hazard += cur_fault_h
            ds_hazard += cur_ds_h

    # Compute 16th and 84th percentile if flag enabled
    percentiles = None
    if calc_percentiles:
        # Retrieving data
        im_values = fault_hazard.index.values
        excd_values, weights = [], []
        for cur_branch in branch_hazard.values():
            assert np.all(cur_branch.fault_hazard.index.values == im_values)
            excd_values.append(cur_branch.total_hazard.values)
            weights.append(cur_branch.branch.weight)
        excd_values, weights = np.asarray(excd_values).T, np.asarray(weights)
        weights = np.repeat(weights[None, ...], im_values.size, 0)

        # Sorting
        sort_ind = np.argsort(excd_values, axis=1)
        excd_values = np.take_along_axis(excd_values, sort_ind, 1)
        weights = np.take_along_axis(weights, sort_ind, 1)

        # Inverse CDF lookup
        cdf_x, cdf_y = excd_values, np.cumsum(weights, axis=1)
        x_values = sha_calc.shared.query_non_parametric_multi_cdf_invs(
            [0.16, 0.84], cdf_x, cdf_y
        )
        x_values = np.stack(x_values, axis=1)
        percentiles = pd.DataFrame(
            data=x_values, columns=["16th", "84th"], index=fault_hazard.index.values
        )

    return EnsembleHazardResult(
        im,
        site_info,
        fault_hazard,
        ds_hazard,
        ensemble,
        list(branch_hazard.values()),
        percentiles=percentiles,
    )

def run_branches_hazard(
    ensemble: gm_data.Ensemble,
    site_info: site.SiteInfo,
    im: IM,
    im_values: Optional[np.ndarray] = None,
) -> Dict[str, BranchHazardResult]:
    """Runs computation of the hazard curve for each of the branches in
    the specified IM-ensemble.

    Parameters
    ----------
    ensemble : Ensemble
        Ensemble to use for calculation
    site_info : SiteInfo
        The site at which to calculate the hazard curve
    im : IM
        IM Object to use for calculations
    im_values: array of floats, optional
        The IM values for which to calculate the hazard for.

    Returns
    -------
    Dict of str : HazardResult, where the key is the branch name
    """
    ensemble.check_im(im)
    im_ensemble = ensemble.get_im_ensemble(im.im_type)

    hazards = {}
    for branch_name, branch in im_ensemble.branches_dict.items():
        hazards[branch_name] = run_branch_hazard(
            branch, site_info, im, im_values=im_values
        )

    return hazards

def run_branch_hazard(
    branch: gm_data.Branch,
    site_info: site.SiteInfo,
    im: IM,
    im_values: Optional[np.ndarray] = None,
) -> BranchHazardResult:
    """Computes the hazard for a single branch

    Parameters
    ----------
    branch: Branch
        The branch for which to calculate the hazard curve
    site_info: SiteInfo
        The site at which to calculate the hazard curve
    im: IM
        IM Object used for calculations
    im_values: np.ndarray, optional
        The IM values for which to calculate the hazard for.

    Returns
    -------
    HazardResult
    """
    im_values = (
        utils.get_im_values(im, n_values=DEFAULT_N_IM_VALUES)
        if im_values is None
        else im_values
    )

    # Fault Hazard
    fault_gm_prob_df = shared.get_gm_prob_df(
        branch,
        site_info,
        im,
        im_values,
        const.SourceType.fault,
        ensemble=branch.im_ensemble.ensemble,
    )
    if fault_gm_prob_df is not None:
        fault_hazard = sha_calc.hazard_curve(
            fault_gm_prob_df, branch.rupture_df_id_ix["annual_rec_prob"]
        )
    else:
        fault_hazard = pd.Series(data=np.zeros(im_values.shape), index=im_values)

    # DS Hazard
    ds_gm_prob_df = shared.get_gm_prob_df(
        branch,
        site_info,
        im,
        im_values,
        const.SourceType.distributed,
        ensemble=branch.im_ensemble.ensemble,
    )
    if ds_gm_prob_df is not None:
        ds_hazard = sha_calc.hazard_curve(
            ds_gm_prob_df, branch.rupture_df_id_ix["annual_rec_prob"]
        )
    else:
        ds_hazard = pd.Series(data=np.zeros(im_values.shape), index=im_values)

    return BranchHazardResult(im, site_info, fault_hazard, ds_hazard, branch)

def run_full_hazard(
    ensemble: gm_data.Ensemble,
    site_info: site.SiteInfo,
    im: IM,
    calc_percentiles: bool = False,
    im_values: Optional[np.ndarray] = None,
) -> Tuple[EnsembleHazardResult, Dict[str, BranchHazardResult]]:
    """Convenience function, computes the ensemble
     and hazard for all branches.

    Parameters
    ----------
    branch: Branch
        The branch for which to calculate the hazard curve
    site_info: SiteInfo
        The site at which to calculate the hazard curve
    im: IM
        IM Object to use for calculations
    calc_percentiles: bool, optional
        True or false for calculating 16th and 84th percentiles
    im_values: np.ndarray, optional
        The IM values for which to calculate the hazard for.

    Returns
    -------
    HazardResult:
        The ensemble hazard
    dict:
        The hazard for each branch, key is the branch name
    """
    branch_hazard = run_branches_hazard(ensemble, site_info, im, im_values=im_values)
    ens_hazard = run_ensemble_hazard(
        ensemble,
        site_info,
        im,
        calc_percentiles=calc_percentiles,
        branch_hazard=branch_hazard,
        im_values=im_values,
    )

    return ens_hazard, branch_hazard

def run_hazard_map(
    ensemble: gm_data.Ensemble, im: IM, exceedance: float, n_procs: Optional[int] = 4
) -> pd.DataFrame:
    """
    Computes the hazard at each station in the ensemble for the
    specified exceedance.

    Parameters
    ----------
    ensemble: Ensemble
    im: IM
        IM Object used for calculations
    exceedance: float
        The exceedance value
    n_procs:
        Number of processes to use

    Returns
    -------
    pd.Series
        format: index = station_name, values: exceedance probability
    """
    # Drop duplicate location stations
    stations_df = ensemble.stations.drop_duplicates(subset=["lon", "lat"])

    n_stations = stations_df.shape[0]
    if n_procs == 1:
        excd_probs = []
        for ix, station_name in enumerate(stations_df.index.values):
            excd_probs.append(
                _get_hazard(ensemble, station_name, im, exceedance, ix, n_stations)
            )
    else:
        with mp.Pool(n_procs) as p:
            excd_probs = p.starmap(
                _get_hazard,
                [
                    (ensemble, station_name, im, exceedance, ix, n_stations)
                    for ix, station_name in enumerate(stations_df.index.values)
                ],
            )

    result_df = stations_df.copy()
    result_df["value"] = excd_probs
    return result_df

def get_exceedance_rate(probability: float, years: int):
    """Gets the exceedance rate for the specified probability
    in number of specified years

    Parameters
    ----------
    probability: float
        The probability of interest (e.g. 50 for 50%)
    years: int
        The number of years

    Returns
    -------
    float
        The exceedance rate
    """
    return -1.0 / years * np.log(1 - (probability / 100))

def exceedance_to_im(
    exceedance: float, im_values: np.ndarray, hazard_values: np.ndarray
):
    """Converts the given exceedance rate to an IM value, based on the
    provided im and hazard values

    Parameters
    ----------
    exceedance: float
        The exceedance value of interest
    im_values: numpy array
        The IM values corresponding to the hazard values
        Has to be the same shape as hazard_values
    hazard_values: numpy array
        The hazard values corresponding to the IM values
        Has to be the same shape as im_values

    Returns
    -------
    float
        The IM value corresponding to the provided exceedance
    """
    return np.exp(
        interp1d(
            np.log(hazard_values) * -1,
            np.log(im_values),
            kind="linear",
            bounds_error=True,
        )(np.log(exceedance) * -1)
    )

def im_to_exceedance(im_value: float, im_values: np.ndarray, hazard_values: np.ndarray):
    """Inverse to exceedance_to_im"""
    return np.exp(
        interp1d(
            np.log(im_values), np.log(hazard_values), kind="linear", bounds_error=True
        )(np.log(im_value))
    )

def _get_hazard(
    ensemble: gm_data.Ensemble,
    station_name: str,
    im: IM,
    exceedance: float,
    ix: int,
    n_stations: int,
):
    """Computes the ensemble hazard curve for the specific station"""
    start_time = time.time()
    site_info = site.get_site_from_name(ensemble, station_name)
    im_value = run_ensemble_hazard(ensemble, site_info, im).exceedance_to_im(exceedance)

    print(
        f"Progress {ix}/{n_stations} - station {station_name} "
        f"- {time.time() - start_time}"
    )
    return im_value

def vs30_update(site_info: site.SiteInfo, hazard_result: BranchHazardResult):
    """Computes the updated hazard for the user specified vs30 value

    Parameters
    ----------
    site_info: SiteInfo
        Site of interest
    hazard_result: HazardResult
        The hazard result for the db vs30 value

    Returns
    -------
    flt_upd_hazard: pd.Series
        The fault based updated hazard
        format: index = IM values, values = exceedance probability
    ds_upd_hazard: pd.DataFrame
        The distributed seismicity based updated hazard
        format: index = IM values, values = exceedance probability
    """
    import empirical.util.classdef as classdef
    import empirical.util.empirical_factory as emp_factory

    branch, ensemble = hazard_result.branch, hazard_result.im_ensemble.ensemble
    im_ensemble = hazard_result.im_ensemble

    # Get IM of interest and the IM values of interest
    im = hazard_result.im
    im_values = hazard_result.im_values

    # Get the recurrance & gm prob dfs
    rec_prob = branch.rupture_df_id["annual_rec_prob"]
    flt_gm_prob_df = shared.get_gm_prob_df(
        branch, site_info, im, im_values, const.SourceType.fault, ensemble=ensemble
    )
    ds_gm_prob_df = shared.get_gm_prob_df(
        branch,
        site_info,
        im,
        im_values,
        const.SourceType.distributed,
        ensemble=ensemble,
    )

    # Get the fault and distributed disagg and combine to get the full disagg
    flt_disagg = sha_calc.disagg_exceedance_multi(
        flt_gm_prob_df, rec_prob, hazard_result.total_hazard
    )
    ds_disagg = sha_calc.disagg_exceedance_multi(
        ds_gm_prob_df, rec_prob, hazard_result.total_hazard
    )
    full_disagg = pd.concat([flt_disagg, ds_disagg])

    # Create distance lookup
    flt_distance_df = site_source.get_distance_df(ensemble.flt_ssddb_ffp, site_info)
    ds_distance_df = site_source.get_distance_df(ensemble.ds_ssddb_ffp, site_info)
    distance_lookup_df = pd.concat([flt_distance_df, ds_distance_df])

    # Create a rupture_id to location name lookup, since the data from the
    # the site-source db uses location names and not rupture ids
    flt_loc_names = site_source.rupture_id_to_loc_name(
        flt_disagg.index.values, const.SourceType.fault
    )
    ds_loc_names = site_source.rupture_id_to_loc_name(
        ds_disagg.index.values, const.SourceType.distributed
    )
    loc_names_lookup = pd.concat([flt_loc_names, ds_loc_names])

    # Compute the mean magnitude
    ruptures = full_disagg.index.values
    flt_ruptures = flt_disagg.index.values
    ds_ruptures = ds_disagg.index.values
    flt_mag_mean_df = shared.compute_contr_mean(
        im_ensemble.rupture_df_id.magnitude.loc[flt_ruptures], full_disagg.loc[flt_ruptures]
    )
    ds_mag_mean_df = shared.compute_contr_mean(
        im_ensemble.rupture_df_id.magnitude.loc[ds_ruptures], full_disagg.loc[ds_ruptures]
    )

    # Create a distance dataframe for the ruptures of interest
    # Note: Have to use reindex since there might be ruptures for which
    # there is no site-source data for the current station (reindex just sets those to nan)
    distance_df = distance_lookup_df.reindex(loc_names_lookup.loc[ruptures].values)
    distance_df = distance_df.set_index(ruptures)

    # Compute the mean rrup & rjb
    flt_rrup_mean_df = shared.compute_contr_mean(
        distance_df.rrup.loc[flt_ruptures], full_disagg.loc[flt_ruptures]
    )
    ds_rrup_mean_df = shared.compute_contr_mean(
        distance_df.rrup.loc[ds_ruptures], full_disagg.loc[ds_ruptures]
    )

    flt_rjb_mean_df = shared.compute_contr_mean(
        distance_df.rjb.loc[flt_ruptures], full_disagg.loc[flt_ruptures]
    )
    ds_rjb_mean_df = shared.compute_contr_mean(
        distance_df.rjb.loc[ds_ruptures], full_disagg.loc[ds_ruptures]
    )

    # Sanity check (can probably remove these at some point)
    assert np.all(np.isclose(flt_mag_mean_df.index.values, im_values))
    assert np.all(np.isclose(ds_mag_mean_df.index.values, im_values))
    assert np.all(np.isclose(flt_rrup_mean_df.index.values, im_values))
    assert np.all(np.isclose(ds_rrup_mean_df.index.values, im_values))
    assert np.all(np.isclose(flt_rjb_mean_df.index.values, im_values))
    assert np.all(np.isclose(ds_rjb_mean_df.index.values, im_values))

    # Compute the vs30 ratios for each IM value of the hazard data
    vs30_ratio, flt_vs30_ratio, ds_vs30_ratio = [], [], []
    for ix, im_value in enumerate(im_values):
        # Create the fault and site objects
        cur_flt_fault = classdef.Fault(
            Mw=flt_mag_mean_df.iloc[ix], rake=-90.0, dip=45.0, zbot=15.0, hdepth=5.0
        )
        cur_ds_fault = classdef.Fault(
            Mw=ds_mag_mean_df.iloc[ix], rake=-90.0, dip=45.0, zbot=15.0, hdepth=5.0
        )

        cur_flt_site_db = classdef.Site(
            rrup=float(flt_rrup_mean_df.iloc[ix]),
            rjb=float(flt_rjb_mean_df.iloc[ix]),
            rx=0,
            vs30=site_info.vs30,
        )
        cur_flt_site_user = classdef.Site(
            rrup=float(flt_rrup_mean_df.iloc[ix]),
            rjb=float(flt_rjb_mean_df.iloc[ix]),
            rx=0,
            vs30=site_info.user_vs30,
        )

        cur_ds_site_db = classdef.Site(
            rrup=float(ds_rrup_mean_df.iloc[ix]),
            rjb=float(ds_rjb_mean_df.iloc[ix]),
            rx=0,
            vs30=site_info.vs30,
        )
        cur_ds_site_user = classdef.Site(
            rrup=float(ds_rrup_mean_df.iloc[ix]),
            rjb=float(ds_rjb_mean_df.iloc[ix]),
            rx=0,
            vs30=site_info.user_vs30,
        )

        # Run the empirical model for using the db and user specified vs30
        flt_im_db, _ = emp_factory.compute_gmm(
            cur_flt_fault,
            cur_flt_site_db,
            classdef.GMM.CB_14,
            str(im),
            period=im.period,
        )
        flt_im_user, _ = emp_factory.compute_gmm(
            cur_flt_fault,
            cur_flt_site_user,
            classdef.GMM.CB_14,
            str(im),
            period=im.period,
        )

        ds_im_db, _ = emp_factory.compute_gmm(
            cur_ds_fault, cur_ds_site_db, classdef.GMM.CB_14, str(im), period=im.period
        )
        ds_im_user, _ = emp_factory.compute_gmm(
            cur_ds_fault,
            cur_ds_site_user,
            classdef.GMM.CB_14,
            str(im),
            period=im.period,
        )

        # Compute the vs30 ratio
        flt_vs30_ratio.append(flt_im_user / flt_im_db)
        ds_vs30_ratio.append(ds_im_user / ds_im_db)

    # Compute the updated IM values
    flt_vs30_updated = im_values * np.asarray(flt_vs30_ratio)
    ds_vs30_updated = im_values * np.asarray(ds_vs30_ratio)

    # Interpolate to return data at the same IM levels
    flt_mask = ~np.isnan(flt_vs30_updated)
    flt_vs30_updated_excd = np.interp(
        im_values,
        flt_vs30_updated[flt_mask],
        hazard_result.fault_hazard.values[flt_mask],
        right=0.0,
    )
    flt_upd_hazard = pd.Series(index=im_values, data=flt_vs30_updated_excd)

    ds_mask = ~np.isnan(ds_vs30_updated)
    ds_vs30_updated_excd = np.interp(
        im_values,
        ds_vs30_updated[ds_mask],
        hazard_result.ds_hazard.values[ds_mask],
        right=0.0,
    )
    ds_upd_hazard = pd.Series(index=im_values, data=ds_vs30_updated_excd)

    return flt_upd_hazard, ds_upd_hazard