import matplotlib
import matplotlib.pyplot as plt
import os
import pdb
import pickle
import copy
import scipy.signal
import scipy.interpolate
import numpy as np
from astropy.modeling import models, fitting
from astropy.nddata import CCDData, StdDevUncertainty
from astropy.io import ascii, fits
from astropy.convolution import convolve, Box1DKernel, Box2DKernel
import pyvista
from pyvista import image
from pyvista import tv
from tools import plots

ROOT = os.path.dirname(os.path.abspath(__file__)) + '/../../'

class SpecData(CCDData) :
    """ Class to include a wavelength array on top of CCDData, with simple read/write/plot methods
    """
    def __init__(self,data,wave=None) :
        if type(data) is str :
            hdulist=fits.open(data)
            self.meta = hdulist[0].header
            self.unit = hdulist[0].header['BUNIT']
            self.data = hdulist[1].data
            self.uncertainty = StdDevUncertainty(hdulist[2].data)
            self.mask = hdulist[3].data
            self.wave = hdulist[4].data
        elif type(data) is CCDData :
            self.unit = data.unit
            self.meta = data.meta
            self.data = data.data
            self.uncertainty = data.uncertainty
            self.mask = data.mask
            self.wave = wave
        else :
            print('Input must be a filename or CCDData object')

    def write(self,file,overwrite=True) :
        hdulist=fits.HDUList()
        hdulist.append(fits.PrimaryHDU(header=self.meta))
        hdulist.append(fits.ImageHDU(self.data))
        hdulist.append(fits.ImageHDU(self.uncertainty.array))
        hdulist.append(fits.ImageHDU(self.mask.astype(np.int16)))
        hdulist.append(fits.ImageHDU(self.wave))
        hdulist.writeto(file,overwrite=overwrite)

    def plot(self,ax,**kwargs) :
        for row in range(self.wave.shape[0]) :
            gd = np.where(self.mask[row,:] == False)[0]
            plots.plotl(ax,self.wave[row,gd],self.data[row,gd],**kwargs)

def get_wavecal(file) :
    """ load a wavecal object from disk file
    """
    with open(file,'rb') as wavecal :
        return pickle.load(wavecal)

class WaveCal() :
    """ Class for wavelength solutions
    """
    def __init__ (self,type='chebyshev',degree=2,ydegree=2,pix0=0,orders=[1]) :
        """ Initialize the wavecal object

            type : type of solution ('poly' or 'chebyshev')
            degree : polynomial degree for wavelength
            ydegree : polynomial degree for  y dimension
            pix0 : reference pixel
            orders : spectral order for each row
            spectrum : spectrum from which fit is derived
        """
        self.type = type
        self.degree = degree
        self.ydegree = ydegree
        self.pix0 = pix0
        self.orders = orders
        self.waves = None
        self.x = None
        self.y = None
        self.weights = None
        self.model = None
        self.ax = None

    def wave(self,pixels=None,image=None) :
        """ Wavelength from pixel using wavelength solution model

            pix : input pixel positions [x] or [y,x]
            image : for input image size [nrows,ncols], return wavelengths at all pixels
            returns wavelength
        """
        if pixels is not None :
            out=np.zeros(len(pixels[0]))
            for i,pixel in enumerate(pixels[0]) :
                if self.type.find('2D') > 0 :
                    order=self.orders[pixels[1][i]]
                    out[i]=self.model(pixel-self.pix0,pixels[1][i])/order
                else :
                    out[i]=self.model(pixel-self.pix0)/self.orders[0]
            return out
        else :
            out=np.zeros(image)
            cols=np.arange(out.shape[-1])
            if out.ndim == 2 :
                for row in range(out.shape[0]) :
                    rows=np.zeros(len(cols))+row
                    try : order = self.orders[row]
                    except : order=self.orders[0]
                    out[row,:] = self.model(cols-self.pix0,rows)/order
            else :
                out= self.model(cols-self.pix0)/self.orders[0]
            return out

    def getmod(self) :
        """ Return model for current attributes
        """

        if self.type == 'poly' :
            mod=models.Polynomial1D(degree=self.degree)
        elif self.type == 'chebyshev' :
            mod=models.Chebyshev1D(degree=self.degree)
        elif self.type == 'chebyshev2D' :
            sz=self.spectrum.data.shape
            mod=models.Chebyshev2D(x_degree=self.degree,y_degree=self.ydegree,
                                   x_domain=[0,sz[1]],y_domain=[0,sz[0]])
        else :
            raise ValueError('unknown fitting type: '+self.type)
            return
        return mod

    def fit(self,plot=True) :
        """ do a wavelength fit
        """
        print("doing wavelength fit")
        # set up fitter and model
        twod='2D' in self.type
        fitter=fitting.LinearLSQFitter()
        mod = self.getmod()

        if not hasattr(self,'ax') : self.ax = None
        if twod :
            nold=-1
            nbd=0
            while nbd != nold :
                nold=nbd
                self.model=fitter(mod,self.pix-self.pix0,self.y,self.waves*self.waves_order,weights=self.weights)
                diff=self.waves-self.wave(pixels=[self.pix,self.y])
                gd = np.where(self.weights > 0)[0]
                print('  rms: {:8.3f}'.format(diff[gd].std()))
                bd = np.where(abs(diff) > 3*diff.std())[0]
                nbd = len(bd)
                print('rejecting {:d} points from {:d} total: '.format(nbd,len(self.waves)))
                self.weights[bd] = 0.

            if self.ax is not None :
                self.ax[1].cla()
                scat=self.ax[1].scatter(self.waves,diff,marker='o',c=self.y,s=2)
                scat=self.ax[1].scatter(self.waves[bd],diff[bd],marker='o',c='r',s=2)
                xlim=self.ax[1].get_xlim()
                self.ax[1].set_ylim(diff.min()-0.5,diff.max()+0.5)
                self.ax[1].plot(xlim,[0,0],linestyle=':')
                self.ax[1].text(0.1,0.9,'rms: {:8.3f}'.format(diff[gd].std()),transform=self.ax[1].transAxes)
                cb_ax = self.fig.add_axes([0.94,0.05,0.02,0.4])
                cb = self.fig.colorbar(scat,cax=cb_ax)
                cb.ax.set_ylabel('Row')
                plt.draw()
                self.fig.canvas.draw_idle()
                input('  See 2D wavecal fit. Hit any key to continue....')

        else :
            self.model=fitter(mod,self.pix-self.pix0,self.waves*self.waves_order,weights=self.weights)
            diff=self.waves-self.wave(pixels=[self.pix])
            print('  rms: {:8.3f} Angstroms'.format(diff.std()))
            if self.ax is not None :
                # iterate allowing for interactive removal of points
                done = False
                ymax = self.ax[0].get_ylim()[1]
                while not done :

                    # do fit
                    gd=np.where(self.weights>0.)[0]
                    bd=np.where(self.weights<=0.)[0]
                    self.model=fitter(mod,self.pix[gd]-self.pix0,self.waves[gd]*self.waves_order[gd],weights=self.weights[gd])
                    diff=self.waves-self.wave(pixels=[self.pix])
                    print('  rms: {:8.3f} Anstroms'.format(diff[gd].std()))

                    # replot spectrum with new fit wavelength scale
                    self.ax[0].cla()
                    self.ax[0].plot(self.wave(image=self.spectrum.data.shape)[0,:],self.spectrum.data[0,:])
                    # plot residuals
                    self.ax[1].cla()
                    self.ax[1].plot(self.waves[gd],diff[gd],'go')
                    self.ax[1].text(0.1,0.9,'rms: {:8.3f} Angstroms'.format(diff[gd].std()),transform=self.ax[1].transAxes)
                    self.ax[1].set_xlabel('Wavelength')
                    self.ax[1].set_ylabel('obs wave - fit wave')
                    if len(bd) > 0 : self.ax[1].plot(self.waves[bd],diff[bd],'ro')
                    self.ax[1].set_ylim(diff[gd].min()-0.5,diff[gd].max()+0.5)
                    for i in range(len(self.pix)) :
                        self.ax[1].text(self.waves[i],diff[i],'{:2d}'.format(i),va='top',ha='center')
                        if self.weights[i] > 0 :
                            self.ax[0].plot([self.waves[i],self.waves[i]],[0,ymax],'g')
                        else :
                            self.ax[0].plot([self.waves[i],self.waves[i]],[0,ymax],'r')
                    plt.draw()

                    # get input from user on lines to remove
                    for i in range(len(self.pix)) :
                        print('{:3d}{:8.2f}{:8.2f}{:8.2f}{:8.2f}'.format(
                               i, self.pix[i], self.waves[i], diff[i], self.weights[i]))
                    i = input('  enter ID of line to remove (-n for all lines<n, +n for all lines>n, O for new degree, return to continue): ')
                    if i == '' :
                        done = True
                    elif i == 'O' :
                        print('  current degree of fit: {:d}'.format(self.degree))
                        self.degree = int(input('  enter new degree of fit: '))
                        mod = self.getmod()
                    elif '+' in i :
                        self.weights[int(i)+1:] = 0.
                    elif '-' in i :
                        self.weights[0:abs(int(i))] = 0.
                    elif int(i) >= 0 :
                        self.weights[int(i)] = 0.
                    else :
                        print('invalid input')

    def set_spectrum(self,spectrum) :
        """ Set spectrum used to derive fit
        """
        self.spectrum = np.atleast_2d(spectrum)

    def get_spectrum(self) :
        """ Set spectrum used to derive fit
        """
        return self.spectrum

    def identify(self,spectrum,file=None,wav=None,wref=None,disp=None,display=None,plot=None,rad=5,thresh=10,
                 xmin=None, xmax=None, lags=range(-300,300), nskip=1) :
        """ Given some estimate of wavelength solution and file with lines,
            identify peaks and centroid
        """

        sz=spectrum.shape
        if len(sz) == 1 :
            spectrum.data = np.atleast_2d(spectrum.data)
            spectrum.uncertainty.array = np.atleast_2d(spectrum.uncertainty.array)
            sz=spectrum.shape
        if xmin is None : xmin=0
        if xmax is None : xmax=sz[-1]
        nrow=sz[0]

        # get initial reference wavelengths if not given
        if wav is None :
            pix=np.arange(sz[-1])
            if self.spectrum is not None :
                # cross correlate with reference image to get pixel shift
                print('  cross correlating with reference spectrum using lags: ', lags)
                fitpeak,shift = image.xcorr(self.spectrum.data,spectrum.data,lags)
                if shift.ndim == 1 :
                    pixshift=(fitpeak+lags[0])[0]
                    print('  Derived pixel shift from input wcal: ',fitpeak+lags[0])
                    if display is not None :
                        display.plotax1.cla()
                        display.plotax1.text(0.05,0.95,'spectrum and reference',transform=display.plotax1.transAxes)
                        for row in range(spectrum.data.shape[0]) :
                            display.plotax1.plot(spectrum.data[row,:],color='m')
                            display.plotax1.plot(self.spectrum.data[row,:],color='g')
                        display.plotax1.set_xlabel('Pixel')
                        display.plotax2.cla()
                        display.plotax2.text(0.05,0.95,'cross correlation: {:8.3f}'.format(pixshift),
                                             transform=display.plotax2.transAxes)
                        display.plotax2.plot(lags,shift)
                        display.plotax1.set_xlabel('Lag')
                        plt.draw()
                        input("  See spectrum and template spectrum (top), cross corrleation(bottom). hit any key to continue")
                    # single shift for all pixels
                    self.pix0 = self.pix0+fitpeak+lags[0]
                    wav=np.atleast_2d(self.wave(image=np.array(sz)))
                else :
                    # different shift for each row
                    wav=np.zeros(sz)
                    cols = np.arange(sz[-1])
                    orders=[]
                    for row in range(wav.shape[0]) :
                        print('  Derived pixel shift from input wcal for row: {:d} {:d}'.format
                               (row,shift[row,:].argmax()+lags[0]),end='\r')
                        rows=np.zeros(len(cols))+row
                        try : order = self.orders[row]
                        except : order=self.orders[0]
                        orders.append(order)
                        pix0 = self.pix0+fitpeak[row]+lags[0]
                        wav[row,:] = self.model(cols-pix0)/order
                    # ensure we have 2D fit
                    self.type = 'chebyshev2D'
                    self.orders = orders
                    print("")
            else :
                # get dispersion guess from header cards if not given in disp
                if disp is None: disp=hd.header['DISPDW']
                if wref is not None :
                    w0=wref[0]
                    pix0=wref[1]
                else:
                    w0=hd.header['DISPWC']
                    pix0=sz[1]/2
                wav=np.atleast_2d(w0+(pix-pix0)*disp)

        # open file with wavelengths and read
        if file is not None :
            f=open(ROOT+'/data/lamps/'+file,'r')
            lines=[]
            for line in f :
                if line[0] != '#' :
                    w=float(line.split()[0])
                    # if we have microns, convert to Angstroms
                    if w<10 : w*=10000
                    if w > wav.min() and w < wav.max() : lines.append(w)
            lines=np.array(lines)
            f.close()
        else :
            lines = self.waves
            weights = self.weights
            gd = np.where(weights >0)[0]
            lines = lines[gd]

        # get centroid around expected lines
        x=[]
        y=[]
        waves=[]
        waves_order=[]
        weight=[]
        diff=[]
        if display is not None and  isinstance(display,pyvista.tv.TV) :
            display.ax.cla()
            display.ax.axis('off')
            display.tv(spectrum.data)
        if plot is not None :
            if type(plot) is matplotlib.figure.Figure :
                plot.clf()
                plt.draw()
                ax1=plot.add_subplot(2,1,1)
                ax2=plot.add_subplot(2,1,2,sharex=ax1)
                plot.subplots_adjust(left=0.05,right=0.92, hspace=1.05)
                ax=[ax1,ax2]
                self.fig = plot
                self.ax = ax
            else :
                fig,ax = plt.subplots(2,1,sharex=True,figsize=(14,7))
                fig.subplots_adjust(hspace=1.05)
                self.fig = fig
                self.ax = ax

        if plot is not None : ax[0].cla()
        for row in range(0,nrow,nskip) :
            print('  identifying lines in row: ', row,end='\r')
            if plot is not None :
                ax[0].plot(wav[row,:],spectrum.data[row,:])
                #ax[0].set_yscale('log')
                ax[0].set_ylim(1.,ax[0].get_ylim()[1])
                ax[0].text(0.1,0.9,'row: {:d}'.format(row),transform=ax[0].transAxes)
                ax[0].set_xlabel('Rough wavelength')
                ax[0].set_ylabel('Intensity')
            for line in lines :
                peak=abs(line-wav[row,:]).argmin()
                if isinstance(display,pyvista.tv.TV) :
                    if (peak > xmin+rad) and (peak < xmax-rad) : display.ax.scatter(peak,row,marker='o',color='r',s=2)
                if ( (peak > xmin+rad) and (peak < xmax-rad) and
                     ((spectrum.data[row,peak-rad:peak+rad]/spectrum.uncertainty.array[row,peak-rad:peak+rad]).max() > thresh) ) :
                    cent = (spectrum.data[row,peak-rad:peak+rad]*np.arange(peak-rad,peak+rad)).sum()/spectrum.data[row,peak-rad:peak+rad].sum()
                    peak = int(cent)
                    cent = (spectrum.data[row,peak-rad:peak+rad]*np.arange(peak-rad,peak+rad)).sum()/spectrum.data[row,peak-rad:peak+rad].sum()
                    if display is not None and  isinstance(display,pyvista.tv.TV) :
                        display.ax.scatter(cent,row,marker='o',color='g',s=2)
                    if plot is not None :
                        ax[0].text(line,1.,'{:7.1f}'.format(line),rotation='vertical',va='top',ha='center')
                    x.append(cent)
                    y.append(row)
                    # we will fit for wavelength*order
                    waves.append(line)
                    try: order = self.orders[row]
                    except: order=self.orders[0]
                    waves_order.append(order)
                    weight.append(1.)
        if plot is not None :
            if self.model is not None :
                # if we have a solution already, see how good it is (after shift)
                diff=self.wave(pixels=[x,y])-np.array(waves)
                ax[1].cla()
                ax[1].scatter(np.array(waves),diff,s=2,c=y)
                ax[1].text(0.1,0.9,'from previous fit, rms: {:8.3f}'.format(diff.std()),transform=ax[1].transAxes)
                xlim=ax[1].get_xlim()
                ax[1].plot(xlim,[0,0],linestyle=':')
                ax[1].set_ylim(diff.min()-0.5,diff.max()+0.5)
                print("  rms from old fit (with shift): {:8.3f}".format(diff.std()))
            plt.figure(plot.number)
            plt.draw()
            input('  See identified lines. hit any key to continue....')
        self.pix=np.array(x)
        self.y=np.array(y)
        self.waves=np.array(waves)
        self.waves_order=np.array(waves_order)
        self.weights=np.array(weight)
        self.spectrum = spectrum
        print('')

    def scomb(self,hd,wav,average=True,usemask=True) :
        """ Resample onto input wavelength grid
        """
        #output grid
        out=np.zeros(len(wav))
        sig=np.zeros(len(wav))
        mask=np.zeros(len(wav),dtype=bool)
        # raw wavelengths
        w=self.wave(image=np.array(np.atleast_2d(hd.data).shape))
        for i in range(np.atleast_2d(hd).shape[0]) :
            sort=np.argsort(w[i,:])
            if usemask :
                gd = np.where(~hd.mask[i,sort])
                sort= sort[gd]
            wmin=w[i,sort].min()
            wmax=w[i,sort].max()
            w2=np.abs(wav-wmin).argmin()
            w1=np.abs(wav-wmax).argmin()
            if average :
                out[w2:w1] += ( np.interp(wav[w2:w1],w[i,sort],np.atleast_2d(hd.data)[i,sort]) /
                                np.interp(wav[w2:w1],w[i,sort],np.atleast_2d(hd.uncertainty.array)[i,sort])**2 )
                sig[w2:w1] += 1./np.interp(wav[w2:w1],w[i,sort],np.atleast_2d(hd.uncertainty.array)[i,sort])**2
            else :
                out[w2:w1] += np.interp(wav[w2:w1],w[i,sort],np.atleast_2d(hd.data)[i,sort])
                sig[w2:w1] += np.interp(wav[w2:w1],w[i,sort],np.atleast_2d(hd.uncertainty.array**2)[i,sort])
        if average :
            out = out / sig
        else :
            sig = np.sqrt(sig)
        return CCDData(out,uncertainty=StdDevUncertainty(sig),mask=mask,header=hd.header,unit='adu')

    def save(self,file) :
        """ Save object to file
        """
        try : delattr(self,'fig')
        except: pass
        try : delattr(self,'ax')
        except: pass
        f=open(file,'wb')
        pickle.dump(self,f)
        f.close()

class Trace() :
    """ Class for spectral traces
    """

    def __init__ (self,inst=None, type='poly',order=2,pix0=0,rad=5,spectrum=None,model=None,sc0=None,rows=None,lags=None,channel=None) :
        self.type = type
        self.order = order
        self.pix0 = pix0
        self.spectrum = spectrum
        self.rad = rad
        if inst == 'TSPEC' :
            self.order = 3
            self.rows = [[135,235],[295,395],[435,535],[560,660],[735,830]]
            self.lags = range(-75,75)
        elif inst == 'DIS' :
            if channel == 0 : self.rows=[[215,915]]
            elif channel == 1 : self.rows=[[100,800]]
            else : raise ValueError('need to specify channel')
            self.lags = range(-300,300)
        elif inst == 'ARCES' :
            self.lags = range(-10,10)
        if rows is not None : self.rows=rows
        if lags is not None : self.lags=lags
        if model is not None : self.model=model
        if sc0 is not None : self.sc0=sc0

    def trace(self,hd,srows,sc0=None,plot=None,thresh=20) :
        """ Trace a spectrum from starting position
        """

        fitter=fitting.LinearLSQFitter()
        if self.type == 'poly' :
            mod=models.Polynomial1D(degree=self.order)
        else :
            raise ValueError('unknown fitting type: '+self.type)
            return

        nrow = hd.data.shape[0]
        ncol = hd.data.shape[1]
        if sc0 is None : self.sc0 = int(ncol/2)
        else : self.sc0 = sc0
        self.spectrum = hd[:,self.sc0]
        self.spectrum.data[self.spectrum.data<0] = 0.
        rows = np.arange(nrow)
        ypos = np.zeros(ncol)
        ysum = np.zeros(ncol)
        yvar = np.zeros(ncol)
        ymask = np.zeros(ncol,dtype=bool)

        # we want to handle multiple traces, so make sure srows is iterable
        if type(srows ) is int or type(srows) is float : srows=[srows]
        oldmodel=copy.copy(self.model)
        self.model=[]
        if plot is not None :
            plot.clear()
            plot.tv(hd)

        rad = self.rad-1
        for irow,srow in enumerate(srows) :
            print('  Tracing row: {:d}'.format(int(srow)),end='\r')
            sr=copy.copy(srow)
            sr=int(round(sr))
            sr=hd.data[sr-rad:sr+rad+1,self.sc0].argmax()+sr-rad
            # march left from center
            for col in range(self.sc0,0,-1) :
                # centroid
                cr=sr-rad+hd.data[sr-rad:sr+rad+1,col].argmax()
                ysum[col] = np.sum(hd.data[cr-rad:cr+rad+1,col])
                ypos[col] = np.sum(rows[cr-rad:cr+rad+1]*hd.data[cr-rad:cr+rad+1,col]) / ysum[col]
                yvar[col] = np.sum(hd.uncertainty.array[cr-rad:cr+rad+1,col]**2)
                ymask[col] = np.any(hd.mask[cr-rad:cr+rad+1,col])
                # if centroid is too far from starting guess, mask as bad
                if np.abs(ypos[col]-sr) > rad/2. : ymask[col] = True
                # use this position as starting center for next if above threshold S/N
                if (not ymask[col]) & np.isfinite(ysum[col]) & (ysum[col]/np.sqrt(yvar[col]) > thresh)  : sr=int(round(ypos[col]))
            sr=copy.copy(srow)
            sr=int(round(sr))
            sr=hd.data[sr-rad:sr+rad+1,self.sc0].argmax()+sr-rad
            # march right from center
            for col in range(self.sc0+1,ncol,1) :
                # centroid
                cr=sr-rad+hd.data[sr-rad:sr+rad+1,col].argmax()
                ysum[col] = np.sum(hd.data[cr-rad:cr+rad+1,col])
                ypos[col] = np.sum(rows[cr-rad:cr+rad+1]*hd.data[cr-rad:cr+rad+1,col]) / ysum[col]
                yvar[col] = np.sum(hd.uncertainty.array[cr-rad:cr+rad+1,col]**2)
                ymask[col] = np.any(hd.mask[cr-rad:cr+rad+1,col])
                if np.abs(ypos[col]-sr) > rad/2. : ymask[col] = True
                # use this position as starting center for next if above threshold S/N
                if (not ymask[col]) & np.isfinite(ysum[col]) & (ysum[col]/np.sqrt(yvar[col]) > thresh)  : sr=int(round(ypos[col]))

            cols=np.arange(ncol)
            gd = np.where((~ymask) & (ysum/np.sqrt(yvar)>thresh) )[0]
            model=(fitter(mod,cols[gd],ypos[gd]))

            # reject outlier points (>1 pixel) and refit
            res = model(cols)-ypos
            gd = np.where((~ymask) & (ysum/np.sqrt(yvar)>thresh) & (np.abs(res)<1))[0]
            model=(fitter(mod,cols[gd],ypos[gd]))
            if len(gd) < 10 :
                print('  failed trace for row: {:d}, using old model'.format(irow))
                model=copy.copy(oldmodel[irow])
            self.model.append(model)

            if plot :
                plot.ax.scatter(cols,ypos,marker='o',color='r',s=4)
                plot.ax.scatter(cols[gd],ypos[gd],marker='o',color='g',s=4)
                plot.ax.plot(cols,model(cols),color='m')
                #plt.pause(0.05)

        self.pix0=0
        print("")
        if plot : input('  See trace. Hit any key to continue....')

    def retrace(self,hd,plot=None,thresh=20) :
        """ Retrace starting with existing model
        """
        self.find(hd)
        srows = []
        for row in range(len(self.model)) :
            srows.append(self.model[row](self.sc0))
        self.trace(hd,srows,plot=plot,thresh=thresh)

    def find(self,hd,lags=None,plot=None) :
        """ Determine shift from existing trace to input frame
        """
        if lags is None : lags = self.lags

        im=copy.deepcopy(hd.data)
        # if we have a window, zero array outside of window
        spec=im[:,self.sc0]
        try:
            spec[:self.rows[0]] = 0.
            spec[self.rows[1]:] = 0.
        except: pass
        fitpeak,shift = image.xcorr(self.spectrum,spec,lags)
        pixshift=(fitpeak+lags[0])[0]
        print('  traces shift: ', fitpeak+lags[0])
        if plot is not None :
            plot.clear()
            plot.tv(im)
            plot.plotax1.cla()
            plot.plotax1.text(0.05,0.95,'obj and ref cross-section',transform=plot.plotax1.transAxes)
            plot.plotax1.plot(self.spectrum.data/self.spectrum.data.max())
            plot.plotax1.plot(im[:,self.sc0]/im[:,self.sc0].max())
            plot.plotax1.set_xlabel('row')
            plot.plotax2.cla()
            plot.plotax2.text(0.05,0.95,'cross correlation {:8.3f}'.format(pixshift),
                              transform=plot.plotax2.transAxes)
            plot.plotax2.plot(lags,shift)
            plot.plotax2.set_xlabel('lag')
            plt.draw()
            input('  See spectra and cross-correlation. Hit any key to continue....')
        self.pix0=fitpeak+lags[0]
        return fitpeak+lags[0]

    def extract(self,hd,rad=None,scat=False,plot=None,medfilt=None) :
        """ Extract spectrum given trace(s)
        """
        if rad is None : rad=self.rad
        nrows=hd.data.shape[0]
        ncols=hd.data.shape[-1]
        spec = np.zeros([len(self.model),hd.data.shape[1]])
        sig = np.zeros([len(self.model),hd.data.shape[1]])
        mask = np.zeros([len(self.model),hd.data.shape[1]],dtype=bool)

        if plot is not None:
            plot.clear()
            plot.tv(hd)

        for i,model in enumerate(self.model) :
            print('  extracting aperture {:d}'.format(i),end='\r')
            cr=model(np.arange(ncols))+self.pix0
            icr=np.round(cr).astype(int)
            rfrac=cr-icr+0.5   # add 0.5 because we rounded
            rlo=[]
            rhi=[]
            for col in range(ncols) :
                r1=icr[col]-rad
                r2=icr[col]+rad
                # sum inner pixels directly, outer pixels depending on fractional pixel location of trace
                if r1>=0 and r2<nrows :
                    spec[i,col]=np.sum(hd.data[r1+1:r2,col])
                    sig[i,col]=np.sum(hd.uncertainty.array[r1+1:r2,col]**2)
                    spec[i,col]+=hd.data[r1,col]*(1-rfrac[col])
                    sig[i,col]+=hd.uncertainty.array[r1,col]**2*(1-rfrac[col])
                    spec[i,col]+=hd.data[r2,col]*rfrac[col]
                    sig[i,col]+=hd.uncertainty.array[r2,col]**2*rfrac[col]
                    sig[i,col]=np.sqrt(sig[i,col])
                    mask[i,col] = np.any(hd.mask[r1:r2+1,col])
                if plot is not None :
                    rlo.append(r1)
                    rhi.append(r2-1)
            if medfilt is not None :
                boxcar = Box1DKernel(medfilt)
                median = convolve(spec[i,:],boxcar,boundary='extend')
                spec[i,:]/=median
                sig[i,:]/=median

            if plot is not None :
                if i%2 == 0 : color='b'
                else : color='m'
                plot.ax.plot(range(ncols),cr,color='g',linewidth=3)
                plot.ax.plot(range(ncols),rlo,color=color,linewidth=1)
                plot.ax.plot(range(ncols),rhi,color=color,linewidth=1)
                plt.draw()
        if plot is not None : input('  See extraction window(s). Hit any key to continue....')
        print("")
        return CCDData(spec,uncertainty=StdDevUncertainty(sig),mask=mask,header=hd.header,unit='adu')

    def extract2d(self,hd,rows=None,plot=None) :
        """  Extract 2D spectrum given trace(s)
             Assumes all requests row uses same trace, just offset, not a 2D model for traces
        """
        nrows=hd.data.shape[0]
        ncols=hd.data.shape[-1]
        out=[]
        if plot is not None:
            plot.clear()
            plot.tv(hd)
        for model in self.model :
            if plot is not None :
                plot.ax.plot([0,ncols],[self.rows[0],self.rows[0]],color='g')
                plot.ax.plot([0,ncols],[self.rows[1],self.rows[1]],color='g')
                plt.draw()
            outrows=np.arange(self.rows[0],self.rows[1])
            noutrows=len(range(self.rows[0],self.rows[1]))
            spec=np.zeros([noutrows,ncols])
            sig=np.zeros([noutrows,ncols])
            cr=model(np.arange(ncols))
            cr-=cr[self.sc0]
            for col in range(ncols) :
                spec[:,col] = np.interp(outrows+cr[col],np.arange(nrows),hd.data[:,col])
                sig[:,col] = np.sqrt(np.interp(outrows+cr[col],np.arange(nrows),hd.uncertainty.array[:,col]**2))
            out.append(CCDData(spec,StdDevUncertainty(sig),unit='adu'))
        if plot is not None: input('  enter something to continue....')

        if len(out) == 1 : return out[0]
        else : return out

    def save(self,file) :
        """ Save object to file
        """
        try : delattr(self,'ax')
        except: pass
        f=open(file,'wb')
        pickle.dump(self,f)
        f.close()

def mash(hd,sp=None,bks=None) :
    """
    Mash image into spectra using requested window
    """
    if sp is None :
        sp=[0,hd.data.shape[0]]
    obj = hd.data[sp[0]:sp[1]].sum(axis=0)
    obj = hd.data[sp[0]:sp[1]].sum(axis=0)

    if bks is not None :
        back=[]
        for bk in bks :
           tmp=np.median(data[bk[0]:bk[1]],axis=0)
           back.append(tmp)
        obj-= np.mean(back,axis=0)

    return obj

def wavecal(hd,file=None,wref=None,disp=None,wid=[3],rad=5,snr=3,degree=2,wcal0=None,thresh=100,type='poly'):
    """
    Get wavelength solution for single 1D spectrum
    """

    # choose middle row +/ 5 rows
    sz=hd.data.shape
    spec=hd.data[int(sz[0]/2)-5:int(sz[0]/2)+5,:].sum(axis=0)
    spec=spec-scipy.signal.medfilt(spec,kernel_size=101)
    pix = np.arange(len(spec))

    fig,ax = plt.subplots(2,1,sharex=True,figsize=(14,6))
    ax[0].plot(spec)

    # get wavelength guess from input WaveCal if given, else use wref and dispersion, else header
    if wcal0 is not None :
        lags=range(-300,300)
        fitpeak,shift = image.xcorr(wcal0.spectrum,spec,lags)
        wnew=copy.deepcopy(wcal0)
        wnew.pix0 = wcal0.pix0+shift.argmax()+lags[0]
        print('  Derived pixel shift from input wcal0: ',shift.argmax()+lags[0])
        wav=wnew.wave(pix)
    else :
        # get dispersion guess from header cards if not given in disp
        if disp is None: disp=hd.header['DISPDW']
        if wref is not None :
            w0=wref[0]
            pix0=wref[1]
            wav=w0+(pix-pix0)*disp
        else:
            w0=hd.header['DISPWC']
            pix0=sz[1]/2
            wav=w0+(pix-pix0)*disp
    ax[1].plot(wav,spec)

    # open file with wavelengths and read
    f=open(file,'r')
    lines=[]
    for line in f :
        if line[0] != '#' :
            w=float(line.split()[0])
            name=line[10:].strip()
            lpix=abs(w-wav).argmin()
            if lpix > 1 and lpix < sz[1]-1 :
                ax[0].text(lpix,0.,'{:7.1f}'.format(w),rotation='vertical',va='top',ha='center')
                lines.append(w)
    lines=np.array(lines)
    f.close()

    # get centroid around expected lines
    cents=[]
    for line in lines :
        peak=abs(line-wav).argmin()
        if (peak > rad) and (peak < sz[1]-rad) and (spec[peak-rad:peak+rad].max() > thresh) :
            print(peak,spec[peak-rad:peak+rad].max())
            cents.append((spec[peak-rad:peak+rad]*np.arange(peak-rad,peak+rad)).sum()/spec[peak-rad:peak+rad].sum())
    cents=np.array(cents)
    print('  cents:', cents)

    waves=[]
    weight=[]
    print('  Centroid  W0  Wave')
    for cent in cents :
        w=wav[int(cent)]
        ax[0].plot([cent,cent],[0,10000],'k')
        print('  {:8.2f}{:8.2f}{:8.2f}'.format(cent, w, lines[np.abs(w-lines).argmin()]))
        waves.append(lines[np.abs(w-lines).argmin()])
        weight.append(1.)
    waves=np.array(waves)
    weight=np.array(weight)

    # set up new WaveCal object
    pix0 = int(sz[1]/2)
    wcal = WaveCal(order=degree,type=type,spectrum=spec,pix0=pix0)

    # iterate allowing for interactive removal of points
    done = False
    ymax = ax[0].get_ylim()[1]
    while not done :
        gd=np.where(weight>0.)[0]
        bd=np.where(weight<=0.)[0]
        wcal.fit(cents[gd],waves[gd],weights=weight[gd])

        # plot
        ax[1].cla()
        ax[1].plot(cents[gd],wcal.wave(cents[gd])-waves[gd],'go')
        if len(bd) > 0 : ax[1].plot(cents[bd],wcal.wave(cents[bd])-waves[bd],'ro')
        diff=wcal.wave(cents[gd])-waves[gd]
        ax[1].set_ylim(diff.min()-1,diff.max()+1)
        for i in range(len(cents)) :
            ax[1].text(cents[i],wcal.wave(cents[i])-waves[i],'{:2d}'.format(i),va='top',ha='center')
            if weight[i] > 0 :
              ax[0].plot([cents[i],cents[i]],[0,ymax],'g')
            else :
              ax[0].plot([cents[i],cents[i]],[0,ymax],'r')
        plt.draw()

        # get input from user on lines to remove
        for i in range(len(cents)) :
            print('  {:3d}{:8.2f}{:8.2f}{:8.2f}{:8.2f}{:8.2f}'.format(
                   i, cents[i], wcal.wave(cents[i]), waves[i], waves[i]-wcal.wave(cents[i]),weight[i]))
        print('  rms: {:8.2f} Anstroms'.format(diff.std()))
        i = input('enter ID of line to remove (-n for all lines<n, +n for all lines>n, return to continue): ')
        if i is '' :
            done = True
        elif '+' in i :
            weight[int(i)+1:] = 0.
        elif '-' in i :
            weight[0:abs(int(i))] = 0.
        elif int(i) >= 0 :
            weight[int(i)] = 0.
        else :
            print('invalid input')

    plt.close()

    return wcal.wave(pix),wcal

def fluxcal(obs,wobs,file=None) :
    """
    flux calibration
    """

    fluxdata=ascii.read(file)
    stan=np.interp(wobs,fluxdata['col1'],fluxdata['col2'])
    return stan/obs

def trace(hd,apertures=None,pix0=1024) :
    """ Get all traces
        apertures is a list of row numbers at pixel 1024
    """
    alltr=[]
    for i in range(len(apertures)) :
        tr=Trace()
        print('tracing aperture {:d}'.format(i),end='\r')
        sr=apertures[i]
        tr.trace(hd,pix0,sr)
        alltr.append(tr)

    return alltr

def extract(hd,apertures) :
    """ Do all extractions
    """
    spec = np.zeros([len(apertures),hd.data.shape[1]])
    for i,order in enumerate(apertures) :
        print('extracting aperture {:d}'.format(i),end='\r')
        spec[i] = order.extract(hd)

    return spec


"""
The purpose of this stage is to simulate event classification sorting the
reconstructed particles into PID-signature channels.

For each PID signature, the input map is transformed by the probability for
events in each of its bins to be ID'd as that signature. Therefore the ouptut
binning is similar to the input binning, but with the added 'pid' dimension,
which has as many bins as PID signatures.
"""

from __future__ import division

from collections import Mapping, OrderedDict

import numpy as np
import scipy as sp

# NOTE: need both versions of the imported names, as eval strings can name
# numpy and scipy either ways
import numpy
import scipy

from pisa.core.stage import Stage
from pisa.core.transform import BinnedTensorTransform, TransformSet
from pisa.utils.fileio import from_file
from pisa.utils.flavInt import flavintGroupsFromString, NuFlavIntGroup
from pisa.utils.hash import hash_obj
from pisa.utils.log import logging
from pisa.utils.profiler import profile

__all__ = ['param']

__author__ = 'L. Schulte, J.L. Lanfranchi, S. Mandalia'

__license__ = '''Copyright (c) 2014-2017, The IceCube Collaboration

 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.'''

def load_pid_energy_param(source):
    """Load pid energy-dependent parameterisation from file or dictionary.

    Parameters
    ----------
    source : string or mapping
        If string, interprete as resource location of the file; if mapping, use
        directly.

    Returns
    -------
    pid_energy_param_dict : OrderedDict
        Keys are `NuFlavIntGroup`s and values are callables of one arg.

    """
    # Get the original dict
    if isinstance(source, basestring):
        orig_dict = from_file(source)
    elif isinstance(source, Mapping):
        orig_dict = source
    else:
        raise TypeError('`source` must either be string or mapping; got %s'
                        ' instead.' % type(source))

    # Build dict with flavintgroups as keys; subdict with signatures as keys
    # and callables as values
    pid_energy_param_dict = OrderedDict()

    for flavintgroup_str, subdict in orig_dict.iteritems():
        flavintgroup = NuFlavIntGroup(flavintgroup_str)

        pid_energy_param_dict[flavintgroup] = OrderedDict()

        for signature, sig_param_spec in subdict.iteritems():
            if isinstance(sig_param_spec, basestring):
                sig_param_func = eval(sig_param_spec)
                if not callable(sig_param_func):
                    raise ValueError(
                        'Group %s PID signature %s param spec "%s" does'
                        ' not evaluate to a callable.'
                        % (xform_flavints, signature, sig_param_spec)
                    )
            elif callable(sig_param_spec):
                sig_param_func = sig_param_spec
            else:
                raise TypeError(
                    'Group %s PID signature %s parameterization is a "%s"'
                    ' but must be a string or callable.'
                    % (xform_flavints, signature, type(sig_param_spec))
                )

            pid_energy_param_dict[flavintgroup][signature] = sig_param_func

    return pid_energy_param_dict

class param(Stage):
    """Parameterised MC PID based on an input json file containing functions
    describing the PID as a function of energy.

    Transforms an input map of the specified particle "signature" (aka ID) into
    a map of the track-like events ('track') and a map of the shower-like events
    ('cascade').

    Parameters
    ----------
    params : ParamSet or sequence with which to instantiate a ParamSet

        Parameters which set everything besides the binning.

        If str, interpret as resource location and load params from resource.
        If dict, set contained params. Format expected is
            {'<param_name>': <Param object or passable to Param()>}

        Parameters required by this service are
            * pid_energy_paramfile : dict or filepath
                json file or equivalent dict containing the PID functions for
                each flavor. The structure should be:
                  {
                    "numu_cc + numubar_cc": {
                      "track" : "lambda E: some function",
                      "cascade" : "lambda E: 1 - some function"
                    },
                    "nue_cc + nuebar_cc": {
                      "track" : "lambda E: some function",
                      "cascade" : "lambda E: 1 - some function"
                    },
                    "nutau_cc + nutaubar_cc": {
                      "track" : "lambda E: some function",
                      "cascade" : "lambda E: 1 - some function"
                    },
                    "nuall_nc + nuallbar_nc": {
                      "track" : "lambda E: some function",
                      "cascade" : "lambda E: 1 - some function"
                    }
                  }

    particles : string

    input_names : sequence of strings

    transform_groups

    sum_grouped_flavints : bool

    input_binning : MultiDimBinning
        Arbitrary number of dimensions accepted. Contents of the input
        `pid_events` parameter defines the possible binning dimensions. Name(s)
        of given binning(s) must match to a reco variable in `pid_events`.

    output_binning : MultiDimBinning

    error_method : None, bool, or string

    transforms_cache_depth : int >= 0

    outputs_cache_depth : int >= 0

    memcache_deepcopy : bool

    debug_mode : None, bool, or string
        Whether to store extra debug info for this service.

    Input Names
    ----------
    The `inputs` container must include objects with `name` attributes:
        * 'nue_cc'
        * 'nuebar_cc'
        * 'numu_cc'
        * 'numubar_cc'
        * 'nutau_cc'
        * 'nutaubar_cc'
        * 'nuall_nc'
        * 'nuallbar_nc'

    Output Names
    ----------
    The `outputs` container generated by this service will be objects with the
    following `name` attribute; pid is added as a binning dimension:
        * 'nue_cc'
        * 'nuebar_cc'
        * 'numu_cc'
        * 'numubar_cc'
        * 'nutau_cc'
        * 'nutaubar_cc'
        * 'nuall_nc'
        * 'nuallbar_nc'

    """
    def __init__(self, params, particles, input_names, transform_groups,
                 sum_grouped_flavints, input_binning, output_binning,
                 memcache_deepcopy, error_method, transforms_cache_depth,
                 outputs_cache_depth, debug_mode=None):
        assert particles in ['muons', 'neutrinos']
        self.particles = particles.strip().lower()
        """Whether stage is instantiated to process neutrinos or muons"""

        self.transform_groups = flavintGroupsFromString(transform_groups)
        """Particle/interaction types to group for computing transforms"""

        self.sum_grouped_flavints = sum_grouped_flavints

        # All of the following params (and no more) must be passed via
        # the `params` argument.
        expected_params = (
            'pid_energy_paramfile'
        )

        if isinstance(input_names, basestring):
            input_names = input_names.replace(' ', '').split(',')

        if self.particles == 'neutrinos':
            if self.sum_grouped_flavints:
                output_names = [str(g) for g in self.transform_groups]
            else:
                output_names = input_names
        elif self.particles == 'muons':
            raise NotImplementedError('%s not implemented.' % self.particles)

        super(self.__class__, self).__init__(
            use_transforms=True,
            params=params,
            expected_params=expected_params,
            input_names=input_names,
            output_names=output_names,
            error_method=error_method,
            outputs_cache_depth=outputs_cache_depth,
            transforms_cache_depth=transforms_cache_depth,
            memcache_deepcopy=memcache_deepcopy,
            input_binning=input_binning,
            output_binning=output_binning,
            debug_mode=debug_mode
        )

        self.include_attrs_for_hashes('particles')
        self.include_attrs_for_hashes('sum_grouped_flavints')
        self.include_attrs_for_hashes('transform_groups')

        self.signatures = output_binning.pid.bin_names
        """PID signatures that this stage generates"""

        # If no bin names are present, use the integer bin indices instead
        if self.signatures is None:
            self.signatures = range(len(output_binning.pid))

        # Define the transform binnning...

        # Note that Numpy broadcasting rules start with last axis and work
        # inwards. We want the input map (say MxN) to automatically be
        # broadcast to multiply into each of L PID bins. Therefore, if
        # we _prepend_ the PID dimension to the transform, we have an MxN input
        # multiplying an LxMxN transform, and Numpy treats this as L separate
        # MxN by MxN multiplies which populate an output array of dimension
        # LxMxN... exactly what we want, and with maximal computational
        # efficiency (for Numpy to handle, at least). If the user's output
        # binning does not follow the same ordering, this is okay, as the
        # output is passed through the `rebin` function each time it is
        # computed, and this takes care of any axis swapping necessary.

        self.transform_output_binning = (
            self.output_binning.pid * self.input_binning
        )

        self.ebin_centers = (
            self.input_binning.reco_energy.weighted_centers.m_as('GeV')
        )

        self.pid_energy_param_dict = None
        self._pid_energy_param_hash = None

    def validate_binning(self):
        """Validate input and output binning"""
        required_input_binning_dims = 'reco_energy', 'reco_coszen'
        required_output_binning_dims = 'reco_energy', 'reco_coszen', 'pid'

        msg = ('%s binning must contain dimensions %s, but has dimensions %s'
               ' instead.')

        if set(self.input_binning.names) != set(required_input_binning_dims):
            raise ValueError(msg % ('Input', required_input_binning_dims,
                                    self.input_binning.names))

        if set(self.output_binning.names) != set(required_output_binning_dims):
            raise ValueError(msg % ('Output', required_output_binning_dims,
                                    self.input_binning.names))

        # While output binning will have a 'pid' dimension, the remaining
        # dimensions must be the same in both input and output binnings
        for dim in self.input_binning.dims:
            if dim != self.output_binning[dim.name]:
                raise NotImplementedError(
                    'Input and output dimensions %s are not equal, but stage'
                    ' %s / service %s does not implement binning up- or'
                    ' downsampling.'
                    % (dim.name, self.stage_name, self.service_name)
                )

    def load_pid_energy_param(self, source):
        """Load pid energy-dependent parameterisation from file or dictionary.

        Parameters
        ----------
        source : string
            Resource location of the file

        """
        this_hash = hash_obj(source)
        if (self._pid_energy_param_hash is not None
                and this_hash == self._pid_energy_param_hash):
            return

        # Invalidate the hash and clear the entry, so we aren't left in an
        # inconsistent state if any of the below fails
        self._pid_energy_param_hash = None
        self.pid_energy_param_dict = None

        # Call external function for basic loading and conversion
        pid_energy_param_dict = load_pid_energy_param(source)

        # Perform validation
        for flavintgroup, subdict in pid_energy_param_dict.iteritems():
            if set(subdict.keys()) != set(self.signatures):
                raise ValueError(
                    'Expected PID specs for %s, but the energy PID'
                    ' parameterization for %s specifies %s instead.'
                    % (self.signatures, flavintgroup, subdict.keys())
                )

        # Transform groups are implicitly defined by keys
        implicit_transform_groups = pid_energy_param_dict.keys()

        # Make sure these match the transform groups specified for the stage
        if set(implicit_transform_groups) != set(self.transform_groups):
            raise ValueError(
                'Transform groups (%s) defined implicitly by `source` "%s" do'
                ' not match those defined as the stage\'s configured'
                ' `transform_groups` (%s).'
                % (implicit_transform_groups, source, self.transform_groups)
            )

        # Verify that each input name--which specifies a flavint or
        # flavintgroup--is wholly encapsulated by one of the transform
        # flavintgroups
        for name in self.input_names:
            if not any(name in group for group in implicit_transform_groups):
                raise ValueError(
                    'Input "%s" either not present in or spans multiple'
                    ' transform groups (transform_groups = %s)'
                    % (name, implicit_transform_groups)
                )

        self.pid_energy_param_dict = pid_energy_param_dict
        self._pid_energy_param_hash = this_hash

    @profile
    def _compute_nominal_transforms(self):
        """Compute new PID transforms."""
        logging.debug('Updating pid.param PID histograms...')

        self.load_pid_energy_param(self.params.pid_energy_paramfile.value)

        nominal_transforms = []
        for xform_flavints in self.transform_groups:
            logging.debug('Working on %s PID', xform_flavints)

            xform_array = np.empty(self.transform_output_binning.shape)

            subdict = self.pid_energy_param_dict[xform_flavints]
            for signature, sig_param_func in subdict.iteritems():
                # Get the PID probabilities vs. energy at the energy bins'
                # (weighted) centers
                pid1d = sig_param_func(self.ebin_centers)

                # Broadcast this 1d array across the reco_coszen dimension
                # since it's independent of reco_coszen
                broadcasted_pid = self.transform_output_binning.broadcast(
                    pid1d, from_dim='reco_energy', to_dims='reco_coszen'
                )

                pid_indexer = (
                    self.transform_output_binning.indexer(pid=signature)
                )

                # Assign the broadcasted array to the correct PID bin
                xform_array[pid_indexer] = broadcasted_pid

            if self.sum_grouped_flavints:
                xform_input_names = []
                for input_name in self.input_names:
                    input_flavs = NuFlavIntGroup(input_name)
                    if set(xform_flavints).intersection(input_flavs):
                        xform_input_names.append(input_name)

                for output_name in self.output_names:
                    if output_name not in xform_flavints:
                        continue
                    xform = BinnedTensorTransform(
                        input_names=xform_input_names,
                        output_name=str(xform_flavints),
                        input_binning=self.input_binning,
                        output_binning=self.transform_output_binning,
                        xform_array=xform_array,
                        sum_inputs=self.sum_grouped_flavints
                    )
                    nominal_transforms.append(xform)

            else:
                for input_name in self.input_names:
                    if input_name not in xform_flavints:
                        continue
                    xform = BinnedTensorTransform(
                        input_names=input_name,
                        output_name=input_name,
                        input_binning=self.input_binning,
                        output_binning=self.transform_output_binning,
                        xform_array=xform_array,
                    )
                    nominal_transforms.append(xform)

        return TransformSet(transforms=nominal_transforms)

    def _compute_transforms(self):
        """There are no systematics in this stage, so the transforms are just
        the nominal transforms. Thus, this function just returns the nominal
        transforms, computed by `_compute_nominal_transforms`..
        """
        return self.nominal_transforms

    def validate_params(self, params):
        """Do checks on the parameters"""
        val = params.pid_energy_paramfile.value
        if not isinstance(val, (basestring, Mapping)):
            raise TypeError(
                'Expecting either a path to a file or a dictionary provided'
                ' as the store of the parameterisations. Got "%s".' % type(val)
            )


#!/usr/bin/env python
""" Where you at? """
import sys,os
import logging
from collections import OrderedDict as odict
from datetime import datetime,timedelta,tzinfo
import dateutil.parser

import mpl_toolkits.basemap as basemap
from matplotlib.patches import Ellipse, Circle
import matplotlib.patheffects as patheffects
from _tkinter import TclError

import numpy as np
import pylab as plt
import ephem

__author__  = "Alex Drlica-Wagner"
__email__   = "kadrlica@fnal.gov"
__version__ = "2.1.3"

MAXREF=5000 # Maximum number of refreshes
DECAM=1.1   # DECam radius (deg)

# Accurate DECam marker size depends on figsize and DPI
# This is a mess...
FIGSIZE=(10.5,8.5)
SCALE=np.sqrt((8.0*6.0)/(FIGSIZE[0]*FIGSIZE[1]))
DPI=80;

FILTERS = ['u','g','r','i','z','Y','VR']
BANDS = FILTERS + ['all']
COLORS = odict([
    ('none','black'),
    ('u','blue'),
    ('g','green'),
    ('r','red'),
    ('i','gold'),
    ('z','magenta'),
    ('Y','black'),
    ('VR','gray'),
])

# Allowed map projections
PROJ = odict([
    ('ortho'  , dict(projection='ortho',celestial=True)),
    ('moll'   , dict(projection='moll',celestial=True)),
    ('mol'    , dict(projection='moll',celestial=True)),
    ('ait'    , dict(projection='hammer',celestial=True)),
    ('mbt'    , dict(projection='mbtfpq',celestial=True)),
    ('mbtfpq' , dict(projection='mbtfpq',celestial=True)),
    ('mcbryde', dict(projection='mbtfpq',celestial=True)),
])

# Derived from telra,teldec of 10000 exposures
SN = odict([
    ('E1',(7.874, -43.010)),
    ('E2',(9.500, -43.999)),
    ('X1',(34.476, -4.931)),
    ('X2',(35.664,-6.413)),
    ('X3',(36.449, -4.601)),
    ('S1',(42.818, 0.000)),
    ('S2',(41.193, -0.991)),
    ('C1',(54.274, -27.113)),
    ('C2',(54.274, -29.090)),
    ('C3',(52.647, -28.101)),
])

SN_LABELS = odict([
    ('SN-E',(8,-41)),
    ('SN-X',(35,-12)),
    ('SN-S',(45,1)),
    ('SN-C',(55,-35)),
])

# The allowed footprint outlines
FOOTPRINTS = ['none','des','des-sn','smash','maglites','bliss','decals','delve']

# CTIO location taken from:
#http://www.ctio.noao.edu/noao/content/Coordinates-Observatories-Cerro-Tololo-and-Cerro-Pachon
#http://arxiv.org/pdf/1210.1616v3.pdf
#(-30h 10m 10.73s, -70h 48m 23.52s, 2213m)

TEL_LON = -70.80653
TEL_LAT = -30.169647
TEL_HEIGHT = 2213

# Create the observatory object
CTIO = ephem.Observer()
CTIO.lon,CTIO.lat = str(TEL_LON),str(TEL_LAT)
CTIO.elevation = TEL_HEIGHT

def get_datadir():
    """ Path to data directory. """
    return os.path.join(os.path.dirname(os.path.realpath(__file__)),'data')

def setdefaults(kwargs,defaults):
    """ set dictionary with defaults. """
    for k,v in defaults.items():
        kwargs.setdefault(k,v)
    return kwargs

def gal2cel(glon, glat):
    """
    Converts Galactic (deg) to Celestial J2000 (deg) coordinates
    """
    glat = np.radians(glat)
    sin_glat = np.sin(glat)
    cos_glat = np.cos(glat)

    glon = np.radians(glon)
    ra_gp = np.radians(192.85948)
    de_gp = np.radians(27.12825)
    lcp = np.radians(122.932)

    sin_lcp_glon = np.sin(lcp - glon)
    cos_lcp_glon = np.cos(lcp - glon)

    sin_d = (np.sin(de_gp) * sin_glat) \
            + (np.cos(de_gp) * cos_glat * cos_lcp_glon)
    ramragp = np.arctan2(cos_glat * sin_lcp_glon,
                         (np.cos(de_gp) * sin_glat) \
                         - (np.sin(de_gp) * cos_glat * cos_lcp_glon))
    dec = np.arcsin(sin_d)
    ra = (ramragp + ra_gp + (2. * np.pi)) % (2. * np.pi)
    return np.degrees(ra), np.degrees(dec)

def cel2gal(ra, dec):
    """
    Converts Celestial J2000 (deg) to Calactic (deg) coordinates
    """
    dec = np.radians(dec)
    sin_dec = np.sin(dec)
    cos_dec = np.cos(dec)

    ra = np.radians(ra)
    ra_gp = np.radians(192.85948)
    de_gp = np.radians(27.12825)

    sin_ra_gp = np.sin(ra - ra_gp)
    cos_ra_gp = np.cos(ra - ra_gp)

    lcp = np.radians(122.932)
    sin_b = (np.sin(de_gp) * sin_dec) \
            + (np.cos(de_gp) * cos_dec * cos_ra_gp)
    lcpml = np.arctan2(cos_dec * sin_ra_gp,
                       (np.cos(de_gp) * sin_dec) \
                       - (np.sin(de_gp) * cos_dec * cos_ra_gp))
    glat = np.arcsin(sin_b)
    glon = (lcp - lcpml + (2. * np.pi)) % (2. * np.pi)
    return np.degrees(glon), np.degrees(glat)

# Stupid timezone definition
ZERO = timedelta(0)
HOUR = timedelta(hours=1)
class UTC(tzinfo):
    """UTC"""
    def utcoffset(self, dt):
        return ZERO

    def tzname(self, dt):
        return "UTC"

    def dst(self, dt):
        return ZERO

def safe_proj(bmap,lon,lat,inverse=False):
    """ Remove points outside of projection

    Parameters:
    -----------
    bmap : basemap
    lon  : longitude
    lat  : latitude
    inverse : inverse projection

    Returns:
    --------
    x,y : projected coordinates
    """
    x,y = bmap(np.atleast_1d(lon),np.atleast_1d(lat),inverse=inverse)
    x[np.abs(x) > 1e29] = None
    y[np.abs(y) > 1e29] = None
    return x,y

def get_boundary(bmap,projection,fact=0.99):
    # Check that point inside boundary
    # Doesn't work for 'ait' and 'moll' projections
    if projection in basemap._pseudocyl:
        # This was estimated by eye...
        rminor=9.00995e6; rmajor = 2*rminor
        boundary = Ellipse((rmajor,rminor),
                           2*(fact*rmajor),2*(fact*rminor))
    else:
        boundary = Ellipse((bmap.rmajor,bmap.rminor),
                           2*(fact*bmap.rmajor),2*(fact*bmap.rminor))

    return boundary

def airmass_angle(x=1.4):
    """ Zenith angle for a given airmass limit """
    return 90.-np.degrees(np.arcsin(1./x))

def load_data(opts):
    """ Load the data (either from DB of file).

    Parameters:
    -----------
    opts : command line options

    Returns:
    --------
    data : numpy recarray
    """
    since = parse_since(opts.since)
    propid = '%' if opts.propid is None else opts.propid
    dtype=[('expnum',int),('telra',float),('teldec',float),('filter',object)]

    if opts.infile is None:
        selection = ['id','telra','teldec','filter']
        #filter = "exposed = TRUE AND flavor LIKE '%s' AND date > '%s' AND propid LIKE '%s' ORDER BY id DESC"%(opts.flavor,since.isoformat(),propid)
        filter = "exposed = TRUE AND flavor SIMILAR TO '%s' AND date > '%s' AND propid LIKE '%s' ORDER BY id DESC"%(opts.flavor,since.isoformat(),propid)
        # Use the FNAL mirror to avoid overloading CTIO
        try: from database import Database
        except ImportError: from pointing.database import Database
        db = Database(dbname='db-'+opts.db)
        db.connect()
        query = "SELECT %s FROM exposure WHERE %s"%(','.join(selection),filter)
        #query = "SELECT id as expnum,telra as ra,teldec as dec,filter as band FROM exposure WHERE exposed = TRUE AND flavor LIKE 'object' and telra between 80 and 82 AND teldec between -71 and -69"
        data = db.execute(query)

        if len(data): ret = np.rec.array(data,dtype=dtype)
        else:         ret = np.rec.recarray(0,dtype=dtype)

        return ret
    else:
        return np.loadtxt(opts.infile,dtype=dtype)

def mjd(datetime):
    """ Modified Julian Date (MJD) """
    mjd_epoch = dateutil.parser.parse('1858-11-17T00:00:00Z')
    mjd_date = (datetime-mjd_epoch).total_seconds()/float(24*60*60)
    return mjd_date

def lmst(observatory):
    """ Calculate Local Mean Sidereal Time (LMST) """
    lmst = np.degrees(observatory.sidereal_time())
    logging.debug('Using pyephem for LMST: %.3f'%lmst)
    return lmst

def moon(datetime):
    """ Moon location

    Parameters:
    -----------
    datetime : the datetime of moon location request

    Returns:
    --------
    (ra, dec), phase : moon parameters [(deg, deg), %]
    """
    moon = ephem.Moon()
    moon.compute(CTIO)
    moon_phase = moon.moon_phase * 100
    moon_ra,moon_dec = np.degrees([moon.ra,moon.dec])
    return (moon_ra, moon_dec),moon_phase

def boolean(string):
    """ Convert strings to booleans for argparse """
    string = string.lower()
    if string in ['0', 'f', 'false', 'no', 'off']:
        return False
    elif string in ['1', 't', 'true', 'yes', 'on']:
        return True
    else:
        raise ValueError()

def splash_screen():
    """ Splash text to print """
    splash = """Running Alex Drlica-Wagner's DECam pointing script..."""
    logging.info(splash)

def parse_utc(value):
    """ Parse isoformat 'utc' option string. """
    if value is None:
        utc = datetime.now(tz=UTC())
    elif isinstance(value,datetime):
        utc = value
    else:
        utc = dateutil.parser.parse(value,tzinfos={'UTC':UTC})
    logging.debug("UTC: %s"%utc.strftime('%Y-%m-%d %H:%M:%S'))
    return utc

def parse_since(value):
    """ Parse isoformat 'since' option string. """
    if value is None:
        since = datetime.now(tz=UTC()) - timedelta(hours=12)
    elif isinstance(value,datetime):
        since = value
    elif value.lower() in ['all','none','forever']:
        since = dateutil.parser.parse('2012-01-01 12:00',tzinfos={'UTC':UTC})
    else:
        since = dateutil.parser.parse(value,tzinfos={'UTC':UTC})
    logging.debug("Since: %s"%since.strftime('%Y-%m-%d %H:%M:%S'))
    return since

def draw_constellation(bmap,name):
    """ Draw a map of the constellations (work in progress). """
    from constellations import CONSTELLATIONS
    points = np.array(CONSTELLATIONS[name])

    drawtype = points[:,0]
    radeg = points[:,1] * 1.0 / 1800 * 15
    decdeg = points[:,2] * 1.0 / 60
    print(radeg,decdeg)
    verts = zip(safe_proj(bmap,radeg,decdeg))
    codes = [XEPHEM2PATH[c] for c in points[:,0]]
    print(x,y)

def draw_milky_way(bmap,width=10,**kwargs):
    """ Draw the Milky Way galaxy. """
    defaults = dict(color='k',lw=1.5,ls='-')
    setdefaults(kwargs,defaults)

    logging.debug("Plotting the Milky Way")
    glon = np.linspace(0,360,500)
    glat = np.zeros_like(glon)
    ra,dec = gal2cel(glon,glat)
    ra -= 360*(ra > 180)

    proj = safe_proj(bmap,ra,dec)
    bmap.plot(*proj,**kwargs)

    if width:
        kwargs.update(dict(ls='--',lw=1))
        for delta in [+width,-width]:
            ra,dec = gal2cel(glon,glat+delta)
            proj = safe_proj(bmap,ra,dec)
            bmap.plot(*proj,**kwargs)

def draw_des(bmap,**kwargs):
    """
    Plot the DES wide-field footprint.

    Parameters:
    -----------
    bmap   : The basemap object
    kwargs : Various plotting arguments

    Returns:
    --------
    None
    """
    # Plot the wide-field survey footprint
    logging.debug("Plotting footprint: %s"%opts.footprint)
    #basedir = os.path.dirname(os.path.abspath(__file__))
    infile = os.path.join(get_datadir(),'des-round19-poly.txt')
    perim = np.loadtxt(infile,dtype=[('ra',float),('dec',float)])
    proj = safe_proj(bmap,perim['ra'],perim['dec'])
    bmap.plot(*proj,**kwargs)

def draw_des_sn(bmap,**kwargs):
    """
    Plot the DES supernova fields.

    Parameters:
    -----------
    bmap   : The basemap object
    kwargs : Various plotting arguments

    Returns:
    --------
    None
    """
    # Plot the SN fields
    logging.debug("Plotting DES supernova fields.")

    boundary = get_boundary(bmap,kwargs.pop('projection',None),fact=0.99)
    for v in SN.values():
        if not boundary.contains_point(bmap(*v)):
            continue
        # This does the projection correctly, but fails at boundary
        bmap.tissot(v[0],v[1],DECAM,100,**kwargs)

    # The SN labels
    sntxt_kwargs = dict(zorder=kwargs['zorder'],fontsize=12,
                        bbox=dict(boxstyle='round,pad=0',fc='w',ec='none',
                                  alpha=0.25))
    for k,v in SN_LABELS.items():
        plt.gca().annotate(k,bmap(*v),**sntxt_kwargs)

def draw_smash(bmap,**kwargs):
    """ Draw the SMASH fields

    Parameters:
    -----------
    bmap   : The basemap object
    kwargs : Various plotting arguments

    Returns:
    --------
    None
    """
    filename = os.path.join(get_datadir(),'smash_fields_final.txt')

    smash=np.genfromtxt(filename,dtype=[('ra',float),('dec',float)],usecols=[4,5])
    smash_x,smash_y = safe_proj(bmap,smash['ra'],smash['dec'])
    kwargs.update(dict(facecolor='none'))
    bmap.scatter(smash_x,smash_y,color='k',**kwargs)

def draw_maglites(bmap,**kwargs):
    """
    Plot the MagLiteS Phase-I footprint.

    Parameters:
    -----------
    bmap   : The basemap object
    kwargs : Various plotting arguments

    Returns:
    --------
    None
    """

    # Plot the wide-field survey footprint
    logging.debug("Plotting MagLiteS footprint")
    infile = os.path.join(get_datadir(),'maglites-poly.txt')
    perim = np.loadtxt(infile,dtype=[('ra',float),('dec',float)])
    proj = safe_proj(bmap,perim['ra'],perim['dec'])
    bmap.plot(*proj,**kwargs)

def draw_maglites2(bmap,**kwargs):
    """
    Plot the MagLiteS Phase-II footprint.

    Parameters:
    -----------
    bmap   : The basemap object
    kwargs : Various plotting arguments

    Returns:
    --------
    None
    """
    # Plot the wide-field survey footprint
    logging.debug("Plotting footprint: %s"%opts.footprint)
    infile = os.path.join(get_datadir(),'maglitesII-poly.txt')
    perim = np.loadtxt(infile,dtype=[('ra',float),('dec',float),('poly',int)])
    for p in np.unique(perim['poly']):
        sel = (perim['poly'] == p)
        proj = safe_proj(bmap,perim[sel]['ra'],perim[sel]['dec'])
        bmap.plot(*proj,**kwargs)

def draw_bliss(bmap,**kwargs):
    """
    Plot the BLISS wide-field footprint.

    Parameters:
    -----------
    bmap   : The basemap object
    kwargs : Various plotting arguments

    Returns:
    --------
    None
    """
    # Plot the wide-field survey footprint
    logging.debug("Plotting footprint: %s"%opts.footprint)
    infile = os.path.join(get_datadir(),'bliss-poly.txt')
    perim = np.loadtxt(infile,dtype=[('ra',float),('dec',float),('poly',int)])
    for p in np.unique(perim['poly']):
        sel = (perim['poly'] == p)
        proj = safe_proj(bmap,perim[sel]['ra'],perim[sel]['dec'])
        bmap.plot(*proj,**kwargs)

def draw_decals(bmap,**kwargs):
    """
    Plot the DECaLS wide-field footprint.

    Parameters:
    -----------
    bmap   : The basemap object
    kwargs : Various plotting arguments

    Returns:
    --------
    None
    """
    # Plot the wide-field survey footprint
    logging.debug("Plotting footprint: %s"%opts.footprint)
    infile = os.path.join(get_datadir(),'decals-poly.txt')
    perim = np.loadtxt(infile,dtype=[('ra',float),('dec',float),('poly',int)])
    for p in np.unique(perim['poly']):
        sel = (perim['poly'] == p)
        proj = safe_proj(bmap,perim[sel]['ra'],perim[sel]['dec'])
        bmap.plot(*proj,**kwargs)

def draw_delve(bmap,**kwargs):
    """ Draw DELVE footprint """
    defaults=dict(color='red', lw=2)
    setdefaults(kwargs,defaults)

    logging.debug("Plotting footprint: %s"%opts.footprint)
    deep = odict([
        ('SextansB', (150.00,   5.33, 3.0)),
        ('IC5152',   (330.67, -51.30, 3.0)),
        ('NGC300',   ( 13.72, -37.68, 3.0)),
        ('NGC55',    (  3.79, -39.22, 3.0)),
    ])
    boundary = get_boundary(bmap,kwargs.pop('projection',None),fact=0.98)

    for ra,dec,radius in deep.values():
        if not boundary.contains_point(bmap(ra,dec)): continue
        # This does the projection correctly, but fails at boundary
        bmap.tissot(ra,dec,radius,100,fc='none',edgecolor=kwargs['color'],lw=kwargs['lw'])

    #for ra,dec,radius in deep.values():
    #    # This doesn't deal with boundaries well
    #    #self.tissot(ra, dec, radius, 100, fc='none',**kwargs)
    #    x,y = safe_proj(bmap,np.array([ra]), np.array([dec]))
    #    bmap.scatter(x,y,facecolor='none',edgecolor=kwargs['color'],lw=2,s=400)

    filename = os.path.join(get_datadir(),'delve-poly.txt')
    perim = np.loadtxt(filename,dtype=[('ra',float),('dec',float),('poly',int)])
    for p in np.unique(perim['poly']):
        sel = (perim['poly'] == p)
        proj = safe_proj(bmap,perim[sel]['ra'],perim[sel]['dec'])
        bmap.plot(*proj,**kwargs)

def plot(opts):
    """
    Core plotting function. Creates the basemap, overplots all of the
    requested features, and returns the map object.

    Parameters:
    -----------
    opts : command line options

    Returns:
    --------
    m : the basemap object
    """
    utc = parse_utc(opts.utc)
    CTIO.date = utc
    since = parse_since(opts.since)

    # Grab the data
    data = load_data(opts)

    # Subselect the data
    sel = np.in1d(data['filter'],FILTERS)
    if opts.band in FILTERS:
        sel &= (data['filter'] == opts.band)
    data = data[sel]

    expnum,telra,teldec,band = data['expnum'],data['telra'],data['teldec'],data['filter']

    # Set the colors
    if opts.color:
        nexp = len(expnum)
        ncolors = len(COLORS)
        color_repeat = np.repeat(COLORS.keys(),nexp).reshape(ncolors,nexp)
        color_idx = np.argmax(band==color_repeat,axis=0)
        color = np.array(COLORS.values())[color_idx]
    else:
        color = COLORS['none']

    # Select the exposure of interest
    if opts.expnum:
        match = np.char.array(expnum).endswith(str(opts.expnum))
        if not match.any():
            msg = "Exposure matching %s not found"%opts.expnum
            raise ValueError(msg)
        idx = np.nonzero(match)[0][0]
    elif len(data)==0:
        idx = slice(None)
    else:
        idx = 0

    # Create the figure
    if plt.get_fignums():
        fig,ax = plt.gcf(),plt.gca()
    else:
        fig,ax = plt.subplots(figsize=FIGSIZE,dpi=DPI)
        fig.canvas.set_window_title("DECam Pointings")
    #fig,ax = plt.subplots()

    # Zenith position
    lon_zen=lmst(CTIO); lat_zen = TEL_LAT
    # Create the Basemap
    proj_kwargs = PROJ[opts.proj]
    # Centering position
    if proj_kwargs['projection'] in basemap._pseudocyl:
        ### This should work, but doesn't.
        ### Compare lon_0=-80.58345277606 to lon_0=-80.6 or lon_0=-80.5
        #lon_0=lon_zen-360*(lon_zen>180),lat_zen=0
        lon_0,lat_0 = 0,0
    else:
        lon_0,lat_0 = -lon_zen, lat_zen # Center position

    proj_kwargs.update(lon_0=lon_0,lat_0=lat_0)

    bmap = basemap.Basemap(**proj_kwargs)
    def format_coord(x,y):
        #Format matplotlib cursor to display RA, Dec
        lon,lat = safe_proj(bmap,x,y,inverse=True)
        lon += 360*(lon < 0)
        return 'ra=%1.3f, dec=%1.3f'%(lon,lat)
    plt.gca().format_coord = format_coord

    parallels = np.arange(-90.,120.,30.)
    bmap.drawparallels(parallels)
    meridians = np.arange(0.,420.,60.)
    bmap.drawmeridians(meridians)
    for mer in meridians[:-1]:
        plt.annotate(r'$%i^{\circ}$'%mer,bmap(mer,5),ha='center')
    plt.annotate('West',xy=(1.0,0.5),ha='left',xycoords='axes fraction')
    plt.annotate('East',xy=(0.0,0.5),ha='right',xycoords='axes fraction')

    # markersize defined at minimum distortion point
    if proj_kwargs['projection'] in basemap._pseudocyl:
        x1,y1=ax.transData.transform(bmap(lon_0,lat_0+DECAM))
        x2,y2=ax.transData.transform(bmap(lon_0,lat_0-DECAM))
    else:
        x1,y1=ax.transData.transform(bmap(lon_zen,lat_zen+DECAM))
        x2,y2=ax.transData.transform(bmap(lon_zen,lat_zen-DECAM))

    # Since markersize defined in "points" in scales with figsize/dpi
    size = SCALE * (y1-y2)**2

    # Scale the marker size to the size of an exposure
    exp_zorder = 10
    exp_kwargs = dict(s=size,marker='H',zorder=exp_zorder,edgecolor='k',lw=1)

    # Projected exposure locations
    x,y = safe_proj(bmap,telra,teldec)

    # Plot exposure of interest
    if len(data):
        logging.debug("Plotting exposure: %i (%3.2f,%3.2f)"%(expnum[idx],telra[idx],teldec[idx]))
        # Hacked path effect (fix if matplotlib is updated)
        bmap.scatter(x[idx],y[idx],color='w',**dict(exp_kwargs,edgecolor='w',s=70,lw=2))
        bmap.scatter(x[idx],y[idx],color=color,**dict(exp_kwargs,alpha=1.0,linewidth=2))

    # Once matplotlib is updated
    #x = bmap.scatter(x[idx],y[idx],color=color,**exp_kwargs)
    #ef = patheffects.withStroke(foreground="w", linewidth=3)
    #x.set_path_effects([ef])

    # Plot previous exposures
    nexp_kwargs = dict(exp_kwargs)
    nexp_kwargs.update(zorder=exp_zorder-1,alpha=0.2,edgecolor='none')#,lw=0)

    exp_slice = slice(None,opts.numexp)
    numexp = len(x[exp_slice])
    logging.debug("Plotting last %s exposures"%(numexp))
    bmap.scatter(x[exp_slice],y[exp_slice],color=color[exp_slice],**nexp_kwargs)

    # Plot zenith position & focal plane scale
    zen_x,zen_y = bmap(lon_zen,lat_zen)
    #zen_kwargs = dict(color='green',alpha=0.75,lw=1,zorder=0)
    zen_kwargs = dict(color='green',alpha=0.75,lw=1,zorder=1000)
    if opts.zenith:
        logging.debug("Plotting zenith: (%.2f,%.2f)"%(lon_zen,lat_zen))
        bmap.plot(zen_x,zen_y,'+',ms=10,**zen_kwargs)
        logging.debug("Plotting focal plane scale.")
        bmap.tissot(lon_zen, lat_zen, DECAM, 100, fc='none', **zen_kwargs)

        # To test exposure size
        #bmap.tissot(lon_zen, lat_zen, DECAM, 100, fc='none', **zen_kwargs)
        #bmap.scatter(*bmap(lon_zen,lat_zen),**nexp_kwargs)
        #bmap.tissot(0, 0, DECAM, 100, fc='none', **zen_kwargs)
        #bmap.scatter(*bmap(0,0),**nexp_kwargs)

    # Plot airmass circle
    if opts.airmass < 1:
        logging.warning("Airmass must be greater than one.")
        opts.airmass = np.nan
    else:
        logging.debug("Plotting airmass: %s"%opts.airmass)
        angle = airmass_angle(opts.airmass)
        bmap.tissot(lon_zen, lat_zen, angle, 100, fc='none',**zen_kwargs)

    # Moon location and phase
    (moon_ra,moon_dec),moon_phase = moon(utc)
    if opts.moon:
        logging.debug("Plotting moon: %i%%,(%.1f,%.1f)"%(moon_phase,moon_ra,moon_dec))
        moon_txt = '%i%%'%moon_phase
        #bbox = dict(boxstyle='circle,pad=0.4',fc='k',ec='k',alpha=0.25,lw=2)
        moon_kwargs = dict(zorder=exp_zorder-1,fontsize=11,va='center',ha='center',weight='bold')
        ax.annotate(moon_txt,bmap(moon_ra,moon_dec),**moon_kwargs)
        # Again old matplotlib making things difficult
        moon_kwargs2 = dict(facecolor='k',alpha=0.25,lw=2,s=2000)
        ax.scatter(*bmap(moon_ra,moon_dec),**moon_kwargs2)

    if opts.mw:
        mw_kwargs = dict(color='k')
        draw_milky_way(bmap,**mw_kwargs)

    # Plot footprint(s)
    fp_zorder=exp_zorder-1
    fp_kwargs=dict(marker='o',mew=0,mfc='none',color='k',lw=2,zorder=fp_zorder)
    if 'none' in opts.footprint:
        opts.footprint = ['none']
    if 'des' in opts.footprint:
        des_kwargs = dict(fp_kwargs,color='b')
        draw_des(bmap,**des_kwargs)
    if 'des' in opts.footprint or 'des-sn' in opts.footprint:
        sn_kwargs = dict(facecolor='none',edgecolor='b',projection=proj_kwargs['projection'],zorder=fp_zorder)
        draw_des_sn(bmap,**sn_kwargs)
    if 'smash' in opts.footprint:
        smash_kwargs = dict(facecolor='none',**exp_kwargs)
        smash_kwargs.update(zorder=exp_zorder+1)
        draw_smash(bmap,**smash_kwargs)
    if 'maglites' in opts.footprint:
        maglites_kwargs = dict(fp_kwargs,color='r')
        draw_maglites(bmap,**maglites_kwargs)
        draw_maglites2(bmap,**maglites_kwargs)
    if 'bliss' in opts.footprint:
        bliss_kwargs = dict(fp_kwargs,color='r')
        draw_bliss(bmap,**bliss_kwargs)
    if 'decals' in opts.footprint:
        decals_kwargs = dict(fp_kwargs,color='m')
        draw_decals(bmap,**decals_kwargs)
    if 'delve' in opts.footprint:
        delve_kwargs = dict(fp_kwargs,color='r')
        draw_delve(bmap,**delve_kwargs)

    # Annotate with some information
    if opts.legend:
        logging.debug("Adding info text.")
        bbox_props = dict(boxstyle='round', facecolor='white')
        textstr= "%s %s\n"%("UTC:",utc.strftime('%Y-%m-%d %H:%M:%S'))
        if len(data):
            textstr+="%s %i (%s)\n"%("Exposure:",expnum[idx],band[idx])
        textstr+="%s %i\n"%("Num. Exp.:",numexp)
        textstr+="%s (%.1f$^{\circ}$, %.1f$^{\circ}$)\n"%("Zenith:",lon_zen,lat_zen)
        textstr+="%s %s\n"%("Airmass:",np.nan_to_num(opts.airmass))
        textstr+="%s %i%% (%.1f$^{\circ}$, %.1f$^{\circ}$)\n"%("Moon:",moon_phase,moon_ra,moon_dec)
        textstr+="%s %s"%("Footprint:",', '.join(opts.footprint))

        ax.annotate(textstr, xy=(0.90,1.05), xycoords='axes fraction',
                    fontsize=10,ha='left',va='top', bbox=bbox_props)

    # Plot filter legend
    if opts.color:
        logging.debug("Adding filter legend.")
        leg_kwargs = dict(scatterpoints=1,fontsize=10,bbox_to_anchor=(0.08,0.20))
        handles, labels = [],[]
        for k in FILTERS:
            if k == 'VR' and not (band=='VR').any(): continue
            labels.append(k)
            handles.append(plt.scatter(None,None,color=COLORS[k],**exp_kwargs))
        plt.legend(handles,labels,**leg_kwargs)

    # Plot the version number
    vers_kwargs = dict(xy=(0.985,0.015),ha='right',va='bottom',
                       xycoords='figure fraction',size=8)
    plt.annotate('pointing v.%s'%__version__,**vers_kwargs)

    # Plot the author's name
    auth_kwargs = dict(xy=(0.015,0.015),ha='left',va='bottom',
                       xycoords='figure fraction',size=8)
    plt.annotate(u'\u00a9'+' %s'%__author__,**auth_kwargs)

    return bmap

if __name__ == "__main__":
    import argparse
    description = __doc__
    parser = argparse.ArgumentParser(description=description,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('expnum',nargs='?',type=int,default=None,
                        help="exposure number to plot")
    parser.add_argument('-a','--airmass',default=1.4,type=float,
                        help='draw airmass limit')
    parser.add_argument('-b','--band',default='all',choices=BANDS,
                        help='draw exposures in specific band')
    parser.add_argument('-c','--color',default=True,type=boolean,
                        help='color corresponding to filter')
    parser.add_argument('--db',default='fnal',choices=['ctio','fnal'],
                        help='database to query for exposures')
    parser.add_argument('-f','--footprint',action='append',choices=FOOTPRINTS,
                        help='footprint to draw')
    parser.add_argument('--flavor',default='object|standard',type=str,
                        help='exposure type [object,zero,dome flat,etc.]')
    parser.add_argument('-i','--infile',default=None,
                        help='list of exposures to draw')
    parser.add_argument('--legend',default=True,type=boolean,
                        help='draw figure legend')
    parser.add_argument('-m','--moon',default=True,type=boolean,
                        help='draw moon location and phase')
    parser.add_argument('--mw',action='store_true',
                        help='draw the Milky Way plane')
    parser.add_argument('-n','--numexp',default=None,type=int,
                        help='number of most recent exposures to plot')
    parser.add_argument('-o','--outfile',default=None,
                        help='output file for saving figure')
    parser.add_argument('--propid',default=None,
                        help='draw exposures from specific propid')
    parser.add_argument('--proj',default='ortho',choices=PROJ.keys(),
                        help='projection for plot')
    parser.add_argument('--refresh',nargs='?',default=None,const=60,type=int,
                        help="refresh interval for figure (seconds).")
    parser.add_argument('--since',default=None,
                        help="UTC for first exposure (defaults to 12 hours)")
    parser.add_argument('--utc',default=None,
                        help="UTC for zenith position (defaults to 'now')")
    parser.add_argument('-v','--verbose',action='store_true',
                        help='output verbosity')
    parser.add_argument('--version',action='version',
                        version='%(prog)s '+__version__)
    parser.add_argument('-z','--zenith',default=True,type=boolean,
                        help="draw zenith position")

    opts = parser.parse_args()

    # Set logging level
    logging.basicConfig(level=logging.DEBUG if opts.verbose else logging.INFO,
                        format='%(message)s',stream=sys.stdout)

    if not opts.footprint: opts.footprint = ['des']

    # Do the plotting
    m = plot(opts)

    # In interactive session
    if sys.flags.interactive: plt.ion()

    if opts.outfile:
        # Save the figure
        logging.debug("Saving figure to: %s"%opts.outfile)
        plt.savefig(opts.outfile,dpi=250)
    elif not opts.refresh:
        # Show plot
        plt.show()
    else:
        # Refresh the plot
        plt.show(block=False)
        for i in range(MAXREF): # safer than while loop
            try:
                plt.pause(opts.refresh)
            except TclError:
                # Catch the TclError thrown when window closed
                break
            logging.debug("Refreshing plot...")
            plt.cla()
            m = plot(opts)
        if i == MAXREF:
            logging.info("Reached max refresh number.")