"""
The following :class:`Filter` and :class:`Wheel` classes can be used to simulate
coronagraph observations in imaging mode.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
import os

class Filter(object):
    """Filter for telescope imaging mode.

    Parameters
    ----------
    name : string
        Name of filter
    bandcenter : float
        Wavelength at bandcenter (um)
    FWHM : float
        Full-width as half-maximum (um)
    wl : array
        Wavelength grid for response function (um)
    response : array
        Filter response function
    notes : string
        Notes of filter, e.g. 'Johnson-Cousins'
    """

    def __init__(self, name=None, bandcenter=None, FWHM=None, wl=None, response=None, notes=''):
        self.name=name
        self.bandcenter=bandcenter
        self.FWHM=FWHM
        self.wl=wl
        self.response=response
        self.notes=notes

    def __str__(self):
        string = 'Filter: \n------------\n'+\
            '- Name : '+"%s" % (self.name)+'\n'+\
            '- Bandcenter (um)  : '+"%s" % (self.bandcenter)+'\n'+\
            '- Full width at half max (um)  : '+"%s" % (self.FWHM)+'\n'+\
            '- Wavelength array (um)  : '+"%s" % (self.wl)+'\n'+\
            '- Response : '+"%s" % (self.response)+'\n'+\
            '- Notes  : '+"%s" % (self.notes)+'\n'
        return string

class Wheel(object):
    """Filter Wheel. Contains different filters as attributes.
    """
    def __init__(self):
        pass

    def add_new_filter(self, filt, name='new_filter'):
        """Adds new filter to wheel

        Parameters
        ----------
        filt : Filter
            New Filter object to be added to wheel
        name : string (optional)
            Name to give new filter attribute
        """
        setattr(self, name, filt)

    def plot(self, ax=None, ylim = None):
        """
        Plot the filter response functions

        Parameters
        ----------
        ax : matplotlib.axis (optional)
            Axis instance to plot on

        Returns
        -------
        ax : ``matplotlib.axis``
            Axis with plot

        Note
        ----
        Only returns an axis if an axis was not provided.
        """

        if ax == None:
            fig = plt.figure(figsize=(14,10))
            gs = gridspec.GridSpec(1,1)
            ax1 = plt.subplot(gs[0])
            ax1.set_ylabel(r"Filter Response")
            ax1.set_xlabel(r"Wavelength [$\mu$m]")
            #ax1.set_ylim([0.0,1.0])
        else:
            ax1 = ax.twinx()
            #ax1.set_ylim([0.0,10.0])
            ax1.axes.get_yaxis().set_visible(False)

        Nfilt = len(self.__dict__)
        #colors,scalarMap,cNorm = scalarmap(np.arange(Nfilt),cmap='Dark2')
        i = 0
        fmax = 1.0
        for attr, value in self.__dict__.items():
            if np.max(value.response) > fmax: fmax = np.max(value.response)
            wl, response = value.wl, value.response
            #ax1.plot(wl,response, lw=3.0, label=value.name, c=colors[i])
            ax1.fill_between(wl,response, color='purple', alpha=0.3)
            i += 1

        if ylim is None:
            ax1.set_ylim([0.0,fmax*10.])
        else:
            ax1.set_ylim(ylim)

        if ax==None:
            return ax1

    def __str__(self):
        string = []
        tdict = sorted(self.__dict__.items(), key=lambda x: x[1].bandcenter)
        for attr, value in tdict:
            string.append(attr)
        print(string)
        return ''#str(string)


def read_jc():
    """
    Read and parse the Johnson-Cousins filter files.

    Returns
    -------
    filters : `numpy.ndarray`
        Array of filter response functions
    filter_names : list
        List of string names for the filters
    bandcenters : `numpy.array`
        Wavelength bandcenters for the filters [microns]
    FWHM : `numpy.array`
        Full-width at half max for the filters
    """
    path = 'filters/UBVRI/'
    # set file path relative to this file
    path = os.path.join(os.path.dirname(__file__), path)
    U = np.genfromtxt(path+'bess-u.pass')
    U[:,0] = U[:,0]/1e4
    B = np.genfromtxt(path+'bess-b.pass')
    B[:,0] = B[:,0]/1e4
    V = np.genfromtxt(path+'bess-v.pass')
    V[:,0] = V[:,0]/1e4
    R = np.genfromtxt(path+'bess-r.pass')
    R[:,0] = R[:,0]/1e4
    I = np.genfromtxt(path+'bess-i.pass')
    I[:,0] = I[:,0]/1e4
    filters = np.array([U,B,V,R,I])
    filter_names = ['U','B','V','R','I']
    bandcenters = np.array([365.6, 435.3, 547.7, 634.9, 800.0]) / 1e3
    FWHM = np.array([34.0, 78.1, 99.1, 106.56, 289.2]) / 1e3
    return filters, filter_names, bandcenters, FWHM

class johnson_cousins(Wheel):
    """
    Instantiate a filter :class:`Wheel` with the Johnson-Cousins filters
    (U, B, V, R, I).

    Example
    -------
    >>> jc = cg.imager.johnson_cousins()
    >>> jc.plot(ylim = (0.0, 1.2))

    .. plot::
      :align: center

      import coronagraph as cg
      jc = cg.imager.johnson_cousins()
      import matplotlib.pyplot as plt
      cg.plot_setup.setup()
      jc.plot(ylim = (0.0, 1.2))
      plt.show()
    """
    def __init__(self):

        filters, filter_names, bandcenters, FWHM = read_jc()

        self.U=Filter(name='U', bandcenter=bandcenters[0], FWHM=FWHM[0], wl=filters[0][:,0], response=filters[0][:,1], notes='Johnson-Cousins')
        self.B=Filter(name='B', bandcenter=bandcenters[1], FWHM=FWHM[1], wl=filters[1][:,0], response=filters[1][:,1], notes='Johnson-Cousins')
        self.V=Filter(name='V', bandcenter=bandcenters[2], FWHM=FWHM[2], wl=filters[2][:,0], response=filters[2][:,1], notes='Johnson-Cousins')
        self.R=Filter(name='R', bandcenter=bandcenters[3], FWHM=FWHM[3], wl=filters[3][:,0], response=filters[3][:,1], notes='Johnson-Cousins')
        self.I=Filter(name='I', bandcenter=bandcenters[4], FWHM=FWHM[4], wl=filters[4][:,0], response=filters[4][:,1], notes='Johnson-Cousins')


def read_landsat():
    """
    Read and parse the LANDSAT filter files.

    Returns
    -------