#! /usr/bin/env python
#
# $Rev::                                                              $:  
# $Author::                                                           $:  
# $LastChangedDate::                                                  $:
#
# make a set of Meshes from a donut root file, applying a number of cuts
#
import numpy
from scipy.spatial import cKDTree
import argparse
import time
import pdb
import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True
from ROOT import TTree, TFile, TCanvas, gStyle, TGraph2D,TH1D

def goodPoints(tree,extnameThis,ifileThis,cutDict):
    """ go through the data, apply the cuts, return a dictionary with good points
    """

    #  spool through the data, and get all the donuts with the desired extname
    #  and nFile (assumes that the order of the tree is by file, and then by extname
    #  so we can spool more quickly
    #
    #  put the output in a numpy record array with all desired pieces of info, for
    #  donuts that pass the simple, individual, cuts
    #  then look collectively at the data to apply the collective cuts
    #  

    # list for structured array of data
    structList = [('zern4',numpy.float32),
                  ('zern5',numpy.float32),
                  ('zern6',numpy.float32),
                  ('zern7',numpy.float32),
                  ('zern8',numpy.float32),
                  ('zern9',numpy.float32),
                  ('zern10',numpy.float32),
                  ('xdecam',numpy.float32),
                  ('ydecam',numpy.float32),
                  ('ix',numpy.int),
                  ('iy',numpy.int),
                  ('nele',numpy.float32),
                  ('chi2',numpy.float32),
                  ('ncalcall',numpy.int),
                  ('istamp',numpy.int),
                  ('weight',numpy.float32),
                  ('extname',str)]

    # cuts
    neleCut = cutDict["neleCut"]
    sumsqzernCut = cutDict["sumsqzernCut"]
    ncalcallCut = cutDict["ncalcallCut"]
    rdecamCut = cutDict["rdecamCut"]
    absz4LoCut = cutDict["absz4LoCut"]
    absz4HiCut = cutDict["absz4HiCut"]
    madCut = cutDict["madCut"]
    kNNCut = cutDict["kNNCut"]

    # lists for temporary storage
    dataList = []
    dataGoodList = []

    # found any? flag
    foundOne = False
    jj = 0

    # loop over donut's in the tree
    for donut in tree:

        # get extname and ifile, and check
        extname = donut.extname
        ifile = int(donut.ifile)
        
        if extname==extnameThis and ifile==ifileThis:

            foundOne = True

            # unpack variables
            zern4 = donut.zern4
            zern5 = donut.zern5
            zern6 = donut.zern6 
            zern7 = donut.zern7 
            zern8 = donut.zern8 
            zern9 = donut.zern9
            zern10 = donut.zern10
            ix = donut.ix
            iy = donut.iy

            chi2 = donut.chi2
            fitstat = donut.fitstat
            nele = donut.nele
            bkgd = donut.bkgd
            ncalcall = donut.ncalcall
            istamp = donut.istamp
            weight = 1.0

            # calculate secondary variables
            xdecam,ydecam = 0.0,0.0
            rdecam = numpy.sqrt(xdecam*xdecam + ydecam*ydecam)

            # apply simple cuts
            if nele > neleCut and numpy.sqrt(zern5*zern5+zern6*zern6+zern7*zern7+zern8*zern8) < sumsqzernCut and ncalcall < ncalcallCut and rdecam < rdecamCut and numpy.abs(zern4) > absz4LoCut and numpy.abs(zern4) < absz4HiCut :

                jj = jj + 1
                print "got to number ",jj
                dummyArray = numpy.zeros(1,structList)
                # now collect these good guys in a list of arrays
                dummyArray['zern4'] =   zern4
                dummyArray['zern5'] =   zern5
                dummyArray['zern6'] =   zern6
                dummyArray['zern7'] =   zern7
                dummyArray['zern8'] =   zern8
                dummyArray['zern9'] =   zern9
                dummyArray['zern10'] =   zern10
                dummyArray['xdecam'] =   xdecam
                dummyArray['ydecam'] =   ydecam
                dummyArray['ix'] =   ix
                dummyArray['iy'] =   iy
                dummyArray['nele'] =   nele
                dummyArray['chi2'] =   chi2
                dummyArray['ncalcall'] =   ncalcall
                dummyArray['istamp'] =   istamp
                dummyArray['extname'] =   extname

                dataList.append(dummyArray)
                                
        else:
            # now the TTree is on to the next guy - so we can just leave now
            if foundOne:
                continue  # want to jump out of for donut in tree Loop (check this)
         


def doMesher(options):
    """ makes Mesh data files    
    """

    # get the TTree
    f = TFile(options.inputRootFile)
    tree = f.Get("t")

    # use goodPoints to make the mesh data files
    cutDict = {"neleCut":3.5,
               "sumsqzernCut":3.0,
               "ncalcallCut":200,
               "rdecamCut":225.,
               "absz4LoCut":4.0,
               "absz4HiCut":15.0,
               "madCut":4.0,
               "kNNCut":3.0}

    extList = options.extNamesInput

    # loop over files
    resultsDict = {}
    for ifile in options.fileList:
        
        goodDataArrayDict = {}
        canvasesDict = {}
        # loop over extname's
        for extname in extList:

            goodDataArray,canvasList = goodPoints(tree,extname,ifile,cutDict)
            goodDataArrayDict["extname"] = goodDataArray
            canvasesDict["extname"] = canvasList

        # done with extensions
        # fill the resultDict, which is a list with two dicts by extname
        # containing the goodDataArrays and the list of two canvases (dats a mouthful)
        resultsDict[ifile] = [goodDataArrayDict,canvasesDict]
        

            
# if calling from the command line, use the argument parser    
if __name__ == '__main__':

    parser = argparse.ArgumentParser(prog='decamMesher')
    parser.add_argument("-i", "--inputRootFile",
                        dest="inputRootFile",
                        default=None,
                        help="input root file ")
    parser.add_argument("-f", "--fileList",
                        dest="fileList",nargs='*',type=int,
                        help="list of file numbers, one mesh per file")
    parser.add_argument("-o", "--outputPrefix",
                        dest="outputPrefix",
                        default=None,
                        help="output File Prefix ")
    parser.add_argument("-e","--extNamesInput",nargs='*',
                    dest="extNamesInput",
                    default=None,
                    help="List of Extensions - overrides defaults")
    parser.add_argument("-doFandA", "--doFandA",
                        dest="doFandA",action="store_true",default=False,
                        help="do Focus and Alignment chips")

    parser.add_argument("-doScience", "--doScience",
                        dest="doScience",action="store_true",default=False,
                        help="do Science chips")


    
    # collect the options 
    options = parser.parse_args()

    # call it!
    doMesher(options)