import ROOT,math,argparse,sys

parser=argparse.ArgumentParser(description="Create recluster animation input")
parser.add_argument('recofile',help='a recofile with recluster info',type=str)
parser.add_argument('event',help='the event to print out',type=int)
parser.add_argument('outfolder',help='the folder to write the output to',type=str)
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=-62.5)
parser.add_argument('--hlp',help='print 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('gROOT->SetStyle("Plain")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")
#set_palette()

rfile=ROOT.TFile(args.recofile)
tree=rfile.Get("cbmsim")

digifile=args.recofile.replace("reco","raw")
digifile=digifile.replace("_recl","")
tree.AddFriend("cbmsim",digifile)

tree.SetBranchStatus("TrackPrePostFit.*",1)
tree.SetBranchStatus("TrackPostFit.*",1)
tree.SetBranchStatus("TpcCluster.*",1)
tree.SetBranchStatus("TpcPreCluster.*",1)
tree.SetBranchStatus("TpcDigi.*",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"
    rfile.Close()
    exit()

clusterfile=open(args.outfolder+'/cluster.inc','w')
reclusterfile=open(args.outfolder+'/recluster.inc','w')
digifile=open(args.outfolder+'/digis.inc','w')
redigifile=open(args.outfolder+'/redigis.inc','w')
planefile=open(args.outfolder+'/planes.inc','w')

colors=['Red','Green','Blue','Yellow','Cyan','Magenta','Coral','Firebrick','NavyBlue','OrangeRed','SpringGreen',
        'BrightGold','Med_Purple']
ev=-1
col1=0
col2=0
col3=0
colstep=0.5
clrad=0.25
clusterfile.write("#include \"colors.inc\"\n#declare cluster=union\n{\n")
digifile.write("#include \"colors.inc\"\n#declare digis=union\n{\n")
reclusterfile.write("#include \"colors.inc\"\n#declare recluster=union\n{\n")
redigifile.write("#include \"colors.inc\"\n#declare redigis=union\n{\n")
planefile.write("#include \"colors.inc\"\n#declare planes=union\n{\n")
for e in tree:
    ev+=1
    if args.event!=ev:
        continue
    for trk in e.TrackPrePostFit:
        cand=trk.getCand()
        detid=cand.getDetIDs()[0]
        hits=cand.getHitIDs(detid)
        colreset=0
        for hit in hits:
            sphit=e.TpcPreSPHit.At(hit)
            cl=e.TpcPreCluster.At(sphit.getClusterID())
            line='#if (frame_number<20)\n'
            line+="sphere {<"+str(cl.pos().X())+","+str(cl.pos().Y())+","+str(cl.pos().Z())+">, "+str(clrad)
            line+="\n texture{ pigment{ color "+colors[colreset]+"}}}\n"
            line+='#else\n'
            line+="sphere {<"+str(cl.pos().X())+","+str(cl.pos().Y())+","+str(cl.pos().Z())+">, "+str(clrad)
            line+="\n texture{ pigment{ color "+colors[colreset]+" transmit 0.7}}}\n"
            line+='#end\n'
            clusterfile.write(line)
            digimap=cl.getDigiMap()
            dc=0
            for dig in digimap:
                dc+=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
                line="sphere {<"+str(dx)+","+str(dy)+","+str(dz)+">, "+str(clrad/2)
                line+="\n texture{ pigment{ color "+colors[colreset]+"}}}\n"

                digifile.write(line)
            
            colreset+=1
            if colreset>len(colors)-1:
                colreset=0

    frame_offset=2
    for trk in e.TrackPostFit:
            cand=trk.getCand()
            detid=cand.getDetIDs()[0]
            hits=cand.getHitIDs(detid)
            mom=trk.getMom()
            ortho=mom.Orthogonal()
            colreset=0
            clcount=-1
            for hit in hits:
                clcount+=1
                if clcount>=1:
                    last_cl=cl
                sphit=e.TpcSPHit.At(hit)
                cl=e.TpcCluster.At(sphit.getClusterID())
                #cluster
                line='#if (frame_number>'+str(clcount-1+frame_offset)+')\n'
                line+="sphere {<"+str(cl.pos().X())+","+str(cl.pos().Y())+","+str(cl.pos().Z())+">, "+str(clrad)
                line+="\n texture{ pigment{ color Black}}}\n"
                line+='#end\n'
                reclusterfile.write(line)
                digimap=cl.getDigiMap()
                dc=-1
                #digis
                redigifile.write('#if (frame_number<='+str(clcount+frame_offset)+')\n')
                for dig in digimap:
                    dc+=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
                    line="sphere {<"+str(dx)+","+str(dy)+","+str(dz)+">, "+str(clrad/2)
                    line+="\n texture{ pigment{ color Gray60}}}\n"
                    redigifile.write(line)
                redigifile.write('#end\n')
                if clcount>=1:
                    diff=cl.pos()-last_cl.pos()
                    halfdist=diff
                    halfdist*=0.5
                    
                    ppos=cl.pos()-halfdist
                    ortho=ROOT.TVector3(1,0,1)

                    corner1=ppos+ortho+ROOT.TVector3(0,0.1,0)
                    corner2=ppos-ortho+ROOT.TVector3(0,0.1,0)
                    #planes
                    line='#if (frame_number='+str(clcount+frame_offset)+' | frame_number='+str(clcount+1)+')\n'
                    line+='box {<'+str(corner1.X())+','+str(corner1.Y())+','+str(corner1.Z())+'> , '
                    line+='<'+str(corner2.X())+','+str(corner2.Y())+','+str(corner2.Z())+'> \n '
                    line+="texture{ pigment{ color rgbt <1,0,0,.8>}}}\n"
                    line+='#end\n'
                    planefile.write(line)

    clusterfile.write("}")
    clusterfile.close()
    reclusterfile.write("}")
    reclusterfile.close()
    digifile.write("}")
    digifile.close()
    redigifile.write("}")
    redigifile.close()
    planefile.write("}")
    planefile.close()

    exit()