import ROOT, math, os, sys
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1')
sys.path.append(pandapath + '/macro/tpc/FOPI/mberger')
from functions import set_palette, openTree, thisIsTheEnd, getSlices
import random
import argparse

class clpoints:
    def __init__(self,name='Nothing'):
        self.xy=ROOT.TPolyMarker()
        self.xz=ROOT.TPolyMarker()
        self.yz=ROOT.TPolyMarker()
        self.allp=[self.xy,self.xz,self.yz]
        self.nCluster=0
        self.filled=0
        self.Name=name
        self.legP=ROOT.TPolyMarker()
        self.extrema=[[999,-999],[999,-999],[999,-999]]

    def addPoint(self,x,y,z):
        self.xy.SetNextPoint(x,y)
        #print "setting point ",x,y,z
        self.xz.SetNextPoint(z,x)
        self.yz.SetNextPoint(z,y)
        self.nCluster+=1
        #min and max X
        self.extrema[0][0]=min(self.extrema[0][0],x)
        self.extrema[0][1]=max(self.extrema[0][1],x)
        #min and max Y
        self.extrema[1][0]=min(self.extrema[1][0],y)
        self.extrema[1][1]=max(self.extrema[1][1],y)
        #min and max Z
        self.extrema[2][0]=min(self.extrema[2][0],z)
        self.extrema[2][1]=max(self.extrema[2][1],z)

        if self.filled==0:
            self.filled=1

    def setStyle(self,marker,size,color):
        for p in self.allp:
            p.SetMarkerStyle(marker)
            p.SetMarkerSize(size)
            p.SetMarkerColor(color)
        self.legP.SetMarkerStyle(marker)
        self.legP.SetMarkerColor(color)
        self.legP.SetMarkerSize(5)

    def isFilled(self):
        return self.filled

def print_da_shit():
    circ=ROOT.TEllipse(0,0,5,5,0,360,0)
    circo=ROOT.TEllipse(0,0,15,15,0,360,0)
    circ.SetFillStyle(4000)
    circo.SetFillStyle(4000)

    c1.cd(1)
    hxy.Draw()

    circo.Draw()
    circ.Draw("SAME")

    c1.cd(2)
    hyz.Draw()

    c1.cd(3)
    hxz.Draw()

    mima=[[999,-999],[999,-999],[999,-999]]

    opt="same"
    for p in allPoints:
        if p.isFilled:
            print "plotting",p.Name
            p.xy.Print()
            c1.cd(1)
            p.xy.Draw(opt)
            #p.xy.Dump()
            c1.cd(2)
            p.yz.Draw(opt)
            c1.cd(3)
            p.xz.Draw(opt)
            opt="same"
            for i in range(3):
                    mima[i][0]=min(mima[i][0],p.extrema[i][0])
                    mima[i][1]=max(mima[i][1],p.extrema[i][1])

    for i in range(3):
        if mima[i][0]<0:
            mima[i][0]*=1.01
        else:
            mima[i][0]*=0.99

        if mima[i][1]<0:
            mima[i][1]*=0.99
        else:
            mima[i][1]*=1.01

    if args.zoom:
        hxy.GetXaxis().SetRangeUser(mima[0][0],mima[0][1])
        hxy.GetYaxis().SetRangeUser(mima[1][0],mima[1][1])

        hyz.GetXaxis().SetRangeUser(mima[2][0],mima[2][1])
        hyz.GetYaxis().SetRangeUser(mima[1][0],mima[1][1])

        hxz.GetXaxis().SetRangeUser(mima[2][0],mima[2][1])
        hxz.GetYaxis().SetRangeUser(mima[0][0],mima[0][1])

    c1.Update()

    if args.c2:
        c2.cd()
        hxy.Draw()
        circ.Draw("SAME")
        circo.Draw("SAME")
        for p in allPoints:
            if p.isFilled:
                p.xy.Draw('same')
        c2.Update()

    #print cluster.nCluster,"cluster and ",recluster.nCluster," recluster and",McResidualsRef.nCluster," refpoints"
   
    print "at event:",ev
    if args.pfile!="None":
        c1.SaveAs(args.pfile+'[')

    u=raw_input("next event?")

    if u=="s" and args.pfile!="None":
        c1.SaveAs(args.pfile)
    if u=="q":
        if args.pfile!="None":
            c1.SaveAs(args.pfile+']')
        exit()

    for h in hists:
        h.GetXaxis().UnZoom()
        h.GetYaxis().UnZoom()

def calc_digilen(phi=-90):
    global xsize
    global ysize
    global zsize
    global lsize
    global p1size
    global p2size

    #print 'd:',dx,dy,dz
    xsize[0]=max(xsize[0],dx)
    ysize[0]=max(ysize[0],dy)
    zsize[0]=max(zsize[0],dz)
    xsize[1]=min(xsize[1],dx)
    ysize[1]=min(ysize[1],dy)
    zsize[1]=min(zsize[1],dz)
#    print 'sizes:',xsize[0],xsize[1],ysize[0],ysize[1],zsize[0],zsize[1]
    #print xsize[0],xsize[1]
    if phi!=-1:
        dPos=ROOT.TVector3(dx,dy,dz)
        dist=dPos-cl.pos()
        alpha=dist.Phi()-phi
        longi=math.cos(alpha)*dist.Perp()
        perp=math.sin(alpha)*dist.Perp()
        #print 'alpha=',alpha*ROOT.TMath.RadToDeg(),'dphi=',dist.Phi()*ROOT.TMath.RadToDeg(),'phi=',phi*ROOT.TMath.RadToDeg()
        #dist.Print()
        lsize[0]=max(lsize[0],longi)
        p1size[0]=max(p1size[0],perp)
        lsize[1]=min(lsize[1],longi)
        p1size[1]=min(p1size[1],perp)
#        print 'longi and perp', longi,perp
#        print lsize[0],lsize[1],p1size[0],p1size[1]
        #print p1size[0],p1size[1]
    return

parser=argparse.ArgumentParser(description='plot cluster before and after reclustering')
parser.add_argument('recofile',help='the recofile',type=str)
parser.add_argument('--runs', help='if data, the runs to use (%(default)s)', type=str, default='3888')
parser.add_argument('--pattern', help='if data, the file name patters ('')', type=str, default='')
parser.add_argument('--ev',help='go to event',type=int,default=-1)

parser.add_argument('--terr',help='plot only if nof tracks and nof retracks is different',action='store_true')
parser.add_argument('--clus',help='Plot standart cluster',action='store_true')
parser.add_argument('--recl',help='Plot the reclusters',action="store_true")
parser.add_argument('--ref',help='plot reference points',action='store_true')
parser.add_argument('--mc',help='also plot mc points',action='store_true')
parser.add_argument('--mcref',help='Draw MC reference points for residuals',action='store_true')
parser.add_argument('--digi',help='plot also the digisfrom cluster',action='store_true')
parser.add_argument('--treedigi',help='plot digis from digitree',action='store_true')
parser.add_argument('--treesample',help='plot samples from sample tree',action='store_true')
parser.add_argument('--treesignal',help='plot signals from signal tree',action='store_true')
parser.add_argument('--genfit',help='plot the genfit track',action='store_true')

parser.add_argument('--track',help='plot only cluster from track',action="store_true")
parser.add_argument('--riemann',help='use riemann tracks',action="store_true")

parser.add_argument('--c2',help='plot hxy in a single canvas',action="store_true")
parser.add_argument('--cpc',help='to long tracks',action="store_true")
parser.add_argument('--pre',help='use pre fits instead of fits',action='store_true')

parser.add_argument('--mcfile',help='the file with the mc points',type=str,default="")
parser.add_argument('--digifile',help='the file with the digis',type=str,default="")
parser.add_argument('--padf',help='the padplane file to use',type=str,default='tpc/parfiles/padPlane_FOPI.dat')

parser.add_argument('--ft0',help='ft for digi z calc',type=float,default=0)
parser.add_argument('--ftbin',help='ftbin for digi z calc',type=float,default=64.3087)
parser.add_argument('--vd',help='drift velocity for digi z calc',type=float,default=0.00237708)
parser.add_argument('--fzGem',help='z gem for digi z calc',type=float,default=0)

parser.add_argument('--zCut',help='only plot events wich are inside cut',type=int,default=[-1,-1],nargs=2)
parser.add_argument('--thetaCut',help='a cut on the theta angles of the track',type=float,default=[-1,-1],nargs=2)
parser.add_argument('--resCut',help='cut on the mean residual of a track',type=float,default=[0,0],nargs=2)
parser.add_argument('--ClMeanAmpCut',help='cut on the cluster mean amp',type=float,default=[-1,-1],nargs=2)
parser.add_argument('--llencut',help='cut on llen',type=float,default=[-1,-1],nargs=2)
parser.add_argument('--specialcut',help='a special cut on clustersize and residual',action='store_true')
parser.add_argument('--zoom',help='automaticly zoom into track',action='store_true')
parser.add_argument('--CutLeng',help='Cut Length',action='store_true')
parser.add_argument('--clustwise',help='cluster wise plotting',action='store_true')
parser.add_argument('--pfile',help='the file to save the canvas to',type=str,default='None')

parser.add_argument('--hlp',help='show help',action='store_true')
args=parser.parse_args()

if args.hlp:
    parser.print_help()
    exit(0)

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


#Infile=ROOT.TFile(args.recofile,"read")
#tree=Infile.Get("cbmsim")
tree=openTree(args.recofile, args.runs, args.pattern)

if args.mc:
    if args.mcfile=="":
        mcfile=args.recofile.replace("reco","mc")
        mcfile=mcfile.replace("_recl","")
    else:
        mcfile=args.mcfile
    tree.AddFriend("cbmsim",mcfile)
if args.digi or args.treedigi or args.treesample or args.treesignal:
    if args.digifile=="":
        digifile=args.recofile.replace("reco","raw")
        digifile=digifile.replace("_recl","")
    else:
        digifile=args.digifile
    tree.AddFriend("cbmsim",digifile)

tree.SetBranchStatus("*",1)
'''
tree.SetBranchStatus("TpcDigi.*",1)
tree.SetBranchStatus("TpcCluster.*",1)
tree.SetBranchStatus("TpcPreCluster.*",1)
tree.SetBranchStatus("TrackPrePostFit.*",1)
tree.SetBranchStatus("TrackPostFit.*",1)
tree.SetBranchStatus("TpcPreSPHit.*",1)
tree.SetBranchStatus("TpcSPHit.*",1)
tree.SetBranchStatus("RiemannTrack.*",1)
tree.SetBranchStatus("RiemannHit.*",1)
tree.SetBranchStatus("PreRiemannTrack.*",1)
tree.SetBranchStatus("PreRiemannHit.*",1)
'''

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


if tree==None:
    print "no cbmsim"
    Infile.Close()
    exit()

c1=ROOT.TCanvas("bla","bla",500,1400)
c3=ROOT.TCanvas("legend","legend",500,500)
#c4=ROOT.TCanvas("infos","infos",500,500)

if args.c2:
    c2=ROOT.TCanvas()
c1.Divide(1,3)
hxy=ROOT.TH2D("hxy","XY-Plane",300,-15,15,300,-15,15)
hyz=ROOT.TH2D("hyz","YZ-Plane",750,0,75,300,-15,15)
hxz=ROOT.TH2D("hxz","XZ-Plane",750,0,75,300,-15,15)
hxy.SetStats(0)
hyz.SetStats(0)
hxz.SetStats(0)
hists=[hxy,hxz,hyz]

ev=-1
outer=ROOT.TArc(.35,0.5,0.5)
outer.SetLineColor(1)
outer.SetLineWidth(2)

circ=ROOT.TEllipse(0,0,5,5,0,360,0)
circo=ROOT.TEllipse(0,0,15,15,0,360,0)
addleg=True
for e in tree:
    ev+=1
    if args.ev!=-1 and args.ev!=ev:
        continue

    cluster=clpoints('Cluster')
    cluster.setStyle(8,0.5,1)
    recluster=clpoints('Recl Cluster')
    recluster.setStyle(8,0.5,2)

    McResidualsRef=clpoints('McResidualRef')
    McResidualsRef.setStyle(8,0.5,3)

    digipoints=clpoints('Digis')
    digipoints.setStyle(29,0.3,6)
    recldigipoints=clpoints('Recl Digis')
    recldigipoints.setStyle(29,0.3,7)
    trkdigipoints=clpoints('Track Digis')
    trkdigipoints.setStyle(29,0.3,11)

    mcpoint=clpoints('MC Points')
    mcpoint.setStyle(8,0.2,4)

    refpoint=clpoints('Ref Points')
    refpoint.setStyle(22,0.5,46)

    trackpoint=clpoints('Genfit Track')
    trackpoint.setStyle(20,0.1,28)
    rtrackpoint=clpoints('Riemann Cluster')
    rtrackpoint.setStyle(28,0.5,34)

    samplepoint=clpoints('Sample Points')
    samplepoint.setStyle(2,0.2,95)

    signalpoint=clpoints('Signal Point')
    signalpoint.setStyle(4,0.2,81)

    allPoints=[]
    allPoints.append(cluster)
    allPoints.append(recluster)
    allPoints.append(McResidualsRef)
    allPoints.append(digipoints)
    allPoints.append(recldigipoints)
    allPoints.append(trkdigipoints)
    allPoints.append(mcpoint)
    allPoints.append(refpoint)
    allPoints.append(trackpoint)
    allPoints.append(rtrackpoint)
    allPoints.append(samplepoint)
    allPoints.append(signalpoint)
    
    if addleg:
        addleg=False
        leg=ROOT.TLegend(0,0,1,1)
        leg.SetFillColor(0)
        for p in allPoints:
            leg.AddEntry(p.legP,p.Name,'p')
        c3.cd()
        leg.Draw()
        c3.Update()

    if args.recl:
        ntrk=e.TrackPrePostFit.GetEntries()
        #ntrk=e.TrackPostFit.GetEntries()
    else:
        ntrk=-1

    nretrk=e.TrackPostFit.GetEntries()


    print ntrk," tracks, ",nretrk," re tracks"
    if (ntrk==nretrk and args.terr):
        continue

    lines=[]

    if args.pre and args.recl:
        dingens=e.TpcPreTrackPreFit
    elif args.recl:
        dingens=e.TrackPrePostFit
    else:
        dingens=e.TrackPostFit
        


    #get track information
    for trk in dingens:

        rep=trk.getCardinalRep()
        flag=int(rep.getStatusFlag())
        print "track fit flag=",flag
        p1=trk.getPos()
        mom=trk.getMom()
        theta=mom.Theta()
        mom.SetMag(10)
        p2=p1+mom
        lines.append(ROOT.TLine(p1.X(),p1.Y(),p2.X(),p2.Y()))
        
        #add cluster from track
        if args.track:
            cand=trk.getCand()
            detid=cand.getDetIDs()[0]
            hits=cand.getHitIDs(detid)
            for hit in hits:
                if args.recl:
                    sphit=e.TpcPreSPHit.At(hit)
                    cl=e.TpcPreCluster.At(sphit.getClusterID())
                else:
                    sphit=e.TpcSPHit.At(hit)
                    cl=e.TpcCluster.At(sphit.getClusterID())

                if args.clus:
                    cluster.addPoint(cl.pos().X(),cl.pos().Y(),cl.pos().Z())

    if args.thetaCut[0]!=-1 and args.thetaCut[0]!=args.thetaCut[1]:
        if theta<args.thetaCut[0] or theta>args.thetaCut[1]:
            continue

            
    #add reclusters from track 
    if args.track and args.recl:
        if args.pre:
            dingens=e.TpcTrackPreFit
        else:
            dingens=e.TrackPostFit

        for trk in dingens:
            cand=trk.getCand()
            detid=cand.getDetIDs()[0]
            hits=cand.getHitIDs(detid)
            for hit in hits:
                sphit=e.TpcSPHit.At(hit)
                cl=e.TpcCluster.At(sphit.getClusterID())
                recluster.addPoint(cl.pos().X(),cl.pos().Y(),cl.pos().Z())
                errorA=cl.sig(3)
                print 'recluster: amp={0}; error A1={1}; A2={2}; A3={3}'.format(cl.amp,errorA.X(),errorA.Y(),errorA.Z())

    #add residual MC ref points
    cut=0                
    trackbranch=e.TrackFitStat_0
    for tr in trackbranch:
        leng=tr.GetLength()
        #print "Length of the track:", leng
        if args.CutLeng:
            if leng > 9 and leng < 10.5:
                cut=1

    if args.mcref:
        for tr in trackbranch:
            resMcRefX=tr.GetMcRefPosX()  
            resMcRefY=tr.GetMcRefPosY() 
            resMcRefZ=tr.GetMcRefPosZ() 
            for i in range(len(resMcRefX)):
                McResidualsRef.addPoint(resMcRefX[i],resMcRefY[i],resMcRefZ[i])

    #add mc points
    if args.mc:
        mean_mcres=0
        num_res=0
        for stat in e.TrackFitStat_0:
            for res in stat.GetResMcZ():
                mean_mcres+=res
                num_res+=1
        if num_res!=0:
            mean_mcres/=num_res
        print "Mean MC Residual in Z:",mean_mcres
        for poi in e.TpcPoint:
            mcpoint.addPoint(poi.GetX(),poi.GetY(),poi.GetZ())

    #add cluster from riemann track (NOT WORKING!)
    if args.riemann:
        for rtrk in e.RiemannTrack:
            nrhits=rtrk.getNumHits()
            #print "number of riemann hits:",nrhits
            for i in range(nrhits):
                rhit=rtrk.getHit(i)
                #print "riemann hit",i,rhit
                rcl=rhit.cluster()
                #print "riemann cluster:",rcl.pos().X(),rcl.pos().Y(),rcl.pos().Z()
                rtrackpoint.addPoint(rcl.pos().X(),rcl.pos().Y(),rcl.pos().Z())

    # add genfit tracks
    if args.genfit:
        for gtrk in e.TrackFitStat_0:
            Gpoints=gtrk.GetGTrack()
            for p in Gpoints:
                trackpoint.addPoint(p.X(),p.Y(),p.Z())

    if args.ref:
        for cl in e.TrackFitStat_0:
            refReP=cl.GetProjectionPoints()
            for p in refReP:
                refpoint.addPoint(p.X(),p.Y(),p.Z())

    #add clusters and digis from Cluster and Recluster Array
    if not args.track:
        if args.recl:
            clbranch=e.TpcPreCluster
        else:
            clbranch=e.TpcCluster
        digi_c=0    
        #add clusters
        meanz=0
        for cl in clbranch:
            if args.clus:
                cluster.addPoint(cl.pos().X(),cl.pos().Y(),cl.pos().Z())
            meanz+=cl.pos().Z()
            #add digis from clusters
            xsize=[cl.pos().X(),cl.pos().X()]
            ysize=[cl.pos().Y(),cl.pos().Y()]
            zsize=[cl.pos().Z(),cl.pos().Z()]

            if args.digi and not args.treedigi:
                digimap=cl.getDigiMap()
                lsize=[0,99]#[cl.pos().Z(),cl.pos().Z()]
                p1size=[0,99]#[cl.pos().Z(),cl.pos().Z()]
                p2size=[0,99]#[cl.pos().Z(),cl.pos().Z()]
                for dig in digimap:
                    digi_c+=1
                    digi=e.TpcDigi.At(dig.first)
                    dx=float(DigiMapper[digi.padId()][0])
                    dy=float(DigiMapper[digi.padId()][1])
                    dz=(digi.t()*args.ftbin+args.ft0)*args.vd+args.fzGem
                    digipoints.addPoint(dx,dy,dz)
                    calc_digilen(-90*ROOT.TMath.DegToRad())

                xlen=xsize[0]-xsize[1]
                ylen=ysize[0]-ysize[1]
                zlen=zsize[0]-zsize[1]
                llen=lsize[0]-lsize[1]
                p1len=p1size[0]-p1size[1]
                p2len=zlen

                #print 'xlen:{0:4.2f} ylen:{1:4.2f} zlen:{2:4.2f} llen:{3:4.2f} p1len:{4:4.2f} p2len:{5:4.2f}'.format(xlen,ylen,zlen,llen,p1len,p2len)

            if args.llencut[0]!=-1 and ( llen>args.llencut[0] and llen<args.llencut[1]):
                if abs(llen-ylen)>0.1:
                    print args.llencut
                    print 'llencut'
                    print_da_shit()


            if args.clustwise:
                print cl.pos().X()
                print_da_shit()
#                u=raw_input('next?')
        if args.treedigi:
            for digi in e.TpcDigi:
                digi_c+=1
                dx=float(DigiMapper[digi.padId()][0])
                dy=float(DigiMapper[digi.padId()][1])
                dz=(digi.t()*args.ftbin+args.ft0)*args.vd+args.fzGem
                digipoints.addPoint(dx,dy,dz)

        if args.treesample:
            for sample in e.TpcSample:
                dx=float(DigiMapper[sample.padId()][0])
                dy=float(DigiMapper[sample.padId()][1])
                dz=(sample.t()*args.ftbin+args.ft0)*args.vd+args.fzGem
                samplepoint.addPoint(dx,dy,dz)
                print 'hit samples:',dx,dy,sample.padId()
                         
        if args.treesignal:
            for signal in e.TpcSignal:
                dx=float(DigiMapper[signal.padId()][0])
                dy=float(DigiMapper[signal.padId()][1])
                dz=(signal.t()+args.ft0)*args.vd+args.fzGem
                signalpoint.addPoint(dx,dy,dz)

        if clbranch.GetEntriesFast()>0:
            meanz/=clbranch.GetEntriesFast()
        if args.zCut[0]!=-1 and (meanz<args.zCut[0] or meanz>args.zCut[1]):
            continue

        #add reclusters
        if args.recl:
            print e.TpcCluster.GetEntries(),"recluster"
            for rcl in e.TpcCluster:
                if args.recl:
                    recluster.addPoint(rcl.pos().X(),rcl.pos().Y(),rcl.pos().Z())
                    errorA=rcl.sig(3)
                    errorXYZ=rcl.sig(2)
                    print 'recluster: amp={0: >11.4f}; error A1={1: >2.4f}; A2={2: >2.4f}; A3={3: >2.4f};pos X={4: >5.2f};Y={5: >5.2f};Z={6: >5.2f}'.format(rcl.amp(),errorA.X(),errorA.Y(),errorA.Z(),rcl.pos().X(),rcl.pos().Y(),rcl.pos().Z())
                    print 'recluster: amp={0: >11.4f}; error X={1: >2.4f}; Y={2: >2.4f}; Z={3: >2.4f};size={4}'.format(rcl.amp(),errorXYZ.X(),errorXYZ.Y(),errorXYZ.Z(),rcl.size())

                #add digis from reclusters
                digirecl_c=0
                if args.digi:
                    digimap=rcl.getDigiMap()
                    for dig in digimap:
                        digirecl_c+=1
                        digi=e.TpcDigi.At(dig.first)
                        dx=float(DigiMapper[digi.padId()][0])
                        dy=float(DigiMapper[digi.padId()][1])
                        dz=(digi.t()*args.ftbin+args.ft0)*args.vd+args.fzGem
                        recldigipoints.addPoint(dx,dy,dz)
                        
    if args.CutLeng:
        if cut==1:
            print_da_shit()
    else:
        print_da_shit() 
#    if args.ev==-1 or args.ev <= ev:
#        if args.c2:
#            c2.cd()
#            hxy.Draw()
#            cluster.xy.Draw("same")
#            recluster.xy.Draw("same")
#            McResidualsRef.xy.Draw("same")