import ROOT, argparse

class clpoints:
    def __init__(self):
        self.xy=ROOT.TPolyMarker()
        self.xz=ROOT.TPolyMarker()
        self.yz=ROOT.TPolyMarker()
        self.allp=[self.xy,self.xz,self.yz]
        self.nCluster=0
    def addPoint(self,x,y,z):
        self.xy.SetNextPoint(x,y)
        self.xz.SetNextPoint(z,x)
        self.yz.SetNextPoint(z,y)
        self.nCluster+=1
    def setStyle(self,marker,size,color):
        for p in self.allp:
            p.SetMarkerStyle(marker)
            p.SetMarkerSize(size)
            p.SetMarkerColor(color)

parser=argparse.ArgumentParser(description='plot digis inside steps and steps')
parser.add_argument('recofile',help='the file with the reclustering output',type=str)
parser.add_argument("--batch",help="batch mode",action="store_true")
parser.add_argument("--nEvts",help="number of events to process",type=int,default=-1)
parser.add_argument("--gif",help="produce animated gif file",action="store_true")
args=parser.parse_args()

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

infile=open(args.recofile,"r")

events=[]
evstart=[]
#reading=True
#while reading:
founddigi=False
for line in infile:
    if line.find("EVENT:")>-1:
        events.append([])
        evstart.append([])
    if line.find("First Plane:")>-1:
        events[-1].append([])
        evstart[-1].append({})
        words=line.split(":")
        evstart[-1][-1]["first c"]=[float(words[1]),float(words[2]),float(words[3])]
    if line.find("First Direction")>-1:
        words=line.split(":")
        evstart[-1][-1]["first n"]=ROOT.TVector3(float(words[1]),float(words[2]),float(words[3]))
    if line.find("ew Plane:")>-1:
        words=line.split(":")
        evstart[-1][-1]["new c"]=[float(words[1]),float(words[2]),float(words[3])]
    if line.find("ew Direction")>-1:
        words=line.split(":")
        evstart[-1][-1]["new n"]=ROOT.TVector3(float(words[1]),float(words[2]),float(words[3]))

    if line.find("Checking if digis inside step:")>-1:
        events[-1][-1].append({})
        events[-1][-1][-1]["digis"]=[]
    if line.find("here center:")>-1:
        words=line.split(":")
        events[-1][-1][-1]["here c"]=[float(words[1]),float(words[2]),float(words[3])]
    if line.find("here normal:")>-1:
        words=line.split(":")
        events[-1][-1][-1]["here n"]=ROOT.TVector3(float(words[1]),float(words[2]),float(words[3]))
    if line.find("next center:")>-1:
        words=line.split(":")
        events[-1][-1][-1]["next c"]=[float(words[1]),float(words[2]),float(words[3])]
    if line.find("next normal:")>-1:
        words=line.split(":")
        events[-1][-1][-1]["next n"]=ROOT.TVector3(float(words[1]),float(words[2]),float(words[3]))


    if line.find("INSIDE")>-1:
#        events[-1][-1]["digis"].append([])
        founddigi=True
        words=line.split(":")
        events[-1][-1][-1]["digis"].append([float(words[1]),float(words[2]),float(words[3])])
#        events[-1][-1]["digis"]=[float(words[1]),float(words[2]),float(words[3])]
        if line.find("NOT")>-1:
            events[-1][-1][-1]["digis"][-1].append(-1)
        else:
            events[-1][-1][-1]["digis"][-1].append(1)

    if founddigi:
        if line.find("BehindHere")>-1:
            words=line.split(":")
            events[-1][-1][-1]["digis"][-1].append(words[1])
        if line.find("BehindNext")>-1:
            founddigi=False
            words=line.split(":")
            events[-1][-1][-1]["digis"][-1].append(words[1])
    if line.find("found some digis, pushing back a cluster")>-1:
        events[-1][-1][-1]["pushcl"]=1
    if line.find("Cluster was to small")>-1:
        words=line.split(":")
        events[-1][-1][-1]["smallsize"]=words[1]
        events[-1][-1][-1]["minsize"]=words[2]

c1=ROOT.TCanvas("bla","bla",500,900) 
c1.Divide(1,2)
hxy=ROOT.TH2D("hxy","XY-Plane",30,-15,15,30,-15,15)
hyz=ROOT.TH2D("hyz","YZ-Plane",75,0,75,30,-15,15)
hxz=ROOT.TH2D("hxz","XZ-Plane",75,0,75,30,-15,15)
hxy.SetStats(0)
hyz.SetStats(0)
hxz.SetStats(0)
hists=[hxy,hxz,hyz]
ev=-1

goto=raw_input("full auto (0) or go to event (evnr)?")

for e in events:
    ev+=1
    print "Event:",ev
    if args.nEvts!=-1 and ev>args.nEvts:
        break
    if int(goto)!=0 and ev!=int(goto):
        continue
    startdirs=[]
    startplanes=[]
    tra=-1
    for tr in e:
        tra+=1
        if evstart[ev][tra].get("first n",0)!=0:
            startdir=evstart[ev][tra]["first n"]
            startdirs.append(ROOT.TLine(evstart[ev][tra]["first c"][0],
                                        evstart[ev][tra]["first c"][1],
                                        evstart[ev][tra]["first c"][0]+startdir.X(),
                                        evstart[ev][tra]["first c"][1]+startdir.Y()))
            startdirs[-1].SetLineColor(3)
            startdirs.append(ROOT.TLine(evstart[ev][tra]["first c"][2],
                                        evstart[ev][tra]["first c"][1],
                                        evstart[ev][tra]["first c"][2]+startdir.Z(),
                                        evstart[ev][tra]["first c"][1]+startdir.Y()))
            startdirs[-1].SetLineColor(3)

            startdir=evstart[ev][tra]["new n"]
            startdirs.append(ROOT.TLine(evstart[ev][tra]["new c"][0],
                                        evstart[ev][tra]["new c"][1],
                                        evstart[ev][tra]["new c"][0]+startdir.X(),
                                        evstart[ev][tra]["new c"][1]+startdir.Y()))
            startdirs[-1].SetLineColor(4)
            startdirs.append(ROOT.TLine(evstart[ev][tra]["new c"][2],
                                        evstart[ev][tra]["new c"][1],
                                        evstart[ev][tra]["new c"][2]+startdir.Z(),
                                        evstart[ev][tra]["new c"][1]+startdir.Y()))
            startdirs[-1].SetLineColor(4)

            plane=evstart[ev][tra]["first n"].Orthogonal()
            startplanes.append(ROOT.TLine(
                    evstart[ev][tra]["first c"][0]-plane.X(),
                    evstart[ev][tra]["first c"][1]-plane.Y(),
                    evstart[ev][tra]["first c"][0]+plane.X(),
                    evstart[ev][tra]["first c"][1]+plane.Y()))
            startplanes[-1].SetLineColor(3)
            startplanes.append(ROOT.TLine(
                    evstart[ev][tra]["first c"][2]-plane.Z(),
                    evstart[ev][tra]["first c"][1]-plane.Y(),
                    evstart[ev][tra]["first c"][2]+plane.Z(),
                    evstart[ev][tra]["first c"][1]+plane.Y()))
            startplanes[-1].SetLineColor(3)

            plane=evstart[ev][tra]["new n"]#.Orthogonal()
#            direct=evstart[ev][tra]["new n"]
            plane.RotateX(90)
            startplanes.append(ROOT.TLine(
                    evstart[ev][tra]["new c"][0]-plane.X(),
                    evstart[ev][tra]["new c"][1]-plane.Y(),
                    evstart[ev][tra]["new c"][0]+plane.X(),
                    evstart[ev][tra]["new c"][1]+plane.Y()))
            startplanes[-1].SetLineColor(4)
            startplanes.append(ROOT.TLine(
                    evstart[ev][tra]["new c"][2]-plane.X(),  #direct.X()/direct.Z(),#-plane.Z(),
                    evstart[ev][tra]["new c"][1]-plane.Y(),
                    evstart[ev][tra]["new c"][2]+plane.X(),# -direct.X()/direct.Z(),#plane.Z(),
                    evstart[ev][tra]["new c"][1]+plane.Y()))
            startplanes[-1].SetLineColor(4)

        alldigis=clpoints()
        alldigis.setStyle(8,0.5,1)
        useddigis=clpoints()
        useddigis.setStyle(8,0.5,2)
        planes=[]
        dirs=[]
        for step in tr:
            print "cluster was pushed:",step.get("pushcl",0)
            print "clustersize was:",step.get("smallsize","ok")
            print "minimum cluster size was:",step.get("minsize","ok")
            extradir=step["here n"]
            dirs.append(ROOT.TLine(step["here c"][0],step["here c"][1],
                                   step["here c"][0]+extradir.X(),step["here c"][1]+extradir.Y()))
            dirs.append(ROOT.TLine(step["here c"][2],step["here c"][1],
                                   step["here c"][2]+extradir.Z(),step["here c"][1]+extradir.Y()))

            plane=step["here n"].Orthogonal()
            planes.append(ROOT.TLine(
                    step["here c"][0]-plane.X(),
                    step["here c"][1]-plane.Y(),
                    step["here c"][0]+plane.X(),
                    step["here c"][1]+plane.Y()))
            planes.append(ROOT.TLine(
                    step["here c"][2]-plane.Z(),
                    step["here c"][1]-plane.Y(),
                    step["here c"][2]+plane.Z(),
                    step["here c"][1]+plane.Y()))

            plane=step["next n"].Orthogonal()
            planes.append(ROOT.TLine(
                    step["next c"][0]-plane.X(),
                    step["next c"][1]-plane.Y(),
                    step["next c"][0]+plane.X(),
                    step["next c"][1]+plane.Y()))
            planes[-1].SetLineColor(2)
            planes.append(ROOT.TLine(
                    step["next c"][2]-plane.Z(),
                    step["next c"][1]-plane.Y(),
                    step["next c"][2]+plane.Z(),
                    step["next c"][1]+plane.Y()))
            planes[-1].SetLineColor(2)

            for d in step["digis"]:
                alldigis.addPoint(d[0],d[1],d[2])
                if d[3]==1:
                    useddigis.addPoint(d[0],d[1],d[2])
#                    print "IN:",d[4],d[5]
#                else:
#                    print "OUT:",d[4],d[5]

            c1.cd(1)
            hxy.Draw()
            startplanes[-4].Draw("same")
            startdirs[-4].Draw("same")
            startplanes[-2].Draw("same")
            startdirs[-2].Draw("same")

            alldigis.xy.Draw("same")
            useddigis.xy.Draw("same")
            planes[-4].Draw("same")
            planes[-2].Draw("same")
            dirs[-2].Draw("same")
            c1.cd(2)
            hyz.Draw()
            startplanes[-3].Draw("same")
            startdirs[-3].Draw("same")
            startplanes[-1].Draw("same")
            startdirs[-1].Draw("same")

            alldigis.yz.Draw("same")
            useddigis.yz.Draw("same")
            planes[-3].Draw("same")
            planes[-1].Draw("same")
            dirs[-1].Draw("same")
#            c1.cd(3)
#            hxz.Draw()
            c1.Update()
            if args.gif:
                c1.Print("Reclustering.gif+20")
            if not args.batch:
                u=raw_input("bla")
                if u=="q":
                    exit()

        for h in hists:
            h.GetXaxis().UnZoom()
            h.GetYaxis().UnZoom()
    if int(goto)!=0:
        goto=raw_input("full auto (0) or go to event (evnr)?")