import numpy as np
import os

from astropy.table import Table
from . import utils

__all__ = ["FilterDefinition", "FilterFile", "ParamFilter"]

VEGA_FILE = os.path.join(utils.path_to_eazy_data(),
                         'alpha_lyr_stis_008.fits')
                         
VEGA = Table.read(VEGA_FILE)
for c in VEGA.colnames:
    VEGA[c] = VEGA[c].astype(float)
    
class FilterDefinition:
    def __init__(self, name=None, wave=None, throughput=None, bp=None):
        """
        Bandpass object
        
        Parameters
        ----------
        name : str
            Label name
        
        wave : array
            Wavelength array, in `astropy.units.Angstrom`.
        
        throughput : array
            Throughput, arbitrary normalization
        
        bp : optional, `pysynphot.obsbandpass` object
            `pysynphot` filter bandpass
        
        """
        self.name = name
        self.wave = wave
        self.throughput = throughput
        self.Aflux = 1.
        
        # pysynphot Bandpass
        if bp is not None:
            self.wave = np.cast[np.double](bp.wave)
            self.throughput =  np.cast[np.double](bp.throughput)
            self.name = bp.name
                    
        self.norm = 1.
        if self.throughput is not None:
            self.norm = np.trapz(self.throughput/self.wave, self.wave)


    def __repr__(self):
        return self.name.__repr__()


    def __str__(self):
        return self.name.__str__()


    def get_extinction(self, EBV=0, Rv=3.1):
        """
        Extinction factor 
        """
        import astropy.units as u
        
        f99 = utils.GalacticExtinction(EBV=EBV, Rv=Rv)
        self.Alambda = f99(self.wave)
        self.Aflux = 10**(-0.4*self.Alambda)


    def extinction_correction(self, EBV, Rv=3.1, mag=True, source_lam=None, source_flux=None):
        """
        Get the MW extinction correction within the filter.  
        
        Optionally supply a source spectrum.
        """
        import astropy.units as u
        try:
            import grizli.utils_c
            interp = grizli.utils_c.interp.interp_conserve_c
        except ImportError:
            interp = utils.interp_conserve
             
        if self.wave is None:
            print('Filter not defined.')
            return False
        
        if source_flux is None:
            source_flux = self.throughput*0.+1
        else:
            source_flux = interp(self.wave, source_lam, source_flux, left=0, right=0)
        
        if (self.wave.min() < 910) | (self.wave.max() > 6.e4):
            Alambda = 0.
        else:
            f99 = utils.GalacticExtinction(EBV=EBV, Rv=Rv)
            Alambda = f99(self.wave)
                         
        delta = np.trapz(self.throughput*source_flux*10**(-0.4*Alambda), self.wave) / np.trapz(self.throughput*source_flux, self.wave)
        
        if mag:
            return 2.5*np.log10(delta)
        else:
            return 1./delta


    @property    
    def ABVega(self):
        """
        Compute AB-Vega conversion
        """
        from astropy.constants import c
        import astropy.units as u
        try:
            import grizli.utils_c
            interp = grizli.utils_c.interp.interp_conserve_c
        except ImportError:
            interp = utils.interp_conserve
        
        # Union of throughput and Vega spectrum arrays
        full_x = np.hstack([self.wave, VEGA['WAVELENGTH']])
        full_x = full_x[np.argsort(full_x)]

        # Vega spectrum, units of f-lambda flux density, cgs
        # Interpolate to wavelength grid, no extrapolation
        vega_full = interp(full_x, VEGA['WAVELENGTH'], VEGA['FLUX'], 
                              left=0, right=0)
                              
        thru_full = interp(full_x, self.wave, self.throughput, 
                              left=0, right=0)        
        
        # AB = 0, same units
        absp = 3631*1e-23*c.to(u.m/u.s).value*1.e10/full_x**2
        
        # Integrate over the bandpass, flam dlam
        num = np.trapz(vega_full*thru_full, full_x)
        den = np.trapz(absp*thru_full, full_x)
        
        return -2.5*np.log10(num/den)


    @property    
    def pivot(self):
        """
        Pivot wavelength
        
        http://pysynphot.readthedocs.io/en/latest/properties.html
        """
        integrator = np.trapz
        
        num = integrator(self.wave, self.wave*self.throughput)
        den = integrator(self.wave, self.throughput/self.wave)
        pivot = np.sqrt(num/den)
        return pivot
    
    @property    
    def equivwidth(self):
        """
        Filter equivalent width

        http://pysynphot.readthedocs.io/en/latest/properties.html
        """
        return np.trapz(self.throughput, self.wave)
    
    
    @property            
    def rectwidth(self):
        """
        Filter rectangular width

        http://pysynphot.readthedocs.io/en/latest/properties.html
        """
        
        rect = self.equivwidth / self.throughput.max()
        return rect
    
    
    @property    
    def ctw95(self):
        """
        95% cumulative throughput width
        
        http://www.stsci.edu/hst/acs/analysis/bandwidths/#keywords
        
        """
        
        dl = np.diff(self.wave)
        filt = np.cumsum((self.wave*self.throughput)[1:]*dl)
        ctw95 = np.interp([0.025, 0.975], filt/filt.max(), self.wave[1:])
        return np.diff(ctw95)[0]
    
    
    def for_filter_file(self, row_str='{i:6} {wave:.5e} {thru:.5e}'):
        """
        Return a string that can be put in the EAZY filter file
        """    
        header = '{0} {1} lambda_c= {2:.4e} AB-Vega= {3:.3f} w95={4:.1f}'
        N = len(self.wave)
        lines = [header.format(N, self.name.split('lambda_c')[0], 
                               self.pivot, self.ABVega, self.ctw95)]
        
        lines += [row_str.format(i=i+1, wave=w, thru=t)
                  for i, (w, t) in enumerate(zip(self.wave, self.throughput))]
        
        return '\n'.join(lines)


class FilterFile:
    def __init__(self, file='FILTER.RES.latest', path='./'):
        """
        Read a EAZY filter file.
        
        .. plot::
            :include-source:
        
            import matplotlib.pyplot as plt
            from eazy.filters import FilterFile
            
            res = FilterFile(path=None)
            print(len(res.filters))
            
            bp = res[205]
            print(bp)
            
            fig, ax = plt.subplots(1,1,figsize=(6,4))
            
            ax.plot(bp.wave, bp.throughput, label=bp.name.split()[0])
            
            ax.set_xlabel('wavelength, Angstroms')
            ax.set_ylabel('throughput')
            ax.legend()
            ax.grid()
            
            fig.tight_layout(pad=0.5)
            
        
        """
        if path is None:
            file_path = os.path.join(os.getenv('EAZYCODE'), 'filters', file)
        else:
            file_path = os.path.join(path, file)
            
        with open(file_path, 'r') as fp:
            lines = fp.readlines()
        
        self.filename = file_path
        
        filters = []
        wave = []
        trans = []
        header = ''
        
        for line in lines:
            if 'lambda_c' in line:
                if len(wave) > 0:
                    # Make filter from lines already read in
                    new_filter = FilterDefinition(name=header,
                                                  wave=np.cast[float](wave), 
                                            throughput=np.cast[float](trans))
                    # new_filter.name = header
                    # new_filter.wave = np.cast[float](wave)
                    # new_filter.throughput = np.cast[float](trans)
                    filters.append(new_filter)

                # Initialize filter
                header = ' '.join(line.split()[1:])
                wave = []
                trans = []
            else:
                lspl = np.cast[float](line.split())
                wave.append(lspl[1])
                trans.append(lspl[2])
                
        # last one
        # new_filter = FilterDefinition()
        # new_filter.name = header
        # new_filter.wave = np.cast[float](wave)
        # new_filter.throughput = np.cast[float](trans)
        new_filter = FilterDefinition(name=header,
                                      wave=np.cast[float](wave), 
                                throughput=np.cast[float](trans))

        filters.append(new_filter)
           
        self.filters = filters


    @property 
    def NFILT(self):
        """
        Number of filters in the list
        """
        return len(self.filters)


    def __getitem__(self, i1):
        """
        Return unit-indexed filter, e.g., 161 = 2mass-j
        """
        return self.filters[i1-1]


    def names(self, verbose=True):
        """
        Print the filter names.
        """
        if verbose:
            for i in range(len(self.filters)):
                print('{0:5d} {1}'.format(i+1, self.filters[i].name))
        else:
            string_list = ['{0:5d} {1}\n'.format(i+1, self.filters[i].name) for i in range(len(self.filters))]
            return string_list


    def write(self, file='xxx.res', verbose=True):
        """
        Dump the filter information to a filter file.
        """
        fp = open(file,'w')
        for filter in self.filters:
            fp.write('{0:6d} {1}\n'.format(len(filter.wave), filter.name))
            for i in range(len(filter.wave)):
                fp.write('{0:6d} {1:.5e} {2:.5e}\n'.format(i+1, filter.wave[i], filter.throughput[i]))
        
        fp.close()
        
        string_list = self.names(verbose=False)
        fp = open(file+'.info', 'w')
        fp.writelines(string_list)
        fp.close()
        
        if verbose:
            print('Wrote <{0}[.info]>'.format(file))


    def search(self, search_string, case=False, verbose=True):
        """ 
        Search filter names for ``search_string``.  If ``case`` is True, then
        match case.
        """
        import re
        
        if not case:
            search_string = search_string.upper()
        
        matched = []
        
        for i in range(len(self.filters)):
            filt_name = self.filters[i].name
            if not case:
                filt_name = filt_name.upper()
                
            if re.search(search_string, filt_name) is not None:
                if verbose:
                    print('{0:5d} {1}'.format(i+1, self.filters[i].name))
                matched.append(i)
        
        return np.array(matched)


class ParamFilter(FilterDefinition):
    def __init__(self, line='#  Filter #20, RES#78: COSMOS/SUBARU_filter_B.txt - lambda_c=4458.276253'):
        
        self.lambda_c = float(line.split('lambda_c=')[1])
        self.name = line.split()[4]
        self.fnumber = int(line.split('RES#')[1].split(':')[0])
        self.cnumber = int(line.split('Filter #')[1].split(',')[0])

[((12, 7, 12, 28), 'astropy.table.Table.read', 'Table.read', ({(12, 18, 12, 27): 'VEGA_FILE'}, {}), '(VEGA_FILE)', False, 'from astropy.table import Table\n'), ((121, 17, 121, 59), 'numpy.hstack', 'np.hstack', ({(121, 27, 121, 58): "[self.wave, VEGA['WAVELENGTH']]"}, {}), "([self.wave, VEGA['WAVELENGTH']])", True, 'import numpy as np\n'), ((136, 14, 136, 51), 'numpy.trapz', 'np.trapz', ({(136, 23, 136, 42): 'vega_full * thru_full', (136, 44, 136, 50): 'full_x'}, {}), '(vega_full * thru_full, full_x)', True, 'import numpy as np\n'), ((137, 14, 137, 46), 'numpy.trapz', 'np.trapz', ({(137, 23, 137, 37): 'absp * thru_full', (137, 39, 137, 45): 'full_x'}, {}), '(absp * thru_full, full_x)', True, 'import numpy as np\n'), ((153, 16, 153, 32), 'numpy.sqrt', 'np.sqrt', ({(153, 24, 153, 31): 'num / den'}, {}), '(num / den)', True, 'import numpy as np\n'), ((163, 15, 163, 51), 'numpy.trapz', 'np.trapz', ({(163, 24, 163, 39): 'self.throughput', (163, 41, 163, 50): 'self.wave'}, {}), '(self.throughput, self.wave)', True, 'import numpy as np\n'), ((187, 13, 187, 31), 'numpy.diff', 'np.diff', ({(187, 21, 187, 30): 'self.wave'}, {}), '(self.wave)', True, 'import numpy as np\n'), ((188, 15, 188, 60), 'numpy.cumsum', 'np.cumsum', ({(188, 25, 188, 59): '(self.wave * self.throughput)[1:] * dl'}, {}), '((self.wave * self.throughput)[1:] * dl)', True, 'import numpy as np\n'), ((358, 15, 358, 32), 'numpy.array', 'np.array', ({(358, 24, 358, 31): 'matched'}, {}), '(matched)', True, 'import numpy as np\n'), ((49, 24, 49, 70), 'numpy.trapz', 'np.trapz', ({(49, 33, 49, 58): 'self.throughput / self.wave', (49, 60, 49, 69): 'self.wave'}, {}), '(self.throughput / self.wave, self.wave)', True, 'import numpy as np\n'), ((99, 16, 99, 83), 'numpy.trapz', 'np.trapz', ({(99, 25, 99, 71): '(self.throughput * source_flux * 10 ** (-0.4 * Alambda))', (99, 73, 99, 82): 'self.wave'}, {}), '(self.throughput * source_flux * 10 ** (-0.4 * Alambda), self.wave)', True, 'import numpy as np\n'), ((99, 86, 99, 134), 'numpy.trapz', 'np.trapz', ({(99, 95, 99, 122): '(self.throughput * source_flux)', (99, 124, 99, 133): 'self.wave'}, {}), '(self.throughput * source_flux, self.wave)', True, 'import numpy as np\n'), ((139, 20, 139, 37), 'numpy.log10', 'np.log10', ({(139, 29, 139, 36): '(num / den)'}, {}), '(num / den)', True, 'import numpy as np\n'), ((190, 15, 190, 29), 'numpy.diff', 'np.diff', ({(190, 23, 190, 28): 'ctw95'}, {}), '(ctw95)', True, 'import numpy as np\n'), ((241, 24, 241, 48), 'os.path.join', 'os.path.join', ({(241, 37, 241, 41): 'path', (241, 43, 241, 47): 'file'}, {}), '(path, file)', False, 'import os\n'), ((102, 23, 102, 38), 'numpy.log10', 'np.log10', ({(102, 32, 102, 37): 'delta'}, {}), '(delta)', True, 'import numpy as np\n'), ((122, 24, 122, 42), 'numpy.argsort', 'np.argsort', ({(122, 35, 122, 41): 'full_x'}, {}), '(full_x)', True, 'import numpy as np\n'), ((239, 37, 239, 58), 'os.getenv', 'os.getenv', ({(239, 47, 239, 57): '"""EAZYCODE"""'}, {}), "('EAZYCODE')", False, 'import os\n'), ((353, 15, 353, 50), 're.search', 're.search', ({(353, 25, 353, 38): 'search_string', (353, 40, 353, 49): 'filt_name'}, {}), '(search_string, filt_name)', False, 'import re\n'), ((133, 26, 133, 39), 'astropy.constants.c.to', 'c.to', ({(133, 31, 133, 38): '(u.m / u.s)'}, {}), '(u.m / u.s)', False, 'from astropy.constants import c\n')]