import numpy as np
import pandas as pd
from scipy import interpolate
import glob
import math
import os
import subprocess
import sys
import gc

import glob
import pickle

import easyaccess as ea

#import AlasBabylon

import fitsio
from astropy.io import fits
import astropy.coordinates as coord
from astropy.coordinates import SkyCoord
import astropy.units as u
from astropy.table import Table, vstack

import tempfile

import matplotlib.pyplot as plt
 
%matplotlib inline

# Useful class to stop "Run All" at a cell 
#  containing the command "raise StopExecution"
class StopExecution(Exception):
    def _render_traceback_(self):
        pass

"""
Explanation: DES Y6 Deep Field Exposures:  Photometric Zeropoints tied to SDSS DR13
1. Setup
End of explanation
"""


verbose = 1
tag_des = 'Y6A2_FINALCUT'      # Official tag for DES Y6A2_FINALCUT
tag_decade = 'DECADE_FINALCUT' # Tag for DECADE
rawdata_dir = '../RawData'
zeropoints_dir='../Zeropoints'

#bandList = ['u', 'g', 'r', 'i', 'z']   # Oops!  Missing transformation eqn for Y-band...
bandList = ['u', 'g', 'r', 'i', 'z']

"""
Explanation: 2. User Input
2.1.  General User Input
End of explanation
"""


do_calc_sdss_zps = True

"""
Explanation: 2.2. Logical Variables to Indicate which Code Cells to Run
End of explanation
"""


region_name_list = [
                    'VVDSF14', 
                    'VVDSF22', 
                    'DEEP2', 
                    'SN-E', 
                    'SN-X_err', 
                    'SN-X', 
                    'ALHAMBRA2', 
                    'SN-S',
                    'SN-C', 
                    'EDFS',
                    'MACS0416',
                    'SN-S_err',
                    'COSMOS'
                    ]

region_ramin = {
                'VVDSF14':208.,
                'VVDSF22':333.,
                'DEEP2':351., 
                'SN-E':6., 
                'SN-X_err':13., 
                'SN-X':32., 
                'ALHAMBRA2':35., 
                'SN-S':39.5,
                'SN-C':50., 
                'EDFS':55.,
                'MACS0416':62.,
                'SN-S_err':83.5,
                'COSMOS':148.
               }

region_ramax = {
                'VVDSF14':212.,
                'VVDSF22':337.,
                'DEEP2':354., 
                'SN-E':12., 
                'SN-X_err':17., 
                'SN-X':38., 
                'ALHAMBRA2':39., 
                'SN-S':44.5,
                'SN-C':56., 
                'EDFS':67.,
                'MACS0416':66.,
                'SN-S_err':88.,
                'COSMOS':153.
               }

region_decmin = {
                'VVDSF14':3.,
                'VVDSF22':-1.5,
                'DEEP2':-1.5, 
                'SN-E':-46., 
                'SN-X_err':-32., 
                'SN-X':-8., 
                'ALHAMBRA2':-0.5, 
                'SN-S':-2.5,
                'SN-C':-31., 
                'EDFS':-51.,
                'MACS0416':-25.5,
                'SN-S_err':-38.,
                'COSMOS':0.5
               }


region_decmax = {
                'VVDSF14':7.,
                'VVDSF22':1.5,
                'DEEP2':1.5, 
                'SN-E':-41., 
                'SN-X_err':-29., 
                'SN-X':-3., 
                'ALHAMBRA2':2.5, 
                'SN-S':1.5,
                'SN-C':-25., 
                'EDFS':-46.,
                'MACS0416':-22.5,
                'SN-S_err':-35.,
                'COSMOS':4.0
               }


for regionName in region_name_list:
    print regionName, region_ramin[regionName], region_ramax[regionName], region_decmin[regionName], region_decmax[regionName]

"""
Explanation: 2.3. Sky Region Definitions
End of explanation
"""


# Check on TMPDIR... 
tempfile.gettempdir()

# Set tmpdir variable to $TMPDIR (for future reference)...
tmpdir = os.environ['TMPDIR']

"""
Explanation: 2.4. Check on Location of TMPDIR...
End of explanation
"""


# Create main Zeropoints directory, if it does not already exist...
if not os.path.exists(zeropoints_dir):
    os.makedirs(zeropoints_dir)

"""
Explanation: 2.5. Create Main Zeropoints Directory (if it does not already exist)...
End of explanation
"""


def DECam_tie_to_sdss(inputFile, outputFile, band, fluxObsColName, fluxerrObsColName, aggFieldColName, verbose):

    import numpy as np 
    import os
    import sys
    import datetime
    import pandas as pd
    from astropy.table import Table, vstack
    
    #validBandsList = ['u', 'g', 'r', 'i', 'z', 'Y']  # Oops!  Missing transform eqn for Y-band!
    validBandsList = ['u', 'g', 'r', 'i', 'z']

    if band not in validBandsList:
        print """Filter band %s is not currently handled...  Exiting now!""" % (band)
        return 1
    
    reqColList = ['psfMag_u','psfMagErr_u', 
                  'psfMag_g','psfMagErr_g',
                  'psfMag_r','psfMagErr_r',
                  'psfMag_i','psfMagErr_i',
                  'psfMag_z','psfMagErr_z',
                  fluxObsColName,fluxerrObsColName,aggFieldColName]        
      
    # Does the input file exist?
    if os.path.isfile(inputFile)==False:
        print """DECam_tie_to_sdss input file %s does not exist.  Exiting...""" % (inputFile)
        return 1

    # Read inputFile into a pandas DataFrame...
    print datetime.datetime.now()
    print """Reading in %s as a pandas DataFrame...""" % (inputFile)
    t = Table.read(inputFile)
    dataFrame = t.to_pandas()
    print datetime.datetime.now()
    print

    reqColFlag = 0
    colList = dataFrame.columns.tolist()
    for reqCol in reqColList:
        if reqCol not in colList:
            print """ERROR:  Required column %s is not in the header""" % (reqCol)
            reqColFlag = 1
    if reqColFlag == 1:
        print """Missing required columns in header of %s...  Exiting now!""" (inputFile)
        return 1

    # Identify column of the standard magnitude for the given band...
    magStdColName = """MAG_STD_%s""" % (band.upper())
    magerrStdColName = """MAGERR_STD_%s""" % (band.upper())
    
    # Transform SDSS mags to DES mags for the given band...
    dataFrame = transSDSStoDES(dataFrame, band, magStdColName, magerrStdColName)

    # Add a 'MAG_OBS' column and a 'MAG_DIFF' column to the pandas DataFrame...
    dataFrame['MAG_OBS'] = -2.5*np.log10(dataFrame[fluxObsColName])
    dataFrame['MAG_DIFF'] = dataFrame[magStdColName]-dataFrame['MAG_OBS']


    ###############################################
    # Aggregate by aggFieldColName
    ###############################################

    # Make a copy of original dataFrame...
    df = dataFrame.copy()

    # Create an initial mask...
    mask1 = ( (df[magStdColName] >= 0.) & (df[magStdColName] <= 25.) )
    mask1 = ( mask1 & (df[fluxObsColName] > 10.) & (df['FLAGS'] < 2) & (np.abs(df['SPREAD_MODEL']) < 0.01))
    if magerrStdColName != 'None':  
        mask1 = ( mask1 & (df[magerrStdColName] < 0.1) )
    magDiffGlobalMedian = df[mask1]['MAG_DIFF'].median()
    magDiffMin = magDiffGlobalMedian - 5.0
    magDiffMax = magDiffGlobalMedian + 5.0
    mask2 = ( (df['MAG_DIFF'] > magDiffMin) & (df['MAG_DIFF'] < magDiffMax) )
    mask = mask1 & mask2

    # Iterate over the copy of dataFrame 3 times, removing outliers...
    #  We are using "Method 2/Group by item" from
    #  http://nbviewer.jupyter.org/urls/bitbucket.org/hrojas/learn-pandas/raw/master/lessons/07%20-%20Lesson.ipynb
    print "Sigma-clipping..."
    niter = 0
    for i in range(3):

        niter = i + 1
        print """   iter%d...""" % ( niter )

        # make a copy of original df, and then delete the old one...
        newdf = df[mask].copy()
        del df
        # group by aggFieldColName...
        grpnewdf = newdf.groupby([aggFieldColName])
        
        # add/update new columns to newdf
        print datetime.datetime.now()
        newdf['Outlier']  = grpnewdf['MAG_DIFF'].transform( lambda x: abs(x-x.mean()) > 3.00*x.std() )
        #newdf['Outlier']  = grpnewdf['MAG_DIFF'].transform( lambda x: abs(x-x.mean()) > 2.00*x.std() )
        print datetime.datetime.now()
        del grpnewdf
        print datetime.datetime.now()
        #print newdf

        nrows = newdf['MAG_DIFF'].size
        print """  Number of rows remaining:  %d""" % ( nrows )

        df = newdf
        mask = ( df['Outlier'] == False )  


    # Perform pandas grouping/aggregating functions on sigma-clipped Data Frame...
    print datetime.datetime.now()
    print 'Performing grouping/aggregating functions on sigma-clipped pandas DataFrame...'
    groupedDataFrame = df.groupby([aggFieldColName])
    magZeroMedian = groupedDataFrame['MAG_DIFF'].median()
    magZeroMean = groupedDataFrame['MAG_DIFF'].mean()
    magZeroStd = groupedDataFrame['MAG_DIFF'].std()
    magZeroNum = groupedDataFrame['MAG_DIFF'].count()
    magZeroErr = magZeroStd/np.sqrt(magZeroNum-1)
    print datetime.datetime.now()
    print

    # Rename these pandas series...
    magZeroMedian.name = 'MAG_ZERO_MEDIAN'
    magZeroMean.name = 'MAG_ZERO_MEAN'
    magZeroStd.name = 'MAG_ZERO_STD'
    magZeroNum.name = 'MAG_ZERO_NUM'
    magZeroErr.name = 'MAG_ZERO_MEAN_ERR'

    # Also, calculate group medians for all columns in df that have a numerical data type...
    numericalColList = df.select_dtypes(include=[np.number]).columns.tolist()
    groupedDataMedian = {}
    for numericalCol in numericalColList:
        groupedDataMedian[numericalCol] = groupedDataFrame[numericalCol].median()
        groupedDataMedian[numericalCol].name = """%s_MEDIAN""" % (numericalCol)

    # Create new data frame containing all the relevant aggregate quantities
    #newDataFrame = pd.concat( [magZeroMedian, magZeroMean, magZeroStd, \
    #                           magZeroErr, magZeroNum], \
    #                           join='outer', axis=1 )
    seriesList = []
    for numericalCol in numericalColList:
        seriesList.append(groupedDataMedian[numericalCol])
    seriesList.extend([magZeroMedian, magZeroMean, magZeroStd, \
                               magZeroErr, magZeroNum])
    #print seriesList
    newDataFrame = pd.concat( seriesList, join='outer', axis=1 )
                               

    #newDataFrame.index.rename('FILENAME', inplace=True)

    # Saving catname-based results to output files...
    print datetime.datetime.now()
    print """Writing %s output file (using pandas to_csv method)...""" % (outputFile)
    newDataFrame.to_csv(outputFile, float_format='%.4f')
    print datetime.datetime.now()
    print

    return 0

# Transform SDSS mags into DES mags for this filter band...
def transSDSStoDES(dataFrame, band, magStdColName, magerrStdColName):

    import numpy as np 
    import pandas as pd
    from collections import OrderedDict as odict

    # Transformation coefficients (updated based on fit to DES).
    # mag_des = mag_sdss + A[mag][0]*color_sdss + A[mag][1]
    # (based on A. Drlica-Wagner's https://github.com/kadrlica/desqr/blob/master/desqr/calibrate.py).
    # Values come from DES-doc#9420 (u-band) and the DES DR2 paper (g-,r-,i-,z-bands).
    SDSS = odict([
            ('u', [-0.350, +0.466, -0.479, ]), # [+0.2 < (g-r)_sdss <= 1.4]
            ('g', [-0.061, +0.008]), # [-1.0 < (g-i)_sdss <= 3.5]
            ('r', [-0.155, -0.007]), # [-0.4 < (r-i)_sdss <= 2.0]
            ('i', [-0.166, +0.032]), # [-0.4 < (r-i)_sdss <= 2.0]
            ('z', [-0.056, +0.027]), # [-0.4 < (r-i)_sdss <= 2.0]
            ]) 
    
    A = SDSS

    if band is 'u':
        # u_des = u_sdss - 0.350*(g-r)_sdss**2 + 0.466*(g-r)_sdss - 0.479    [+0.2 < (g-r)_sdss <= 1.4]
        dataFrame[magStdColName] = dataFrame['psfMag_u']+\
            A[band][0]*(dataFrame['psfMag_g']-dataFrame['psfMag_r'])*(dataFrame['psfMag_g']-dataFrame['psfMag_r'])+\
            A[band][1]*(dataFrame['psfMag_g']-dataFrame['psfMag_r'])+\
            A[band][2]
        dataFrame[magerrStdColName] = dataFrame['psfMagErr_u']  # temporary
        mask  = ( (dataFrame['psfMag_g']-dataFrame['psfMag_r']) > +0.2)
        mask &= ( (dataFrame['psfMag_g']-dataFrame['psfMag_r']) <= 1.4)
    elif band is 'g':
        # g_des = g_sdss - 0.061*(g-i)_sdss + 0.008   [-1.0 < (g-i)_sdss <= 3.5]
        dataFrame[magStdColName] = dataFrame['psfMag_g']+\
            A[band][0]*(dataFrame['psfMag_g']-dataFrame['psfMag_i'])+A[band][1]
        dataFrame[magerrStdColName] = dataFrame['psfMagErr_g']  # temporary
        mask  = ( (dataFrame['psfMag_g']-dataFrame['psfMag_i']) > -1.0)
        mask &= ( (dataFrame['psfMag_g']-dataFrame['psfMag_i']) <= 3.5)
    elif band is 'r':
        # r_des = r_sdss - 0.155*(r-i)_sdss - 0.007   [-0.4 < (r-i)_sdss <= 2.0]
        dataFrame[magStdColName] = dataFrame['psfMag_r']+\
            A[band][0]*(dataFrame['psfMag_r']-dataFrame['psfMag_i'])+A[band][1]
        dataFrame[magerrStdColName] = dataFrame['psfMagErr_r']  # temporary
        mask  = ( (dataFrame['psfMag_r']-dataFrame['psfMag_i']) > -0.4)
        mask &= ( (dataFrame['psfMag_r']-dataFrame['psfMag_i']) <= 2.0)
    elif band is 'i':
        # i_des = i_sdss - 0.166*(r-i)_sdss + 0.032   [-0.4 < (r-i)_sdss <= 2.0]
        dataFrame[magStdColName] = dataFrame['psfMag_i']+\
            A[band][0]*(dataFrame['psfMag_r']-dataFrame['psfMag_i'])+A[band][1]
        dataFrame[magerrStdColName] = dataFrame['psfMagErr_i']  # temporary
        mask  = ( (dataFrame['psfMag_r']-dataFrame['psfMag_i']) > -0.4)
        mask &= ( (dataFrame['psfMag_r']-dataFrame['psfMag_i']) <= 2.0)
    elif band is 'z':
        # z_des = z_sdss - 0.056*(r-i)_sdss + 0.027   [-0.4 < (r-i)_sdss <= 2.0]
        dataFrame[magStdColName] = dataFrame['psfMag_z']+\
            A[band][0]*(dataFrame['psfMag_r']-dataFrame['psfMag_i'])+A[band][1]
        dataFrame[magerrStdColName] = dataFrame['psfMagErr_z']  # temporary
        mask  = ( (dataFrame['psfMag_r']-dataFrame['psfMag_i']) > -0.4)
        mask &= ( (dataFrame['psfMag_r']-dataFrame['psfMag_i']) <= 2.0)
    else:
        msg = "Unrecognized band: %s "%band
        raise ValueError(msg)

    dataFrame = dataFrame[mask].copy()

    return dataFrame

"""
Explanation: 3. Useful Modules
End of explanation
"""


%%time 

if do_calc_sdss_zps: 
    
    fluxObsColName =    'FLUX_PSF' 
    fluxerrObsColName = 'FLUXERR_PSF' 
    aggFieldColName =   'FILENAME' 
    
    subdir = """DES_%s""" % (tag_des)

    tmpdir = os.environ['TMPDIR']


    for regionName in region_name_list:
    
        print
        print """# # # # # # # # # # # # # # #"""
        print """Working on region %s""" % (regionName)
        print """# # # # # # # # # # # # # # #"""
        print
        
        for band in bandList:
                    
            input_file_template = """cat_%s.%s.?.%s.sdss.fits""" % (subdir, regionName, band)
            input_file_template = os.path.join(rawdata_dir, 'ExpCatFITS', subdir, input_file_template)
            input_file_list = glob.glob(input_file_template)
            input_file_list = np.sort(input_file_list)

            if np.size(input_file_list) == 0:
                print "No files matching template %s" % (input_file_template)
    
            for inputFile in input_file_list:
            
                print inputFile 

                if os.path.exists(inputFile):

                    #outputFile = os.path.splitext(inputFile)[0] + '.zps.csv'
                    #print outputFile
                    
                    outputFile = os.path.splitext(os.path.basename(inputFile))[0]+'.csv'
                    outputFile = 'zps_' + outputFile[4:]
                    outputFile = os.path.join(zeropoints_dir, outputFile)

                    print outputFile    

                    status = DECam_tie_to_sdss(inputFile, outputFile, 
                                                    band, 
                                                    fluxObsColName, fluxerrObsColName, 
                                                    aggFieldColName, verbose)
                    if status > 0:
                        print 'ERROR:  %s FAILED!  Continuing...'
            
                else:
            
                    print """%s does not exist...  skipping...""" % (inputFile)
                
                print

"""
Explanation: 4. Zeropoints by tying to DES-transformed SDSS DR13 Stars
We will first work with the DES data, and then we will repeat for the DECADE data.
DES:
Calculate zeropoints region by region...
End of explanation
"""


%%time 

if do_calc_sdss_zps: 
        
    subdir = """DES_%s""" % (tag_des)

    tmpdir = os.environ['TMPDIR']

    
    for band in bandList:
                    
        print
        print """# # # # # # # # # # # # # # #"""
        print """Working on band %s""" % (band)
        print """# # # # # # # # # # # # # # #"""
        print

        outputFile = """zps_%s.%s.sdss.csv""" % (subdir, band)
        outputFile = outputFile = os.path.join(zeropoints_dir, outputFile)
    
        input_file_template = """zps_%s.*.?.%s.sdss.csv""" % (subdir, band)
        input_file_template = os.path.join(zeropoints_dir, input_file_template)
        input_file_list = glob.glob(input_file_template)
        input_file_list = np.sort(input_file_list)

        if np.size(input_file_list) == 0:
            print "No files matching template %s" % (input_file_template)
            continue
    
        df_comb = pd.concat(pd.read_csv(inputFile) for inputFile in input_file_list)

        outputFile = """zps_%s.%s.sdss.csv""" % (subdir, band)
        outputFile = outputFile = os.path.join(zeropoints_dir, outputFile)
        print outputFile
        df_comb.to_csv(outputFile, index=False)
        
        del df_comb


"""
Explanation: Combine region-by-region results into a single file...
End of explanation
"""


%%time 

if do_calc_sdss_zps: 
    
    fluxObsColName =    'FLUX_PSF' 
    fluxerrObsColName = 'FLUXERR_PSF' 
    aggFieldColName =   'FILENAME' 
    
    subdir = """%s""" % (tag_decade)

    tmpdir = os.environ['TMPDIR']


    for regionName in region_name_list:
    
        print
        print """# # # # # # # # # # # # # # #"""
        print """Working on region %s""" % (regionName)
        print """# # # # # # # # # # # # # # #"""
        print
        
        for band in bandList:
                    
            input_file_template = """cat_%s.%s.?.%s.sdss.fits""" % (subdir, regionName, band)
            input_file_template = os.path.join(rawdata_dir, 'ExpCatFITS', subdir, input_file_template)
            input_file_list = glob.glob(input_file_template)
            input_file_list = np.sort(input_file_list)

            if np.size(input_file_list) == 0:
                print "No files matching template %s" % (input_file_template)
    
            for inputFile in input_file_list:
            
                print inputFile 

                if os.path.exists(inputFile):

                    #outputFile = os.path.splitext(inputFile)[0] + '.zps.csv'
                    #print outputFile
                    
                    outputFile = os.path.splitext(os.path.basename(inputFile))[0]+'.csv'
                    outputFile = 'zps_' + outputFile[4:]
                    outputFile = os.path.join(zeropoints_dir, outputFile)

                    print outputFile    

                    status = DECam_tie_to_sdss(inputFile, outputFile, 
                                                    band, 
                                                    fluxObsColName, fluxerrObsColName, 
                                                    aggFieldColName, verbose)
                    if status > 0:
                        print 'ERROR:  %s FAILED!  Continuing...'
            
                else:
            
                    print """%s does not exist...  skipping...""" % (inputFile)
                
                print

"""
Explanation: DECADE:
Calculate zeropoints region by region...
End of explanation
"""


%%time 

if do_calc_sdss_zps: 
        
    subdir = """%s""" % (tag_decade)

    tmpdir = os.environ['TMPDIR']

    
    for band in bandList:
                    
        print
        print """# # # # # # # # # # # # # # #"""
        print """Working on band %s""" % (band)
        print """# # # # # # # # # # # # # # #"""
        print

        outputFile = """zps_%s.%s.sdss.csv""" % (subdir, band)
        outputFile = outputFile = os.path.join(zeropoints_dir, outputFile)
    
        input_file_template = """zps_%s.*.?.%s.sdss.csv""" % (subdir, band)
        input_file_template = os.path.join(zeropoints_dir, input_file_template)
        input_file_list = glob.glob(input_file_template)
        input_file_list = np.sort(input_file_list)

        if np.size(input_file_list) == 0:
            print "No files matching template %s" % (input_file_template)
            continue
    
        df_comb = pd.concat(pd.read_csv(inputFile) for inputFile in input_file_list)

        outputFile = """zps_%s.%s.sdss.csv""" % (subdir, band)
        outputFile = outputFile = os.path.join(zeropoints_dir, outputFile)
        print outputFile
        df_comb.to_csv(outputFile, index=False)
        
        del df_comb


"""
Explanation: Combine region-by-region results into a single file...
End of explanation
"""


