import ROOT, argparse, math, sys, os, copy,time
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/mberger')
from functions import Divideh

def setMarker(vec,style,color,size):
    marker=ROOT.TPolyMarker()
    marker.SetPoint(0,vec.X(),vec.Y())
    marker.SetMarkerStyle(style)
    marker.SetMarkerColor(color)
    marker.SetMarkerSize(size) 
    return marker

def setLine(start,stop,color=1,size=1):
    line=ROOT.TLine(start.X(),
                    start.Y(),
                    stop.X(),
                    stop.Y())
    line.SetLineColor(color)
    line.SetLineWidth(size)
    return line

def calcPOCA(thisrefpos,mom,pos):
    tmp1=ROOT.TVector3(mom)
    tmp1.SetMag(1)
    tmp2=ROOT.TVector3(pos-thisrefpos)
    fac=tmp2*mom
    fac/=mom.Mag()
    tmp1*= fac
    poca=ROOT.TVector3(thisrefpos+tmp1)
    return poca

def recalcPos():
    global clpos_new
    global weights
    global dirs
    global pocas

    weights=[]
    dirs=[]
    pocas=[]
     #get maxdigiamp:
    maxdigiamp=0
    for idigi in  range(len(digipos)):
        maxdigiamp=max(maxdigiamp,digiamp[idigi])

    ampweighting=clamp
    #if args.calcmod==2:
    #    ampweighting=maxdigiamp
    
    for idigi in range(len(digipos)):
        #dirToPad=digipos[idigi]-startpos
        pocas.append(calcPOCA(startpos,refmom,digipos[idigi]))
        dirToPad=digipos[idigi]-pocas[-1]
        dirs.append(dirToPad)

        if dirToPad.Mag()==0:
            weight=0
            print 'mag=0, setting weight to 0'
        else:
            weight=padplane.GetPad(hitpads[idigi]).GetWeight(digiamp[idigi]/ampweighting,
                                                             dirToPad.X(),
                                                             dirToPad.Y())
        tmp=dirToPad
        tmp*=weight
        weights.append(weight)

        print 'weight={0}, ampfrac={1}, amp={4}, dir=({2},{3})'.format(weight,digiamp[idigi]/ampweighting,dirToPad.X(),dirToPad.Y(),digiamp[idigi])
        if args.calcmod==0:
            digipos_new.append(ROOT.TVector3(tmp+startpos))
        if args.calcmod==1:
            digipos_new.append(ROOT.TVector3(tmp)) 
        if args.calcmod==2:
            digipos_new.append(digipos[idigi])

    #calculate new cluster position
    if args.calcmod==1:
        clpos_new=ROOT.TVector3(clpos_unw)
    sumAmp=0
    for idigi in range(len(digipos_new)):

        if args.calcmod==0:
            tmp=ROOT.TVector3(digipos_new[idigi])
            tmp*=digiamp[idigi]
            sumAmp+=digiamp[idigi]
            
        if args.calcmod==1:
            tmp=ROOT.TVector3(digipos_new[idigi])
            sumAmp+=digiamp[idigi]

        if args.calcmod==2:
            tmp=ROOT.TVector3(digipos_new[idigi])
            if weights[idigi]==0:
                weights[idigi]=0
            else:
                weights[idigi]=1./weights[idigi]
            tmp*=digiamp[idigi]*weights[idigi]
            sumAmp+=digiamp[idigi]*weights[idigi]
            print 'weighted Amp={0}, Amp={1}, weight{2}'.format(digiamp[idigi]*weights[idigi],digiamp[idigi],weights[idigi])
            #tmp*=weights[idigi]
            #sumAmp+=weights[idigi]

        clpos_new+=tmp

    if args.calcmod==0:
        clpos_new*=1./sumAmp
    if args.calcmod==2:
        clpos_new*=1./sumAmp
        clpos_new=clpos_new-startpos
        clpos_new*=0.5
        clpos_new+=startpos

def calc_cov2D(positions,amps,cog):
    cov=ROOT.TMatrixD(2,2)
    amp=0
    #print 'calc_cov2D'
    #cov.Print()
    for idig in range(len(positions)):
        c=ROOT.TMatrixD(2,1)
        c[0][0]=positions[idig].X()-cog.X()
        c[1][0]=positions[idig].Y()-cog.Y()
        c_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,c)
        acov=ROOT.TMatrixD(c,ROOT.TMatrixD.kMult,c_t)
        #print 'acov:',amps[idig]
        #acov.Print()
        acov*=amps[idig]
        amp+=amps[idig]
        cov+=acov

    cov*=1./float(amp)
    #print 'cov:',amp
    #cov.Print()
    return cov

def draw_cov2D(cov,cog,col):
    tcov=ROOT.TMatrixDSym(2,cov.GetMatrixArray())
    eigenproblem=ROOT.TMatrixDEigen(ROOT.TMatrixD(tcov))
    eigenwert=eigenproblem.GetEigenValues()
    eigenvector=eigenproblem.GetEigenVectors()
    teigenvector=ROOT.TVector2(eigenvector[0][0],eigenvector[1][0])
    teigenvector2=ROOT.TVector2(eigenvector[0][1],eigenvector[1][1])
    eigenvector_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,eigenvector)
    #print "draw cov2D"
    #teigenvector.Print()
    #teigenvector2.Print()
    #print eigenwert[0][0],eigenwert[1][1]
    #cog.Print()
    if eigenwert[0][0]!=0 and eigenwert[1][1]!=0:
        pcov=ROOT.TEllipse(cog.X(),cog.Y(),math.sqrt(eigenwert[0][0]),math.sqrt(eigenwert[1][1]),0,360,ROOT.TMath.RadToDeg()*teigenvector.Phi())
    else:
        return False
    pcov.SetFillStyle(4000)
    pcov.SetLineColor(col)
    return pcov


def draw_all():
    canv.cd()
    legend=ROOT.TLegend(0.1,0.7,0.4,0.9)
    legend.SetFillColor(0)
    plane.Draw()
    plane.GetXaxis().SetRangeUser(clmeanpos.X()-1,clmeanpos.X()+1)
    plane.GetYaxis().SetRangeUser(clmeanpos.Y()-1,clmeanpos.Y()+1)

    print 'residual_new={0}, residual={1}'.format((clpos_new-ppos).Mag(),(clpos-ppos).Mag())
    daval=[]
    global daval
    counter=-1

    for aval in avalanches:
        counter+=1
        apos=ROOT.TVector3(aval.x(),aval.y(),aval.t())
        daval.append(setMarker(apos,20,3,1*mcweights[counter]))
        daval[-1].Draw()

    if len(daval)>0:
        legend.AddEntry(daval[0],'Avalanche','p')
    ddelec=[]
    global ddelec
    for elec in driftedele:
        edpos=ROOT.TVector3(elec.x(),elec.y(),elec.t())
        ddelec.append(setMarker(edpos,7,2,5))
        ddelec[-1].Draw()

    if len(ddelec)>0:
        legend.AddEntry(ddelec[0],'Drifted Electron','P')

    ddigi_new=[]
    global ddigi_new
    ddigi_line=[]
    global ddigi_line
    counter=-1
    for d in digipos_new:
        counter+=1
        ddigi_new.append(setMarker(d,28,2,1))
        ddigi_new[-1].Draw()
        ddigi_line.append(setLine(digipos[counter],d))
        ddigi_line[-1].Draw()
        
    if len(ddigi_new)>0:
        legend.AddEntry(ddigi_new[0],'New Digipos','P')

    #for i in range(len(DigiMapper)-1):
        #padplane.GetPad(i).Draw()
    #momline=setLine(refpoint-refmom,refpoint+refmom)
    momline=setLine(startpos-refmom,startpos+refmom)
    momline.Draw()

    pocalines=[]
    counter=-1
    print 'drawing pocas'
    if len(pocas)>0 and len(digipos_new)>0:
        for poc in pocas:
            counter+=1
            print 'poc'
            pocalines.append(setLine(poc,digipos_new[counter]))
            pocalines[-1].Draw()

    hcounter=-1
    dtext=[]
    global dtext
    for hp in hitpads:
        hcounter+=1
        #print 'pad id:',hp
        padplane.GetPad(hp).Draw(2)
        dtext.append(ROOT.TText(digipos[hcounter].X(),
                                digipos[hcounter].Y(),
                                str(round(digiamp[hcounter],3))
                                )
                     )
        dtext[-1].SetTextSize(0.02)
        dtext[-1].Draw('same')

    dclpos_new=setMarker(clpos_new,28,1,1)
    global dclpos_new
    dclpos_new.Draw()
    legend.AddEntry(dclpos_new,'New Cluster Pos','p')

    dclpos_unw=setMarker(clpos_unw,24,1,0.5)
    global dclpos_unw
    dclpos_unw.Draw()
    legend.AddEntry(dclpos_unw,'Unweighted clpos','p')

    dclpos=setMarker(clpos,7,1,5)
    global dclpos
    dclpos.Draw()
    legend.AddEntry(dclpos,'Cluster Pos','p')

    dppos=setMarker(ppos,7,4,5)
    global dppos
    dppos.Draw()
    
    legend.Draw()
    canv.Update()
    u=raw_input("wait?")

parser=argparse.ArgumentParser(description='check pad response wit mc and reco files')
parser.add_argument('mcfile',help='the mc file',type=str)
parser.add_argument('digifile',help='the digi file',type=str)
parser.add_argument('recofile',help='the reco file',type=str)
parser.add_argument('--padf',help='the padplane file to use',type=str,default='tpc/parfiles/padPlane_FOPI.dat')
parser.add_argument('--pshape',help='pad shape file',type=str,default='tpc/parfiles/PadShapePrototype.dat')
parser.add_argument('--padres',help='padresponse file anme',type=str,default='tpc/FOPI/par/padresponse.txt')
parser.add_argument('--old',help='use old padresponse function (gauss int)',action='store_true')
parser.add_argument('--clpos',help='use clusterpos as start position for recalculation',action='store_true')
parser.add_argument('--maxpos',help='use max digi as start position for recalculation',action='store_true')
parser.add_argument('--evalRes',help='evaluate residual distribution for N events',type=int,default=-1)
parser.add_argument('--calcmod',help='the modus to calculate the new pos',type=int,default=0)
parser.add_argument('--startAt',help='go to event N to start',type=int,default=0)
parser.add_argument('--pfile',help='the file to save pictures to',type=str,default='./check_padresponse')
parser.add_argument('--iterative',help='recalculated new digipos until clusterpos is not moving anymore',action='store_true')
parser.add_argument('--sigCut',help='cut on error',type=float,default=[0,1000],nargs=2)
parser.add_argument('--eigenCut',help='cut on eigenval on plane',type=float,default=[0,1000],nargs=2)
parser.add_argument('--mcVerbose',help='plot more mc values',action='store_true')
args=parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") 
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

DigiMapper = []
file = open(args.padf, "r")
lines=file.readlines()
file.close()
lines.reverse()
for line in lines:
    koord = line.split(" ") 
    DigiMapper.append([float(koord[3]), float(koord[4])])

gem=ROOT.TpcGem(2000,0.02)
gem=ROOT.TpcGem(2000,0.02)
if args.old:
    padShapePool=ROOT.TpcPadShapePool(args.pshape,gem,0.1,0.005,0.01)
else:
    padShapePool=ROOT.TpcPadShapePool(args.pshape,args.padres)

#padShapePool=ROOT.TpcPadShapePool(args.pshape,gem,0.1,0.005,0.01)
padplane=ROOT.TpcPadPlane(args.padf,padShapePool)

infile=ROOT.TFile(args.recofile,"read")
tree=infile.Get("cbmsim")


tree.AddFriend("cbmsim",args.mcfile)
tree.AddFriend("cbmsim",args.digifile)

tree.SetBranchStatus("*",1)
#tree.SetBranchStatus("TpcCluster.*",1)
#tree.SetBranchStatus("TpcSPHit.*",1)
#tree.SetBranchStatus("TpcDigi.*",1)

plane=ROOT.TH2D("padplane","padplane",300,-15,15,300,-15,15)
plane.SetStats(0)
canv=ROOT.TCanvas('pads','pads',1000,1000) 

if args.evalRes!=-1:
    rescanv=ROOT.TCanvas("res",'res',1000,500)
    rescanv.Divide(2,1)
    
residual=ROOT.TH1D('residual','residual',400,-1,1)
residual.SetLineColor(2)
residual_new=ROOT.TH1D('residual_new','residual new',400,-1,1)
residual_new.SetLineColor(1)

digiresidual=ROOT.TH1D('dresidual','digi residual',400,-1,1)
digiresidual.SetLineColor(2)
digiresidual_new=ROOT.TH1D('dresidual_new','digi residual new',400,-1,1)
digiresidual_new.SetLineColor(1)

digiresidual2D=ROOT.TH2D('dres2D','digi residual 2D',73,0,73,400,-1,1)
digiresidual2D_new=ROOT.TH2D('dres2D_new','digi residual 2D new',73,0,73,400,-1,1)
canv2D=ROOT.TCanvas('2D','2D',1000,750)
canv2D.Divide(1,2)

canv_misc=ROOT.TCanvas('cmisc','misc',1000,1000)
hnumIter=ROOT.TH1D('hnumIter','NUmber of iterations',200,0,200)

global clpos_new
global weights 
global ddigi_new
global dppos
global dclpos
global dclpos_unw
global ddelec
global daval
global pocas
pocas=[]

evcount=-1
for e in tree:
    evcount+=1
    clmeanpos=ROOT.TVector3(0,0,0)
    if args.evalRes!=-1 and args.evalRes<evcount:
        break
    if evcount<args.startAt:
        continue

    for cl in e.TpcCluster:
        clmeanpos+=cl.pos()

    if e.TpcCluster.GetEntries()>0:
        clmeanpos*=1./e.TpcCluster.GetEntries()
        
    trcount=-1
    for tr in e.CutCosmics:
        trcount+=1
        for icl in range(tr.nhits()):
            print 'Event: {1}, track: {2}, cluster: {0}'.format(icl,evcount,trcount)
            refpoint=tr.GetProjectionPoint(icl)
            refmom=tr.GetProjectionPointsMom(icl)
            sphit=e.TpcSPHit.At(icl)
            cl=e.TpcCluster.At(sphit.getClusterID())
            clpos=ROOT.TVector3(cl.pos())
            clmeanpos=ROOT.TVector3(clpos)
            cl_sigA=cl.sig(3)
            if cl_sigA.Y()<args.sigCut[0] or cl_sigA.Y()>args.sigCut[1]:
                print 'SKIPPING',cl_sigA.Y()
                continue

            if cl.eigenValOnPlane().Z()<args.eigenCut[0] or cl.eigenValOnPlane().Z()>args.eigenCut[1]:
                print 'eigenval:', cl.eigenValOnPlane().Z(),args.eigenCut[0],args.eigenCut[1], cl.eigenValOnPlane().Z()<args.eigenCut[0], cl.eigenValOnPlane().Z()>args.eigenCut[1]
                continue

            if cl.isGeom().Mag()>1 or not cl.fitGood():
                print 'geomErr:',cl.hasGeomErr(), 'fitgood:',cl.fitGood()
                #continue

            clamp=cl.amp()
            ppos=ROOT.TVector3(tr.GetMCRefPos(icl))

            digimap=cl.getDigiMap()
            hitpads=[]
            mcids=[]
            mcweights=[]
            avalanches=[]
            driftedele=[]
            digipos=[]
            digiamp=[]
            refpos=[]

            maxdigipos=ROOT.TVector3(0,0,0)
            maxdigiamp=0

            for dig in digimap:
                digi=e.TpcDigi.At(dig.first)
                hitpads.append(digi.padId())
                digipos.append(ROOT.TVector3(DigiMapper[hitpads[-1]][0],DigiMapper[hitpads[-1]][1],clpos.Z()))
                digiamp.append(digi.amp())
                mcidcol=digi.mcId()
                if digiamp[-1]>maxdigiamp:
                    maxdigiamp=digiamp[-1]
                    maxdigipos=digipos[-1]

                if args.mcVerbose:
                    for imcid in range(mcidcol.nIDs()):
                        mcids.append(mcidcol.ID(imcid))
                        mcweights.append(mcids[-1].weight())
                        if mcids[-1].index()>=e.TpcDriftedElectron.GetEntriesFast():
                            continue
    
                        driftedele.append(e.TpcDriftedElectron.At(mcids[-1].index()))
                        avalanches.append(e.TpcAvalanche.At(mcids[-1].index()))


            #calculate unweighted cluster position
            clpos_unw=ROOT.TVector3(0,0,0)
            clpos_old=ROOT.TVector3(0,0,0)
            idig=-1
            for idpos in digipos:
                idig+=1
                clpos_unw+=idpos
                wpos=ROOT.TVector3(idpos)
                wpos*=digiamp[idig]
                clpos_old+=wpos
            #print 'length',len(digipos)
            clpos_unw*=1./len(digipos)
            clpos_old*=1./clamp



            #calculate reweighted clusterposition
            #first calculated new positions for the amplitudes
            digipos_new=[]
            

            clpos_new=ROOT.TVector3(0,0,0)

            startpos=clpos_unw
            if args.clpos:
                startpos=clpos
            if args.maxpos:
                startpos=maxdigipos

            startpos_old=ROOT.TVector3(clpos_new)
            lastchange=0
            change=1
            iterCount=-1
            clpos_new=startpos
            if args.evalRes==-1:
                draw_all()
                u=raw_input("start steping?")
            while (abs(change)>0.001):
            #for i in range(3):
                iterCount+=1
                lastchange=change
                digipos_new=[]
                startpos_old=ROOT.TVector3(startpos)
                clpos_new=ROOT.TVector3(0,0,0)
                recalcPos()
                startpos=ROOT.TVector3(clpos_new)
                if not args.iterative:
                    change=0
                else:
                    change=(startpos_old-startpos).Mag()
                    print 'change=',change
                if args.evalRes==-1:
                    draw_all()
                    u='l'
                    if args.iterative:
                        u=raw_input("next step?")
                        if u=='q':
                            exit()
                    time.sleep(0.1)

            hnumIter.Fill(iterCount)
            diff=clpos-ppos
            #residual.Fill(diff.Perp()*diff.X()/abs(diff.X()))
            residual.Fill(diff.X())
            diff=clpos_new-ppos
            #residual_new.Fill(diff.Perp()*diff.X()/abs(diff.X()))
            residual_new.Fill(diff.X())

            for idp in digipos:
                diff=idp-ppos
                digiresidual.Fill(diff.X())
                digiresidual2D.Fill(idp.Z(),diff.X())
            for idp in digipos_new:
                diff=idp-ppos
                digiresidual_new.Fill(diff.X())
                digiresidual2D_new.Fill(idp.Z(),diff.X())

            eq_digiamp=[]
            digiamp_new=[]
            i=-1
            for damp in digiamp:
                i+=1
                eq_digiamp.append(1)
                if args.calcmod==2:
                    digiamp_new.append(damp*weights[i])
                else:
                    digiamp_new.append(damp)

            orig_cov=calc_cov2D(digipos,digiamp,clpos_old)
            orig_cov_noAmp=calc_cov2D(digipos,eq_digiamp,clpos_old)
            new_cov=calc_cov2D(digipos_new,digiamp_new,clpos_new)
            new_cov_unwMeam=calc_cov2D(digipos_new,digiamp_new,clpos_unw)

            if args.evalRes==-1:
                draw_all()
                canv.cd()
                pcov1=draw_cov2D(orig_cov,clpos_old,3)
                pcov2=draw_cov2D(new_cov,clpos_new,4)
                pcov3=draw_cov2D(orig_cov_noAmp,clpos_old,8)
                #pcov4=draw_cov2D(new_cov_unwMeam,clpos_new,6)
                if pcov1:
                    pcov1.Draw("same")
                if pcov2:
                    pcov2.Draw("same")
                if pcov3:
                    pcov3.Draw('same')
                #if pcov4:
                #    pcov4.Draw('same')
                canv.Update()

                '''
                ddpos=[]
                for dd in digipos:
                    ddpos.append(setMarker(dd,20,7,1))
                    ddpos[-1].Draw()
                    '''



                u='w'
                u=raw_input("next cluster?")
                canv.Clear()
                if u=='q':
                    exit()
        print "next track",evcount

if args.evalRes!=-1:
    rescanv.cd(1)
    residual_new.Draw()
    residual.Draw("same")
    rescanv.cd(2)
    digiresidual_new.Draw()
    digiresidual.Draw("same")

    canv2D.cd(1)
    digiresidual2D.Draw('colz')
    canv2D.cd(2)
    digiresidual2D_new.Draw('colz')    

    canv_misc.cd()
    hnumIter.Draw()

    rescanv.SaveAs(args.pfile+".pdf(")
    canv2D.SaveAs(args.pfile+".pdf")
    canv_misc.SaveAs(args.pfile+".pdf)")

    outf=ROOT.TFile(args.pfile+".root",'recreate')
    rescanv.Write()
    canv2D.Write()
    canv_misc.Write()
    residual_new.Write()
    residual.Write()
    digiresidual_new.Write()
    digiresidual.Write()
    digiresidual2D_new.Write()
    digiresidual2D.Write()
    hnumIter.Write()

    outf.Close()

    rescanv.Update()
    canv2D.Update()
    '''
    rfile=ROOT.TFile("./padresponse.root",'recreate')
    rescanv.Write()
    canv2D.Write()
    digiresidual_new.Write()
    digiresidual.Write()
    residual_new.Write()
    residual.Write()    
    digiresidual2D.Write()
    digiresidual2D_new.Write()
    rfile.Close()
    '''
    u=raw_input("done?")