import ROOT, argparse,sys,math


def TVector3ToPov(vec):
    return "<"+str(vec.X())+","+str(vec.Y())+","+str(vec.Z())+">"

def TVector3ToRot(vec,dtheta=0,dphi=0):
    return "rotate <"+str(vec.Theta()*ROOT.TMath.RadToDeg()+dtheta)+",0,"+str(vec.Phi()*ROOT.TMath.RadToDeg()+dphi)+">"
    #return "rotate <"+str(vec.Theta()*ROOT.TMath.RadToDeg())+",0,0>"

def TVector3ToTrans(vec):
    return "translate <"+str(vec.X())+","+str(vec.Y())+","+str(vec.Z())+">"

def PovLine(vec1,vec2,rad,col,trans):
    return "cylinder {"+TVector3ToPov(vec1)+","+TVector3ToPov(vec2)+","+str(rad)+"\n texture{ pigment{ color "+str(col)+" transmit "+str(trans)+"}}}\n"

def PovSphere(vec,rad,texture):
    texture.replace("{","{{")
    texture.replace("}","}}")
    return "sphere {{ {0},{1} \n texture {{ {2} }} }}\n".format( TVector3ToPov(vec),rad,texture )

def PovArrow(length,rad,texture,transformations=""):
    texture.replace("{","{{")
    texture.replace("}","}}")
    
    cyllen=length-0.08    
    line='union \n{ cylinder\n'
    line+='{\n'
    line+=' <0,0,0>, <{0},0,0>,{1}\n'.format(cyllen,rad)
    line+=' texture {{ {0} }}'.format(texture)
    line+='}\n' #close cylinder
    line+='cone\n'
    line+='{\n'
    line+='<{0},0,0>, {1}, <{2},0,0>, 0\n'.format(cyllen,rad*2,length)
    line+=' texture {{ {0} }}'.format(texture)
    line+='}\n' #close cone
    line+='\n {0} \n'.format(transformations)
    line+='}' #close union
    return line


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,default=23)
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=0)
parser.add_argument('--hlp',help='print help',action='store_true')
args=parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

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)
tree.SetBranchStatus("CutCosmics.*",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()


trackfile=open(args.outfolder+'/genfit_track.inc','w')
digifile=open(args.outfolder+'/digis.inc','w')
planefile=open(args.outfolder+'/planes.inc','w')
reclplanefile=open(args.outfolder+'/recl_planes.inc','w')
digisonplanefile=open(args.outfolder+'/digisonplane.inc','w')
projectfile=open(args.outfolder+'/projections.inc','w')
camerafile=open(args.outfolder+'/camera.inc','w')
camerafile2=open(args.outfolder+'/camera2.inc','w')
lightfile=open(args.outfolder+'/light.inc','w')
spotlightfile=open(args.outfolder+'/spotlight.inc','w')
#spotlightfile2=open(args.outfolder+'/spotlight.inc','w')
clusterfile=open(args.outfolder+'/cluster.inc','w')
planeerrorfile=open(args.outfolder+'/plane_error.inc','w')
errorfile=open(args.outfolder+'/error.inc','w')
eigen2Dfile=open(args.outfolder+'/eigen2D.inc','w')
eigen3Dfile=open(args.outfolder+'/eigen3D.inc','w')

colors=['Red','Green','Blue','Yellow','Cyan','Magenta','Coral','Firebrick','NavyBlue','OrangeRed','SpringGreen',
        'BrightGold','Med_Purple']

ev=-1
digifile.write("#include \"colors.inc\"\n#declare digis=union\n{\n")
planefile.write("#include \"colors.inc\"\n #include \"my_textures.inc\"\n")
trackfile.write("#include \"colors.inc\"\n#declare track=union\n{\n")
planefile.write("#declare dplane=\nbox{<1,1,0.001> <-1,-1,0.001> \n texture{  pigment { color MediumTurquoise transmit 0.6 }\n finish  { ambient 0.15 diffuse 0.85 phong 0.8 phong_size 50 } }}\n\n")
planefile.write("#declare dplanes=union\n{\n")
planesize=0.08
reclplanefile.write('#declare reclplane=\nbox{{<{0},{0},0.001> <-{0},-{0},0.001> \n texture{{ pigment{{ color CadetBlue transmit 0.5 }} }} }}\n\n'.format(planesize))
reclplanefile.write("#declare reclplanes=union\n{\n")
digisonplanefile.write("#declare digisOnPlane=union\n{\n")
projectfile.write("#declare projections=union\n{\n")
lightfile.write("#declare mainlight=\nlight_source\n{\n")
spotlightfile.write("#declare spotlight1=\nlight_source\n{\n")
camerafile.write("camera\n { up <0,1,0> \nright <1,0,0>\n")
camerafile2.write("camera\n {\n")
clusterfile.write("#declare cluster=union\n{\n")
planeerrorfile.write("#declare plane_error=\ncylinder{<0,0,0>,<0,0,0.0015>,1 \n texture { pigment { color DarkOrchid transmit 0.5 } \n finish  { ambient 0.15 diffuse 0.85 phong 0.8 phong_size 50 } } }\n")
planeerrorfile.write('#declare plane_errors=union\n{\n')
errorfile.write('#declare error3D=\nsphere{<0,0,0>,1\ntexture { pigment { color DarkOrchid transmit 0.5 } \n finish  { ambient 0.15 diffuse 0.85 phong 0.8 phong_size 50 } } }\n')
errorfile.write('#declare errors3D=union\n{\n')
eigen2Dfile.write('#declare eigen2D=union\n{\n')
eigen3Dfile.write('#declare eigen3D=union\n{\n')

trackrad=0.01
clrad=0.25
for e in tree:
    ev+=1
    print 'event',ev
    if args.event!=ev:
        continue


    for tr in e.CutCosmics:
        print 'track'
        gtrack=tr.GetGTrack()
        planesU=tr.GetPlaneU()
        planesV=tr.GetPlaneV()
        planesO=tr.GetPlaneO()
        #planesW=tr.GetPlaneW()
        #digisOnPlane=tr.GetDigisOnPlane()
        normals=[]
        mom=tr.GetMom()
        mom.SetMag(1)
        meanz=tr.GetMeanZ()

        ncluster=tr.nhits()

        hitids=tr.GetHitIDs()


        #loop over genfit track and write to povray inc file
        pcounter=-1
        for p in gtrack:
            pcounter+=1
            trackfile.write(PovSphere(p,trackrad,"pigment  { color SlateBlue }\nfinish   { ambient 0.15 diffuse 0.85 phong 1 phong_size 100 reflection 0.5}"))
            if pcounter<len(gtrack)-1:
                direction=gtrack[pcounter+1]-p
                step=ROOT.TVector3(direction)
                step*=1./100.
                for i in range(100):
                    pos=ROOT.TVector3(step)
                    pos*=i
                    pos+=p
                    trackfile.write(PovSphere(pos,trackrad,"pigment  { color SlateBlue }\nfinish   { ambient 0.15 diffuse 0.85 phong 1 phong_size 100 reflection 0.5}"))



        clcount=-1
        for clid in hitids:
            print 'at cluster',clcount
            clcount+=1
            if clcount>=1:
                last_cluster=cluster
            sphit=e.TpcSPHit.At(clid)
            cluster=e.TpcCluster.At(sphit.getClusterID())
            digimap=cluster.getDigiMap()

            clusterfile.write('#if (frame_number={0} | frame_number>{1})\n'.format(clcount+1,ncluster-1))
            clusterfile.write(PovSphere(cluster.pos(),clrad/4,"Bronze_Metal"))
            clusterfile.write('#end')

            #create detplanes for projection
            ipl=clcount
            offset=ROOT.TVector3(planesO[ipl])
            U=ROOT.TVector3(planesU[ipl])
            V=ROOT.TVector3(planesV[ipl])
            #W=ROOT.TVector3(planesW[ipl])
            normal=ROOT.TVector3( V.Cross(U) )
            normal.SetMag(1)
            normals.append(normal)
            line=""
            line+='#if (frame_number={0} | frame_number>{1})\n'.format(clcount+1,ncluster-1)
            line+="object \n{\ndplane "
            line+=TVector3ToRot(normal,0,90)
            line+="\n"
            line+=TVector3ToTrans(offset)
            line+="\n"
            line+="}\n"
            line +='#end\n'
            planefile.write(line)

            axis1=ROOT.TVector3(cluster.axis(1))
            axis2=ROOT.TVector3(cluster.axis(2))
            rotangle=V.Angle(axis1)*ROOT.TMath.RadToDeg()

            planeerrorfile.write('\n#if (frame_number={0} | frame_number>{1})\n'.format(clcount+1,ncluster-1))
            planeerrorfile.write('object {{plane_error \nscale <{2},{3},{4}>\nrotate <0,0,0>  \n{0}\n{1} }}'.format(TVector3ToRot(normal,0,90),TVector3ToTrans(offset),math.sqrt(cluster.sig(3).Y()),math.sqrt(cluster.sig(3).Z()),1,rotangle ) )
            planeerrorfile.write('#end\n')

            errorfile.write('\n#if (frame_number={0} | frame_number>{1})\n'.format(clcount+1,ncluster-1))
            errorfile.write('object {{error3D scale <{2},{3},{4}> {0} {1} }}'.format(TVector3ToRot(normal,0,90),TVector3ToTrans(offset),math.sqrt(cluster.sig(3).Y()),math.sqrt(cluster.sig(3).Z()),math.sqrt(cluster.sig(3).X() ) ) )
            errorfile.write('#end\n')

            eigen2Dfile.write('\n#if (frame_number={0} | frame_number>{1})\n'.format(clcount+1,ncluster-1))
            eigentexture='pigment { color DarkOrchid } \n finish  { ambient 0.15 diffuse 0.85 phong 0.8 phong_size 10 }'
            axis1=ROOT.TVector3(cluster.axis(1))
            axis1.SetMag(math.sqrt(cluster.sig(3).Y()))
            eigen2Dfile.write(PovArrow(math.sqrt(cluster.sig(3).Y()),clrad/16.,eigentexture,'\n{0}\n{1}'.format(TVector3ToRot(normal,0,90),TVector3ToTrans(offset))))
            eigen2Dfile.write('\n')
            eigen2Dfile.write(PovArrow(math.sqrt(cluster.sig(3).Z()),clrad/16.,eigentexture,'\nrotate<0,0,90>\n{0}\n{1}'.format(TVector3ToRot(normal,0,90),TVector3ToTrans(offset))))
            eigen2Dfile.write('\n#end\n')

            eigen3Dfile.write('\n#if (frame_number={0} | frame_number>{1})\n'.format(clcount+1,ncluster-1))
            eigen3Dfile.write(PovArrow(math.sqrt(cluster.sig(3).Y()),clrad/16.,eigentexture,'\n{0}\n{1}'.format(TVector3ToRot(normal,0,90),TVector3ToTrans(offset))))
            eigen3Dfile.write('\n')
            eigen3Dfile.write(PovArrow(math.sqrt(cluster.sig(3).Z()),clrad/16.,eigentexture,'\nrotate<0,0,90>\n{0}\n{1}'.format(TVector3ToRot(normal,0,90),TVector3ToTrans(offset))))
            eigen3Dfile.write('\n')
            eigen3Dfile.write(PovArrow(math.sqrt(cluster.sig(3).X()),clrad/16.,eigentexture,'\nrotate<0,90,0>\n{0}\n{1}'.format(TVector3ToRot(normal,0,90),TVector3ToTrans(offset))))
            eigen3Dfile.write('\n#end\n')
            
            #get previous cluster and create reclustering plane
            if clcount>=1:
                diff=cluster.pos()-last_cluster.pos()
                halfdist=ROOT.TVector3(diff)
                halfdist*=0.5
                ppos=cluster.pos()-halfdist
                line=""
                line+="object \n{\nreclplane "
                line+=TVector3ToRot(diff,0,90)
                line+="\n"
                line+=TVector3ToTrans(ppos)
                line+="\n"
                line+="}\n"
                reclplanefile.write(line)

            #create all digis
            dc=-1
            line='#if (frame_number={0} | frame_number>{1})\n'.format(clcount+1,ncluster-1)
            digifile.write(line)
            digisonplanefile.write(line)
            projectfile.write(line)

            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
                dpos=ROOT.TVector3(dx,dy,dz)
                line="sphere {<"+str(dx)+","+str(dy)+","+str(dz)+">, "+str(clrad/8)
                line+="\n texture{ pigment{ color DarkOrchid }\nfinish   { ambient 0.15 diffuse 0.85 phong 0.8 phong_size 50 reflection 0.2}}}\n"
                digifile.write(line)
                
                dipos=ROOT.TVector3(planesO[clcount])
                dir1=ROOT.TVector3(planesU[clcount])
                dir1*=((dpos-planesO[clcount])*planesU[clcount])
                dir2=ROOT.TVector3(planesV[clcount])
                dir2*=((dpos-planesO[clcount])*planesV[clcount])
                dipos+=dir1+dir2
                dipos2=ROOT.TVector3(dpos-dipos)
                dipos2.SetMag(0.002)
                dipos2+=dipos
                line='cylinder {{ <0,0,-0.001> , <0,0,0.001>, {0}\n texture {{ pigment {{ color Black transmit 0}} \n finish   {{ ambient 0.15 diffuse 0.85 phong 0.8 phong_size 50 reflection 0.2}}}}\n {1} {2} }}\n'.format(clrad/8.,TVector3ToRot(normal,0,90),TVector3ToTrans(dipos))
                digisonplanefile.write(line)
                
                pline=""
                pline+=PovLine(dipos,dpos,clrad/48,"Black",0.5)
                projectfile.write(pline)

            digifile.write('#end')
            digisonplanefile.write('#end')
            projectfile.write('#end')

            print 'writing camera'
            #first the camera shifted perp to track
            planenum=clcount
            lookat=ROOT.TVector3(planesO[planenum])
            location=ROOT.TVector3(planesV[planenum])
            location*=2
            location+=lookat
            momoff=ROOT.TVector3(mom)
            momoff*=0.9
            location+=momoff
            line='#if ( frame_number={0})\n'.format(clcount+1)
            line+="location "+TVector3ToPov(location)+" \n"
            line+="look_at "+TVector3ToPov(lookat)+"\n"
            line+="#end\n"
            camerafile.write(line)
            
            #now the camera shifted along track
            lookat=ROOT.TVector3(cluster.pos())
            location=ROOT.TVector3(cluster.pos())
            momoff=ROOT.TVector3(mom)
            momoff*=2
            location+=momoff
            camerafile2.write('#if ( frame_number={0})\n look_at {1}\n location {2}\n#end\n'.format(clcount+1,TVector3ToPov(lookat),TVector3ToPov(location)))

            spotlightfile.write('#if ( frame_number={0})\n {2}, color White\n spotlight \n point_at {1}\n#end\n'.format(clcount+1,TVector3ToPov(lookat),TVector3ToPov(location)))
            print 'writing light'
            lightpos=ROOT.TVector3(planesV[planenum])
            lightpos*=1.
            lightpos+=planesO[planenum]
            lightmomoff=ROOT.TVector3(mom)
            lightmomoff*=1.5
            lightpos+=lightmomoff
            #lightdir1=ROOT.TVector3(planesU[planenum])
            #lightdir2=ROOT.TVector3(mom)
            lightdir1=ROOT.TVector3( (planesO[planenum]-lightpos).Orthogonal() )
            lightdir2=ROOT.TVector3( lightdir1.Cross( planesO[planenum]-lightpos ) ) 

            lightdir1.SetMag(0.05)
            lightdir2.SetMag(0.05)
            
            off1=ROOT.TVector3(lightdir1)
            off2=ROOT.TVector3(lightdir2)
            off1*=-1
            off2*=-1
            #off1.SetMag(1.25)
            #off2.SetMag(1.25)
            #lightpos-=off1
            #lightpos-=off2
            
            line='#if ( frame_number={0})\n'.format(clcount+1)
            line+=TVector3ToPov(lightpos)
            lightfile.write(line)

            lightfile.write("\n color White\n area_light")
            lightfile.write(TVector3ToPov(lightdir1)+" "+TVector3ToPov(lightdir2)+",5,5\n adaptive 1 \n jitter\n ")
            #lightfile.write("{0}\n".format(TVector3ToTrans(off1+off2) ))
            lightfile.write("#end\n")

            
           
    campos=ROOT.TVector3(planesV[int(ncluster/2)])
    campos*=-20
    campos+=planesO[int(ncluster/2)]
    momadd=ROOT.TVector3(mom)
    momadd*=4
    campos+=momadd
    camerafile.write("#if (frame_number>{0} )\n look_at {1}\n location {2} \n #end\n".format(ncluster-1,TVector3ToPov(planesO[int(ncluster/2)]),TVector3ToPov(campos)))
    camerafile.write("}")
    camerafile.close()

    camerafile2.write("#if (frame_number>{0} )\n look_at {1}\n location {2} \n #end\n".format(ncluster-1,TVector3ToPov(planesO[int(ncluster/2)]),TVector3ToPov(campos)))
    camerafile2.write("}")
    camerafile2.close()

    lightfile.write("#if (frame_number>{0} )\n <0,0,{1}>\n color White\n area_light <5,0,0> <0,5,0>,5,5\n adaptive 1 \n jitter\n #end\n".format(ncluster-1,meanz))
    lightfile.write("\n}\n")
    lightfile.close()


    spotlightfile.write("\n}\n")
    spotlightfile.close()

    clusterfile.write("}")
    clusterfile.close()
    
    planefile.write("}")
    planefile.close()

    reclplanefile.write("}")
    reclplanefile.close()

    trackfile.write("}")
    trackfile.close()

    digifile.write("}")
    digifile.close()

    digisonplanefile.write("}")
    digisonplanefile.close()

    projectfile.write("}")
    projectfile.close()

    planeerrorfile.write('}')
    planeerrorfile.close()

    errorfile.write('}')
    errorfile.close()

    eigen2Dfile.write('}')
    eigen2Dfile.close()

    eigen3Dfile.write('}')
    eigen3Dfile.close()

    exit()