###############################################################################
# apogee.spec.plot: various way to plot APOGEE spectra
###############################################################################
from functools import wraps
import numpy
from galpy.util import bovy_plot
from matplotlib import pyplot
import matplotlib.ticker as ticker
from matplotlib.ticker import NullFormatter
from matplotlib.backends.backend_pdf import PdfPages
import apogee.spec.window as apwindow
import apogee.tools.read as apread
from apogee.tools import air2vac, atomic_number,apStarWavegrid,_apStarPixelLimits,_aspcapPixelLimits

_LAMBDASUB= 15000
_STARTENDSKIP= 30
# Good, clean Lines, mainly from Smith et al. (2013)
line_labels= {}
line_labels['fe']= r'$\mathrm{Fe\kern 0.1em I}$'
line_labels['mg']= r'$\mathrm{Mg\kern 0.1em I}$'
line_labels['al']= r'$\mathrm{Al\kern 0.1em I}$'
line_labels['si']= r'$\mathrm{Si\kern 0.1em I}$'
line_labels['k']= r'$\mathrm{K\kern 0.1em I}$'
line_labels['ca']= r'$\mathrm{Ca\kern 0.1em I}$'
line_labels['ti']= r'$\mathrm{Ti\kern 0.1em I}$'
line_labels['cr']= r'$\mathrm{Cr\kern 0.1em I}$'
line_labels['ni']= r'$\mathrm{Ni\kern 0.1em I}$'
line_labels['na']= r'$\mathrm{Na\kern 0.1em I}$'
line_labels['mn']= r'$\mathrm{Mn\kern 0.1em I}$'
line_labels['s']= r'$\mathrm{S\kern 0.1em I}$'
line_labels['v']= r'$\mathrm{V\kern 0.1em I}$'
line_labels['cob']= r'$\mathrm{Co\kern 0.1em I}$'
line_labels['cu']= r'$\mathrm{Cu\kern 0.1em I}$'
line_labels['oh']= r'$\mathrm{OH}$'
line_labels['co']= r'$^{12}\!\mathrm{CO}$'
line_labels['cn']= r'$\mathrm{CN}$'
line_labels['13co']= r'$^{13}\!\mathrm{CO}$'
line_labels['hbrpi']= r'$\mathrm{Br-}\pi$'
line_labels['hbrla']= r'$\mathrm{Br-}\lambda$'
line_labels['hbr']= r'$\mathrm{H[Br]}$'
line_labels['dib']= r'$\mathrm{DIB}$'
# From Table 2 in Smith et al. (2013)
_FEI_lines= [air2vac(l) for l in [15194.492,15207.526,15395.718,15490.339,
                                  15648.510,15964.867,16040.657,16153.247,
                                  16165.032]]
_FEI_lines.append(16697.635) # one more from Shetrone
# From Table 5
_MGI_lines= [air2vac(l) for l in [15740.716,15748.9,15765.8,15879.5,
                                  15886.2,15954.477]]
_ALI_lines= [air2vac(l) for l in [16718.957,16750.564286,16763.359]]
_SII_lines= [air2vac(l) for l in [15361.161,15376.831,15833.602,15960.063,
                                  16060.009,16094.787,16215.670,16680.770,
                                  16828.159]]
_KI_lines= [air2vac(l) for l in [15163.067,15168.376]]
_CAI_lines= [air2vac(l) for l in [16136.823,16150.763,16155.236,16157.364]]
_TII_lines= [air2vac(l) for l in [15543.756,15602.842,15698.979,15715.573,
                                  16635.161]]
_VI_lines= [air2vac(15925.)]
_CRI_lines= [air2vac(l) for l in [15680.063,15860.214]]
_MNI_lines= [air2vac(l) for l in [15159.,15217.85,15262.4]]
_COI_lines= [air2vac(16757.7)]
_NII_lines= [air2vac(l) for l in [15605.680,16584.439,16589.295,
                                  16673.711,16815.471,16818.760]]
_CUI_lines= [air2vac(16005.7)]
# From Katia Cunha
_NAI_lines= [air2vac(16388.85)]
# From Matthew Shetrone
_SI_lines= [15406.540,15426.490,15474.043,15482.712]
# From Table 4 in Smith et al. (2013), with a few tweaks
_OH_lines= [air2vac(l) for l in [15279.5,15391.,15505.5,15570.5]]
_CO_lines= [air2vac(l) for l in [15582.,15780.5,15988.,16189.5]]
_CN_lines= [air2vac(l) for l in [15260.,15321.,15397.,15332.,15410.8,
                                 15447.,15466.,15472.,15482.]]
_13CO_lines= [air2vac(l) for l in [16122.5,16743.5]]
#The hydrogen bracket series
_HBRPI_lines= [15196.005]
_HBRLA_lines= [15704.960]
_HBR_lines= [15004.970,15043.157,15086.906,15137.367,15264.717,
             15345.992,15443.148,15560.708,15884.888,16113.721,
             16411.681,16811.117]

def specPlotInputDecorator(func):
    """Decorator to parse input to spectral plotting"""
    @wraps(func)
    def input_wrapper(*args,**kwargs):
        if len(args) >= 2 and isinstance(args[0],(list,numpy.ndarray)) \
                and isinstance(args[1],(list,numpy.ndarray)):
            # wavelength, spectrum
            return func(args[0],args[1],*args[2:],**kwargs)
        elif len(args) >= 1 and isinstance(args[0],(list,numpy.ndarray)):
            # spectrum on standard re-sampled wavelength grid
            lam=apStarWavegrid()
            apStarBlu_lo,apStarBlu_hi,apStarGre_lo,apStarGre_hi,apStarRed_lo,apStarRed_hi = _apStarPixelLimits(dr=None)    
            aspcapBlu_start,aspcapGre_start,aspcapRed_start,aspcapTotal = _aspcapPixelLimits(dr=None)
            if len(args[0]) == aspcapTotal: # Input is on ASPCAP grid
                spec= numpy.zeros(len(lam))
                spec[apStarBlu_lo:apStarBlu_hi]= args[0][:aspcapGre_start]
                spec[apStarGre_lo:apStarGre_hi]= args[0][aspcapGre_start:aspcapRed_start]
                spec[apStarRed_lo:apStarRed_hi]= args[0][aspcapRed_start:]
            else:
                spec= args[0]
            return func(lam,spec,*args[1:],**kwargs)
        elif isinstance(args[0],(int,numpy.short,str)) \
                and isinstance(args[1],str):
            # location ID and APOGEE ID (loc ID can be string for 1m sample)
            if kwargs.get('apStar',False):
                spec, hdr= apread.apStar(args[0],args[1],header=True,
                                         ext=kwargs.pop('ext',1))
                spec= spec[numpy.amin([kwargs.pop('apStarIndx',1),
                                       len(spec)-1])]
            else: #aspcapStar
                spec, hdr= apread.aspcapStar(args[0],args[1],header=True,
                                             ext=kwargs.pop('ext',1))
            lam= 10.**numpy.arange(hdr['CRVAL1'],
                                   hdr['CRVAL1']+len(spec)*hdr['CDELT1'],
                                   hdr['CDELT1'])
            return func(lam,spec,*args[2:],**kwargs)
    return input_wrapper

@specPlotInputDecorator
def waveregions(*args,**kwargs):
    """
    NAME:
       waveregions
    PURPOSE:
       plot selected regions of the spectrum in one row
    INPUT:
       Either:
          (a) wavelength, spectrum (\AA,spectrum units)
          (b) spectrum (assumed on standard APOGEE re-sampled wavelength grid)
          (c) location ID, APOGEE ID (default loads aspcapStar, loads extension ext(=1); apStar=True loads apStar spectrum)
    KEYWORDS:
       File loading:
          ext= (1) extension to load
          apStar= (False) if True, load the apStar spectrum
          apStarIndx= (1) index in the apStar spectrum to load
       Chunks position:
          startlams, endlams= start and end wavelength in \AA of the various chunks (takes precedence over startindxs, endindxs)
          startindxs, endindxs= star and end index in the wavelength array of the various chunks
       Plotting-specific keywords
          labelLines= (True) label some lines
          noMolecLines= (False) don't label the molecules
          cleanZero= (True) replace <= zero entries with NaN
          labelID= A string ID that will be placed in the top-left corner
          labelTeff, labellogg, labelmetals, labelafe= parameter labels that will be placed in the top-right corner
          noxlabel= (False) if True, don't label the x axis
          pyplot.plot args and kwargs
    OUTPUT:
       plot to output
       The final axes allow one to put additional labels on the plot, e.g., for adding the APOGEE ID:
       bovy_plot.bovy_text(r'$\mathrm{%s}$' % '2M02420597+0837017',top_left=True)       
       Note that an ID (e.g., the apogee ID) and Teff, logg, metallicity, and alpha-enhancement labels can be added using the keywords label* above
    HISTORY:
       2015-01-18 - Written (based on older code) - Bovy (IAS)
    """
    # Grab non-pyplot.plot kwargs
    apStar= kwargs.pop('apStar',False)
    labelLines= kwargs.pop('labelLines',not 'overplot' in kwargs)
    noMolecLines= kwargs.pop('noMolecLines',False)
    cleanZero= kwargs.pop('cleanZero',True)
    noxticks= kwargs.pop('_noxticks',False)
    noxlabel= kwargs.pop('noxlabel',False)
    noskipdiags= kwargs.pop('_noskipdiags',False)
    labelwav= kwargs.pop('_labelwav',False)
    plotw= kwargs.pop('_plotw',None)
    markLines= kwargs.pop('markLines',False)
    markwav= kwargs.pop('_markwav',None)
    # Labels
    labelID= kwargs.pop('labelID',None)
    labelTeff= kwargs.pop('labelTeff',None)
    labellogg= kwargs.pop('labellogg',None)
    labelmetals= kwargs.pop('labelmetals',None)
    labelafe= kwargs.pop('labelafe',None)
    # Clean bad lines
    if cleanZero:
        args[1][args[1] <= 0.]= numpy.nan
    # Chunk parameters
    if 'startlams' in kwargs:
        # Turn startlams into a startindxs and similar for endlams
        startlams= kwargs.pop('startlams')
        endlams= kwargs.pop('endlams')
        startindxs= []
        endindxs= []
        for ii in range(len(startlams)):
            startindxs.append(numpy.argmin(numpy.fabs(startlams[ii]-args[0])))
            endindxs.append(numpy.argmin(numpy.fabs(endlams[ii]-args[0])))
    else:
        startindxs= kwargs.pop('startindxs',
                               [322,1794,2707,3850,4740,5820,7185])
        endindxs= kwargs.pop('endindxs',
                             [590,1940,2857,4025,5070,5955,7400])
    nregions= len(startindxs)
    # Calculate the width of the plot
    dx= numpy.array([args[0][numpy.amin([len(args[0])-1,endindxs[ii]])]\
                         -args[0][numpy.amax([0,startindxs[ii]-1])] \
                         for ii in range(nregions)],
                    dtype='float')
    # Adjust 0 (and -1) to start (end) a little further
    startendskip= kwargs.pop('_startendskip',_STARTENDSKIP)
    dx[0]= args[0][numpy.amin([len(args[0])-1,endindxs[0]])]\
        -args[0][numpy.amax([0,startindxs[0]-startendskip])] 
    dx[-1]= args[0][numpy.amin([len(args[0])-1,endindxs[-1]+startendskip])]\
        -args[0][numpy.amax([0,startindxs[-1]-1])] 
    if nregions == 1: #special case
        dx= [args[0][numpy.amin([len(args[0])-1,endindxs[0]+startendskip])]\
                 -args[0][numpy.amax([0,startindxs[0]-startendskip])]]
    # Determine a good step for the tickmarks
    tickStepTmp= numpy.log10(numpy.sum(dx)/10.) % 1
    if tickStepTmp > numpy.log10(1.5) and tickStepTmp < numpy.log10(3.5):
        tickStep= 2.*10.**int(numpy.log10(numpy.sum(dx)/10.))
    elif tickStepTmp > numpy.log10(3.5) and tickStepTmp < numpy.log10(7.5):
        tickStep= 5.*10.**int(numpy.log10(numpy.sum(dx)/10.))
    else:
        tickStep= 10.**int(numpy.log10(numpy.sum(dx)/10.))
    dx/= numpy.sum(dx)
    if noxticks:
        totdx= 0.825
    else:
        totdx= 0.85
    skipdx= kwargs.pop('skipdx',0.015)
    dx*= (totdx-(nregions-1)*skipdx)
    # Setup plot
    overplot= kwargs.pop('overplot',False)
    if not overplot:
        bovy_plot.bovy_print(fig_width=kwargs.pop('fig_width',8.4),
                             fig_height=kwargs.pop('fig_height',2.5),
                             axes_labelsize=16,text_fontsize=14,
                             legend_fontsize=12,
                             xtick_labelsize=12,ytick_labelsize=12)
        pyplot.figure()
    if overplot:
        yrange= numpy.array(pyplot.gca().get_ylim())
        kwargs.pop('yrange',None) # pop if there
    elif apStar:
        yrange= kwargs.pop('yrange',[0.,1.1*numpy.nanmax(args[1])])
    else:
        yrange= kwargs.pop('yrange',[0.2,1.2])
    # Deal with the label
    if apStar:
        ylabel= kwargs.pop('ylabel',r'$f_\lambda(\lambda)\,(10^{-17}\,\mathrm{erg\, s}^{-1}\,\mathrm{cm}^{-2}\,\AA^{-1})$')
    else:
        ylabel= kwargs.pop('ylabel',r'$f/f_c(\lambda)$')
    kwargs['zorder']= kwargs.get('zorder',10)
    for ii in range(nregions):
        # Setup the axes
        if ii == 0:
            left, bottom, width, height= 0.1+(0.85-totdx)*2., 0.125, dx[ii],0.8
        else:
            left, bottom, width, height= 0.1+(0.85-totdx)*2.+numpy.cumsum(dx)[ii-1]+skipdx*ii,\
                0.125, dx[ii], 0.8
        thisax= pyplot.axes([left,bottom,width,height])
        fig= pyplot.gcf()
        fig.sca(thisax)
        startindx, endindx= startindxs[ii], endindxs[ii]
        if ii == 0 and nregions == 1:
            xrange=[args[0][numpy.amax([0,startindx-startendskip])]-_LAMBDASUB,
                    args[0][numpy.amin([len(args[0])-1,endindx+startendskip])]-_LAMBDASUB]
        elif ii == 0:
            xrange=[args[0][numpy.amax([0,startindx-startendskip])]-_LAMBDASUB,
                    args[0][numpy.amin([len(args[0])-1,endindx])]-_LAMBDASUB]
        elif ii == (nregions-1):
            xrange=[args[0][numpy.amax([0,startindx-1])]-_LAMBDASUB,
                    args[0][numpy.amin([len(args[0])-1,endindx+startendskip])]-_LAMBDASUB]
        else:
            xrange=[args[0][numpy.amax([0,startindx-1])]-_LAMBDASUB,
                    args[0][numpy.amin([len(args[0])-1,endindx])]-_LAMBDASUB]
        thisax.plot(args[0][startindx:endindx]-_LAMBDASUB,
                    args[1][startindx:endindx],
                    *args[2:],**kwargs)
        if not plotw is None:
            thisax.plot(args[0][startindx:endindx]-_LAMBDASUB,
                        plotw[startindx:endindx],
                        '-',lw=2.,color='0.65',zorder=1)
        thisax.set_xlim(xrange[0],xrange[1])
        thisax.set_ylim(yrange[0],yrange[1])
        if noxticks:
            nullfmtx= NullFormatter()         # no labels, assume 1\AA
            thisax.xaxis.set_major_formatter(nullfmtx)
            thisax.xaxis.set_major_locator(ticker.MultipleLocator(2.))
        else:
            thisax.xaxis.set_major_locator(ticker.MultipleLocator(tickStep))
        bovy_plot._add_ticks(xticks=True^noxticks)
        if ii > 0:
            nullfmt   = NullFormatter()         # no labels
            thisax.yaxis.set_major_formatter(nullfmt)
        elif not overplot:
            pyplot.ylabel(ylabel)
        # Remove spines between different wavelength regions
        if ii == 0 and not nregions == 1:
            thisax.spines['right'].set_visible(False)
            thisax.tick_params(right=False,which='both')
        elif ii == (nregions-1) and not nregions == 1:
            thisax.spines['left'].set_visible(False)
            thisax.tick_params(labelleft='off')
            thisax.tick_params(left=False,which='both')
        elif not nregions == 1:
            thisax.spines['left'].set_visible(False)
            thisax.spines['right'].set_visible(False)
            thisax.tick_params(labelleft='off')
            thisax.tick_params(left=False,which='both')
            thisax.tick_params(right=False,which='both')
        # Plot cut-out markers
        cutOutkwargs = dict(transform=thisax.transAxes,color='k',
                            clip_on=False)
        if not noskipdiags:
            d = .015 # how big to make the diagonal lines in axes coordinates
            slope= 1./(dx[ii]+0.2*skipdx)/3.
            if ii == 0 and not nregions == 1:
                thisax.plot((1-slope*d,1+slope*d),(-d,+d), **cutOutkwargs)
                thisax.plot((1-slope*d,1+slope*d),(1-d,1+d), **cutOutkwargs)
            elif ii == (nregions-1) and not nregions == 1:
                thisax.plot((-slope*d,+slope*d),(-d,+d), **cutOutkwargs)
                thisax.plot((-slope*d,+slope*d),(1-d,1+d), **cutOutkwargs)
            elif not nregions == 1:
                thisax.plot((1-slope*d,1+slope*d),(-d,+d), **cutOutkwargs)
                thisax.plot((1-slope*d,1+slope*d),(1-d,1+d), **cutOutkwargs)
                thisax.plot((-slope*d,+slope*d),(-d,+d), **cutOutkwargs)
                thisax.plot((-slope*d,+slope*d),(1-d,1+d), **cutOutkwargs)