import numpy
from scipy import signal
import sys, re
sys.argv.append( '-b-' )
import ROOT, array

##from jungfrauZeug import getNextEventSingle
from psana import *

from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

from cluster import Cluster, BuildClusters


camera = 0
path = "/reg/d/psdm/xcs/xcs11016/usrdaq/jungfrau/"
resPath = "/reg/d/psdm/cxi/cxi11216/results/jungfrau"

##psiFileA = sys.argv[1] ## assume A
##psiFileA = re.sub("000000.dat", "00000%d.dat" %(rank), psiFileA)
##fA = open(path+psiFileA,"rb")
run = eval(sys.argv[1]) 
## point to xtc
darkRun0 = eval(sys.argv[2]) ## point to ped?
darkRun1 = darkRun0 + 3
darkRun2 = darkRun0 + 4
ds  = DataSource('exp=cxi11216:run=%d:smd' %(run))
cspad = Detector('Dg3CsPad2x2', ds.env())
##jungfrau = Detector('Jungfrau', ds.env())
imp = Detector('Dg3Imp', ds.env())
fir = numpy.loadtxt('fir8.fir')
config = ds.env().configStore()
evrs = []
for key in config.keys():
    if key.type() == EvrData.ConfigV7:
        evrs.append(key.src())

try:
    label = "_" + sys.argv[3]
    if 'lumi' in label:
        lumi = True
    else:
        lumi = False
    if 'noCM' in label:
        doCM = False
    else:
        doCM = True
    if 'noClus' in label:
        doClusters = False
    else:
        doClusters = True
    if 'gainCorr' in label:
        corrGain = True
    else:
        corrGain = False
    if 'stitch' in label:
        stitch = True
    else:
        stitch = False
##    if 'noCM' in label:
##        doCM = False
##    else:
##        doCM = True
    
except:
    label = ""
    print "good luck having this work"

nonStandardGains = []
## not e.g. 31

ped = True
try:
    ped0 = (numpy.load("r%d_stats_gain0_ped.npy" %(darkRun0)).astype(int))&0x3fff
    ped1 = (numpy.load("r%d_stats_gain1_ped.npy" %(darkRun1)).astype(int))&0x3fff
    ped2 = (numpy.load("r%d_stats_gain2_ped.npy" %(darkRun2)).astype(int))&0x3fff
    g0cut = 1<<14
    g1cut = 2<<14
    g2cut = 3<<14
    ped0[256:, 256:] = 0
    ped1[256:, 256:] = 0
    ped2[256:, 256:] = 0

except:
    print "no pedestal found"
    ped = None

##localPath = "/reg/neh/home4/philiph/psana/jungfrau/"
localPath = "/reg/neh/home4/philiph/psana/jungfrau/singleModule/"
try:
    a = numpy.load(localPath + "summary_40_regionA_m_csnVsJF.npy")
    b = numpy.load(localPath + "summary_40_regionB_m_csnVsJF.npy")
    rAp0 = []
    rAp1 = []
    rBp0 = []
    rBp1 = []
    for gains in range(3):
        rAp0.append(a[:, gains*4])
        rAp1.append(a[:, gains*4+1])
        rBp0.append(b[:, gains*4])
        rBp1.append(b[:, gains*4+1])
    rAp0 = numpy.array(rAp0)
    rAp1 = numpy.array(rAp1)
    rBp0 = numpy.array(rBp0)
    rBp1 = numpy.array(rBp1)
except:
    print "could not open %s or something like that" %(localPath + "summary_40_regionA_m_csnVsJF.npy")

##tFile = ROOT.TFile('clusters_rank%d_r%d_d%d.root' %(rank, run, pedRun),'RECREATE')
##tFile = ROOT.TFile('/reg/d/psdm/xcs/xcs11016/results/philiph/clusters_2sigma_rank%d_r%d.root' %(rank, run),'RECREATE')
tname = resPath + '/clusters%s_rank%d_r%d.root' %(label, rank, run)
if not doClusters:
    tname = re.sub("/clusters", "/regions", tname)
tFile = ROOT.TFile(tname, 'RECREATE')

tTree = ROOT.TTree('ntup','ntup')
maxClusters = 100000
tnEvent = array.array('i', [0])
tnClusters = array.array('i', [0])
tCamera = array.array('i',maxClusters*[0])
tcOffset = array.array('f',maxClusters*[0])
trOffset = array.array('f',maxClusters*[0])
tcFlat = array.array('f',maxClusters*[0])
trFlat = array.array('f',maxClusters*[0])
tnPixels = array.array('i',maxClusters*[0])
tnMaskedPixels = array.array('i',maxClusters*[0])
tenergy = array.array('f',maxClusters*[0])
tseedEnergy = array.array('f',maxClusters*[0])
tseedCol = array.array('i', maxClusters*[0]) 
tseedRow = array.array('i', maxClusters*[0]) 
teTotalNoCuts = array.array('f', maxClusters*[0])
teSecondaryPixelNoCuts = array.array('f', maxClusters*[0])
tisSquare = array.array('i', maxClusters*[0])
if lumi:
    tnEvent = array.array('i', [0])
    tSeconds = array.array('i', [0])
    tNSeconds = array.array('i', [0])
    tFiducial = array.array('i', [0])
    tKicked = array.array('i', [0])
    tFlux = array.array('f', [0])
    tCsN = array.array('i',[0])
    tCsTot = array.array('f', [0])
    tCsAve25 = array.array('f', [0])
    regionSize = 20
    regionPixels = regionSize**2
    tRegionA = array.array('i', regionPixels*[0])
    tRAGain = array.array('i', regionPixels*[0])
    tRAstitch = array.array('f', regionPixels*[0])
    tRegionB = array.array('i', regionPixels*[0])
    tRBGain = array.array('i', regionPixels*[0])
    tRBstitch = array.array('f', regionPixels*[0])


tTree.Branch('nEvent', tnEvent, 'nEvent/I')
tTree.Branch('nClusters', tnClusters, 'nClusters/I')
tTree.Branch('camera', tCamera, 'camera[nClusters]/I')
tTree.Branch('energy', tenergy, 'energy[nClusters]/F')
tTree.Branch('cOffset', tcOffset, 'cOffset[nClusters]/F')
tTree.Branch('rOffset', trOffset, 'rOffset[nClusters]/F')
tTree.Branch('cFlat', tcFlat, 'cFlat[nClusters]/F')
tTree.Branch('rFlat', trFlat, 'rFlat[nClusters]/F')
tTree.Branch('nPixels', tnPixels, 'nPixels[nClusters]/I')
tTree.Branch('nMaskedPixels', tnMaskedPixels, 'nMaskedPixels[nClusters]/I')
tTree.Branch('seedEnergy', tseedEnergy, 'seedEnergy[nClusters]/F')
tTree.Branch('seedCol', tseedCol, 'seedCol[nClusters]/I')
tTree.Branch('seedRow', tseedRow, 'seedRow[nClusters]/I')
tTree.Branch('eTotalNoCuts', teTotalNoCuts, 'eTotalNoCuts[nClusters]/F')
tTree.Branch('eSecondaryPixelNoCuts', teSecondaryPixelNoCuts, 'eSecondaryPixelNoCuts[nClusters]/F')
tTree.Branch('isSquare', tisSquare, 'isSquare[nClusters]/I')

if lumi:
    tTree.Branch('nEvent', tnEvent, 'nEvent/I')
    tTree.Branch('seconds', tSeconds, 'seconds/I')
    tTree.Branch('nSeconds', tNSeconds, 'nSeconds/I')
    tTree.Branch('fiducial', tFiducial, 'fiducial/I')
    tTree.Branch('kicked', tKicked, 'kicked/I')
    tTree.Branch('flux', tFlux, 'flux/F')
    tTree.Branch('csN', tCsN, 'csN/I')
    tTree.Branch('csTot', tCsTot, 'csTot/F')
    tTree.Branch('csAve25', tCsAve25, 'csAve25/F')
    tTree.Branch('regionA', tRegionA, 'regionA[%d]/I' %(regionPixels))
    tTree.Branch('rAGain', tRAGain, 'rAGain[%d]/I' %(regionPixels))
    tTree.Branch('rAstitch', tRAstitch, 'rAstitch[%d]/F' %(regionPixels))
    tTree.Branch('regionB', tRegionB, 'regionB[%d]/I' %(regionPixels))
    tTree.Branch('rBGain', tRBGain, 'rBGain[%d]/I' %(regionPixels))
    tTree.Branch('rBstitch', tRBstitch, 'rBstitch[%d]/F' %(regionPixels))

    if True:
        xA = 123
        yA = 71
        xB = 356
        yB = 154
    else:
        xA = 740
        yA = 420
        xB = 740
        yB = 460

singles = ROOT.TH1F("singles", "singles", 500, 0, 1000)
squares = ROOT.TH1F("squares", "squares", 500, 0, 1000)

gainCorr = None
if run not in nonStandardGains:
    seedCut = 100 ##???
    neighborCut = 33 ##???
    if corrGain:
        gainCorrFile = "gainPixelCorr_passZeroL_r5_c0.npy"
        gainCorr = numpy.load(gainCorrFile)
        print "using gain correction %s, rank %d" %(gainCorrFile, rank)
else:
    ## high gain: Cu at 778 ADU
    seedCut = 100 ## bit under half of 220 for Mo L line ##400 ##???
    neighborCut = 54 ## 3sigma
##    neighborCut = 36 ## 2sigma
##    gainCorrFile = "gainPixelCorr_passZeroL_r21_c0.npy"
##    gainCorr = numpy.load(gainCorrFile)
##    print "using gain correction %s" %(gainCorrFile)

centroid = None
rows = 1024/2
cols = 1024
banks = 4
bankN = cols/banks



nSkippedForFlux = 0
##maxNevents = 0
maxNevents = 400000

nEvt = 0
for nEvent,evt in enumerate(ds.events()):
    if nEvent%size != rank:
        continue
    if nEvent>maxNevents:
        break
##    epixSum = numpy.zeros(cameras)
##    normedSum = numpy.zeros(cameras)
##    eArray = []
##    nArray = []
##    fArray = []
##    for c in range(cameras):
##        eArray.append([])
##        nArray.append([])
        
    nGoodEvt = 0
    nClusters = 0 ## make ntuple happy outside doClusters mode
    flux = signal.lfilter(fir, 1, imp.waveform(evt)[0])
##    if nEvent%100==0:
##        print "event 100 mod:", nEvent, flux
##        print flux[len(fir):].max(), flux[len(fir):].argmax()
    flux = flux[len(fir):].max()
##    if nEvent%100==0:
##        print "event:", nEvent, flux
    csAve = cspad.calib(evt).mean()
##    print nEvent, csAve
    jungfrau = evt.get(Jungfrau.ElementV1, Source('Jungfrau'))
##    frame = jungfrau.frame()[0]
    frame = numpy.copy(jungfrau.frame()[0])
    frame[256:, 256:] = 0
##    print frame
    
    
##    if csAve>100:
##        nSkippedForFlux += 1
##        continue
##    fArray.append(csAve)


    if True:
        rawFrame = numpy.copy(frame)
        fG0 = frame<g0cut
        fG1 = (frame>=g0cut) & (frame<g1cut)
        fG2 = frame>=g2cut
        fGval = fG0*1 + fG1*2 + fG2*3
        frame = (frame&0x3fff).astype('int')
        frame[fG0] = frame[fG0] - ped0[fG0]
        frame[fG1] = ped1[fG1] - frame[fG1]
        frame[fG2] = ped2[fG2] - frame[fG2]
        if gainCorr is not None:
            frame = (frame*gainCorr).astype('int')

    if ped is not None:
        if doCM:
            for r in range(rows):
                colOffset = 0
                for b in range(1, banks+1):
                    try:
                        rowCM = numpy.median(frame[r, colOffset:colOffset + bankN][frame[r, colOffset:colOffset + bankN]<seedCut])
                        if not numpy.isnan(rowCM):
                            ## e.g. found no pixels below seedCut
                            ##                    print r, rowCM, colOffset, b, bankN
                            frame[r, colOffset:colOffset + bankN] -= rowCM*(rowCM<seedCut-1)
                    except:
                        rowCM = -666
                        print "rowCM problem"
                        print frame[r, colOffset:colOffset + bankN]
                    colOffset += bankN


        if doClusters:##csAve >0.:
            bc = BuildClusters(frame, seedCut, neighborCut)
            fc = bc.findClusters()
            if camera == 0:
                nClusters = 0
                tnEvent[0] = nEvent
            for c in fc:
                if c.goodCluster and nClusters<maxClusters:
                    tCamera[nClusters] = camera
                    seedCol = c.seedCol
                    seedRow = c.seedRow
                    centroidRow, centroidCol = c.centroid()
                    tcOffset[nClusters] = centroidCol-seedCol
                    trOffset[nClusters] = centroidRow-seedRow
                    sC = 0
                    sR = 0
                    if (centroid is not None):
                        sC, sR = centroid.getSmoothedCentroid(tcOffset[nClusters], trOffset[nClusters])
                    tcFlat[nClusters] = sC
                    trFlat[nClusters] = sR
                    tnPixels[nClusters] = c.nPixels
                    tnMaskedPixels[nClusters] = c.nMaskedPixels
                    tenergy[nClusters] = c.eTotal
                    tseedEnergy[nClusters] = c.seedEnergy
                    tseedCol[nClusters] = seedCol
                    tseedRow[nClusters] = seedRow
                    teTotalNoCuts[nClusters] = c.eTotalNoCuts
                    teSecondaryPixelNoCuts[nClusters] = c.eSecondaryPixelNoCuts
                    tisSquare[nClusters] = int(c.isSquare())
                    nClusters += 1
                    if c.nPixels == 1:
                        singles.Fill(c.eTotal)
                    if c.isSquare:
                        squares.Fill(c.eTotal)
        else:
            pass ##?
        tnClusters[0] = nClusters
        if lumi:
            tnEvent[0] = nEvent
            evtId = evt.get(EventId)
            tSeconds[0] = evtId.time()[0]
            tNSeconds[0] = evtId.time()[1]
            tFiducial[0] = evtId.fiducials()
            tKicked[0] = 0
            if evrs:
                evr = evt.get(EvrData.DataV4, evrs[0])
                for ec in evr.fifoEvents():
                    if ec.eventCode() == 162:
                        print "found a kicked event"
                        tKicked[0] = 1
##            f3 = ipm03.channel(evt)
##            for i in range(4):## do I need to do this?
            tFlux[0] = flux
            tCsTot[0] = -666
            tCsAve25[0] = -666
            tCsN[0] = - 666
            try:
                data = cspad.calib(evt)
                tCsTot[0] = data[(data)>25].sum()/25.
                tCsN[0] = ((data)>25).sum()
                tCsAve25[0] = 0.
                if tCsN[0] > 0:
                    tCsAve25[0] = data[(data)>25].mean()
            except:
                print "no cspad object"
        
            if False:
                rfAf = rawFrame[yA:yA+regionSize, xA:xA+regionSize].flatten()
                rfBf = rawFrame[yB:yB+regionSize, xB:xB+regionSize].flatten()
            else:
                rfAf = frame[yA:yA+regionSize, xA:xA+regionSize].flatten()
                rfAfGain = fGval[yA:yA+regionSize, xA:xA+regionSize].flatten()
                rfBf = frame[yB:yB+regionSize, xB:xB+regionSize].flatten()
                rfBfGain = fGval[yB:yB+regionSize, xB:xB+regionSize].flatten()
            if stitch:
                tempRAstitch = rfAf*0.
                tempRBstitch = rfAf*0.
##                rGs = []
##                for i in range(3):
## smarter to loopicize this...
                rG0 = (rfAfGain==1)
                rG1 = (rfAfGain==2)
                rG2 = (rfAfGain==3)
                tempRAstitch[rG0] = rAp1[0][rG0]*rfAf[rG0]+rAp0[0][rG0]
                tempRAstitch[rG1] = rAp1[1][rG1]*rfAf[rG1]+rAp0[1][rG1]
                tempRAstitch[rG2] = rAp1[2][rG2]*rfAf[rG2]+rAp0[2][rG2]
                rG0 = (rfBfGain==1)
                rG1 = (rfBfGain==2)
                rG2 = (rfBfGain==3)
                tempRBstitch[rG0] = rBp1[0][rG0]*rfBf[rG0]+rBp0[0][rG0]
                tempRBstitch[rG1] = rBp1[1][rG1]*rfBf[rG1]+rBp0[1][rG1]
                tempRBstitch[rG2] = rBp1[2][rG2]*rfBf[rG2]+rBp0[2][rG2]
            for i in range(regionPixels):## figure out how to copy instead
                tRegionA[i] = rfAf[i]
                tRAGain[i] = rfAfGain[i]
                tRegionB[i] = rfBf[i]
                tRBGain[i] = rfBfGain[i]
                if stitch:
                    tRAstitch[i] = tempRAstitch[i] 
                    tRBstitch[i] = tempRBstitch[i] 

##      fullAve = fullSum/ROIrows/ROIcols
            if camera == 0:
                nGoodEvt += 1
    tTree.Fill()


    nEvt += 1
    if nEvt%100==0:
        ##if nEvt>0:
        print "event %d" %(nEvt)
##        print csAve
##        if nEvt>50:
##            break

tFile.Write()
tFile.Close()