import os,sys
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1')
import argparse
sys.path.append('$PANDAPATH/macro/tpc/FOPI/mberger')
from root_pov_tools import *

parser=argparse.ArgumentParser(description="Create cluster picture from recofile")
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=22)
parser.add_argument('outfolder',help='the folder to write the output to',type=str)
parser.add_argument('--trackrad',help='radius of the track',type=float,default=0.01)
args=parser.parse_args()


import ROOT
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")
tree.SetBranchStatus('*',0)
tree.SetBranchStatus('CutCosmics.*',1)

trackfile=open(args.outfolder+'/genfit_track.inc','w')
planefile=open(args.outfolder+'/planes.inc','w')
camerafile=open(args.outfolder+'/camera.inc','w')
lightfile=open(args.outfolder+'/light.inc','w')
clusterfile=open(args.outfolder+'/cluster.inc','w')
mcpointfile=open(args.outfolder+'/mcpoint.inc','w')
residualfile=open(args.outfolder+'/residuals.inc','w')
mcresidualfile=open(args.outfolder+'/mcresiduals.inc','w')

trackfile.write("#include \"colors.inc\"\n")
trackfile.write("#include \"transforms.inc\"\n")
trackfile.write("#declare track_spline=\nspline{\n natural_spline\n")

planefile.write("#declare dplane=\nbox{<1,1,0.001> <-1,-1,0.001> \n texture{ My_Ruby_Glass }}\n\n")
planefile.write("#declare dplanes=union\n{\n")

lightfile.write("#declare mainlight=\nlight_source\n{\n")

camerafile.write("camera\n {\n")

clusterfile.write("#declare cluster=union\n{\n")
mcpointfile.write("#declare mcpoint=union\n{\n")

residualfile.write('#declare residual=\nunion{')
mcresidualfile.write('#declare mcresidual=\nunion{')

trackrad=args.trackrad
clrad=args.trackrad*2

ev=-1
for e in tree:
    ev+=1
    print 'event',ev
    if args.event!=ev:
        continue

    done=False
    trcounter=-1
    print 'number of tracks:',e.CutCosmics.GetEntries()
    for tr in e.CutCosmics:
        trcounter+=1
        #if trcounter>0:
        #    continue
        mom=tr.GetMom()
        print mom.Mag()
        #if mom.Mag()>0.04:
        #    continue

        ncluster=tr.nhits()
        print "number of cluster:",ncluster
        gtrack=tr.GetGTrack()


        #loop over genfit track and write to povray inc file
        pcounter=-1
        tstepsize=1./(len(gtrack)-1)
        direction=gtrack[1]-gtrack[0]
        trackfile.write("{0},{1},\n".format(-tstepsize,TVector3ToPov(gtrack[0]-direction)))
        for p in gtrack:
            pcounter+=1
            if pcounter<len(gtrack)-1:
                direction=gtrack[pcounter+1]-p
            trackfile.write("{0},{1},\n".format(pcounter*tstepsize,TVector3ToPov(p)))
        pcounter+=1
        trackfile.write("{0},{1}\n}}".format(pcounter*tstepsize,TVector3ToPov(gtrack[-1]+direction)))
        trackfile.write("\n#declare track=\nunion\n{\n#local Nr=0;\n#local EndNr=1;\n")
        trackfile.write("#while(Nr<=EndNr)\nobject\n{{\nsphere\n{{\n<0,0,0>,{0}\ntexture\n{{\n".format(trackrad))
        trackfile.write("pigment{{ color SlateBlue transmit 0.2}}\nfinish {{ ambient 0.15 diffuse 0.85 phong 1 phong_size 100 reflection 0.2}} \n }} \ntranslate track_spline(Nr)\n}}\n}}\n".format())
        #trackfile.write("Spline_Trans(track_spline,Nr,y,0.05,0.7)\n}\n#local Nr=Nr+0.001;\n#end\n}\n")
        trackfile.write("#local Nr=Nr+0.001;\n#end\n}\n")


        #loop over cluster
        for clid in range(ncluster):
            pos=tr.GetHitPosition(clid)
            mcpos=tr.GetMCRefPos(clid)
            trpos=tr.GetProjectionPoint(clid)
            trmom=tr.GetProjectionPointsMom(clid)
            trmom.SetMag(1)
            residual=ROOT.TVector3(pos-trpos)
            residual_half=ROOT.TVector3(residual)
            residual_half*=0.5
            mcresidual=ROOT.TVector3(pos-mcpos)

            #create cluster spheres
            print 'creating cluster'
            clusterfile.write('#if (frame_number={0} | frame_number>{1})\n'.format(clid,ncluster-1))
            clusterfile.write(PovSphere(pos,clrad/4,"Bronze_Metal"))
            clusterfile.write('#end')
            #create MCpoints
            print 'creating mcpoint'
            mcpointfile.write('#if (frame_number={0} | frame_number>{1})\n'.format(clid,ncluster-1))
            mcpointfile.write(PovSphere(mcpos,clrad/4,"Chrome_Metal"))
            mcpointfile.write('#end')

            #create residuals
            residualfile.write('#if (frame_number={0} | frame_number>{1})\n'.format(clid,ncluster-1))
            mcresidualfile.write('#if (frame_number={0} | frame_number>{1})\n'.format(clid,ncluster-1))
            residualfile.write(PovLine(trpos,pos,0.001,'Black',0))
            residualfile.write('#end')
            mcresidualfile.write(PovLine(mcpos,pos,0.001,'DarkSlateBlue',0))
            mcresidualfile.write('#end')


            #create cameras
            print 'creating camera'
            lookat=ROOT.TVector3(pos)
            trmom.SetMag(0.1)
            residualoffset=ROOT.TVector3(residual)
            residualoffset*=0.5
            location=pos+trmom+residualoffset
            trmom.SetMag(1)
            camerafile.write('#if ( frame_number={0})\n look_at {1}\n location {2}\n#end\n'.format(clid,TVector3ToPov(lookat),TVector3ToPov(location)))

            #create light
            print "creating light"
            lightpos=location
            lightdir1=ROOT.TVector3(residual)
            lightdir2=ROOT.TVector3(residual.Cross(trmom))
            lightdir1.SetMag(0.5)
            lightdir2.SetMag(0.5)
            lightfile.write('#if ( frame_number={0})\n{1}'.format(clid,TVector3ToPov(lightpos)))
            lightfile.write("\n color White\n area_light")
            lightfile.write("{0} {1} ,5,5\n adaptive 1\n jitter\n".format(TVector3ToPov(lightdir1),TVector3ToPov(lightdir2)))
            lightfile.write("#end\n")
            
           

        done=True
        break
    if done:
        break

camerafile.write("}")
camerafile.close()

trackfile.close()

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

mcpointfile.write("}")
mcpointfile.close()

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

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

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