import sys,os
sys.path.append('/home/mberger/fopiroot/fopiroot_dev/macro/tpc/FOPI/mberger')
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1')
import glob,math,argparse,ROOT
from functions import set_palette

def drawHist(c,h,fopt,u,v,cl,cog,mc):
    latex=ROOT.TLatex()
    text1=[]
    hproj1=h.ProjectionY()
    hproj2=h.ProjectionX()
    hproj1.SetTitle('Projection 1')
    hproj2.SetTitle('Projection 2')
    c.cd(1)
    h.SetStats(0)
    h.Draw('colz')
    c.cd(2)
    fitfunc.SetParameter(0,hproj1.GetMaximum())
    fitfunc.SetParameter(1,hproj1.GetMean())
    fitfunc.SetParameter(2,hproj1.GetRMS())
    hproj1.Fit(fitfunc,fopt)
    text1.append('Projection 1 Fit:')
    text1.append('Mean: {0:.3} #pm {1:.3}'.format(fitfunc.GetParameter(1),fitfunc.GetParError(1)))
    hproj1.Draw()
    c.cd(3)
    hproj2.Draw()
    fitfunc.SetParameter(0,hproj2.GetMaximum())
    fitfunc.SetParameter(1,hproj2.GetMean())
    fitfunc.SetParameter(2,hproj2.GetRMS())
    hproj2.Fit(fitfunc,fopt)
    c.cd(4)
    text1.append('Projection 2 Fit:')
    text1.append('Mean: {0:.3} #pm {1:.3}'.format(fitfunc.GetParameter(1),fitfunc.GetParError(1)))
    text1.append('U: theta={0:.3} Phi={1:.3}'.format(u.Theta()*ROOT.TMath.RadToDeg(),u.Phi()*ROOT.TMath.RadToDeg()))
    text1.append('V: theta={0:.3} Phi={1:.3}'.format(v.Theta()*ROOT.TMath.RadToDeg(),v.Phi()*ROOT.TMath.RadToDeg()))
    text1.append('Cluster shift: {0:.3} ({1:.3})'.format((cl-cog).Perp(),(cl-cog).Mag()))
    text1.append('Cluster Res: {0:.3} ({1:.3})'.format((cl-mc).Perp(),(cl-mc).Mag()))
    text1.append('Cog Res: {0:.3} ({1:.3})'.format((cog-mc).Perp(),(cog-mc).Mag()))

    x=0
    y=.9
    for t in text1:
        latex.DrawLatex(x,y,t)
        y-=0.1
    c.Update()

parser=argparse.ArgumentParser(description="plots digisposition projected on detector plane")
parser.add_argument('recofile',help='the reco file',type=str,default='./')
parser.add_argument('--pfile',help='the pdf to write into',type=str,default='None')
parser.add_argument('--events',help='number of events to run',type=int,default=-1)

args=parser.parse_args()

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

if args.recofile=='./':
    print 'no recofile given'
    exit()


RFile=ROOT.TFile(args.recofile,'read')
tree=RFile.Get('cbmsim')
if tree==None:
    print 'no tree found'
    RFile.Close()
    exit()


canv=ROOT.TCanvas('canv','digis on plane',1000,1000)
canv.Divide(2,2)

#hproj1=ROOT.TH1D('hproj1','Projection 1',100,-2,2)
#hproj2=ROOT.TH1D('hproj2','Projection 2',100,-2,2)
fitfunc=ROOT.TF1('gaus','gaus',-2,2)
fitfunc.SetLineWidth(1)
fitfunc.SetLineColor(2)

if args.pfile!='None':
    canv.SaveAs(args.pfile+'[')

ev=-1
for e in tree:
    ev+=1
    if args.events!=-1 and ev==args.events:
        if args.pfile!='None':
            canv.SaveAs(args.pfile+']')
        break

    trackcount=-1
    for tr in e.TrackFitStat_0:
        trackcount+=1

        print 'track'
        Umean=ROOT.TVector3(0,0,0)
        Vmean=ROOT.TVector3(0,0,0)
        U=tr.GetPlaneU()
        V=tr.GetPlaneV()
        O=tr.GetPlaneO()

        first=True
        for c in range(tr.nhits()):
            clPos=ROOT.TVector3(tr.GetHitPositionsX(c),tr.GetHitPositionsY(c),tr.GetHitPositionsZ(c))
            clCogPos=tr.GetCogPos(c)
            mcPos=tr.GetMCRefPos(c)
            print 'cluster pos', clPos.X(),clPos.Y(),clPos.Z()
            print 'clCogPos',clCogPos.X(),clCogPos.Y(),clCogPos.Z()
            print 'mcPos',mcPos.X(),mcPos.Y(),mcPos.Z()

            clPosUV=ROOT.TVector3((clPos-O[c])*U[c]/U[c].Mag(),(clPos-O[c])*V[c]/V[c].Mag(),0)
            clCogPosUV=ROOT.TVector3((clCogPos-O[c])*U[c]/U[c].Mag(),(clCogPos-O[c])*V[c]/V[c].Mag(),0)
            mcPosUV=ROOT.TVector3((mcPos-O[c])*U[c]/U[c].Mag(),(mcPos-O[c])*V[c]/V[c].Mag(),0)

            print 'clusterpos UV',clPosUV.X(),clPosUV.Y()
            print 'cluster Cog pos UV',clCogPosUV.X(),clCogPosUV.Y()
            print 'mc pos UV',mcPosUV.X(),mcPosUV.Y()
            hdigisOnPlane=ROOT.TH2D('hdigisOnPlane','Digis on plane (cluster)',
                                    100,clPosUV.X()-1,clPosUV.X()+1,20,clPosUV.Y()-1,clPosUV.Y()+1)
            if first:
                hdigisOnPlaneT=ROOT.TH2D('hdigisOnPlaneT','Digis on plane (track)',
                                         100,clPosUV.X()-1,clPosUV.X()+1,20,clPosUV.Y()-1,clPosUV.Y()+1)
                first=False

            digisOnPlane=tr.GetDigisOnPlane(c)
            digisOnPlaneAmp=tr.GetDigisOnPlaneAmp(c)

            for idigi in range(len(digisOnPlane)):
                digipos=digisOnPlane[idigi]
                digiamp=digisOnPlaneAmp[idigi]
                hdigisOnPlane.Fill(digipos.X(),digipos.Y(),digiamp)
                hdigisOnPlaneT.Fill(digipos.X(),digipos.Y(),digiamp)
                print 'digipos:', digipos.X(),digipos.Y()

            fitopt='MEIQ'
            drawHist(canv,hdigisOnPlane,fitopt,U[c],V[c],clPos,clCogPos,mcPos)
            canv.cd(1)
            clusterPoint=ROOT.TPolyMarker()
            clusterPoint.SetPoint(0,clPosUV.X(),clPosUV.Y())
            clusterPoint.SetMarkerStyle(2)
            clusterPoint.SetMarkerSize(1)
            clusterPoint.Draw("same")

            clusterCogPoint=ROOT.TPolyMarker()
            clusterCogPoint.SetPoint(0,clCogPosUV.X(),clCogPosUV.Y())
            clusterCogPoint.SetMarkerStyle(5)
            clusterCogPoint.SetMarkerSize(1)
            clusterCogPoint.Draw("same")
            
            MCPoint=ROOT.TPolyMarker()
            MCPoint.SetPoint(0,mcPosUV.X(),mcPosUV.Y())
            MCPoint.SetMarkerStyle(4)
            MCPoint.SetMarkerSize(1)
            #MCPoint.SetMarkerColor
            MCPoint.Draw("same")
            canv.Update()


            Umean+=U[c]
            Vmean+=V[c]

            if args.pfile!='None':
                canv.SaveAs(args.pfile)
            
            if args.events==-1:
                u=raw_input('next cluster?')
                if u=='q':
                    if args.pfile!='None':
                        canv.SaveAs(args.pfile+']')
                    exit()
                
            hdigisOnPlane.Reset()
            canv.cd(4).Clear()


        Umean*=1./tr.nhits()
        Vmean*=1./tr.nhits()
        drawHist(canv,hdigisOnPlaneT,fitopt,Umean,Vmean,clPos,clCogPos,mcPos)

        if args.pfile!='None':
            canv.SaveAs(args.pfile)

        if args.events==-1:
            u=raw_input('next track?')
            if u=='q':
                if args.pfile!='None':
                    canv.SaveAs(args.pfile+']')
                exit()

        hdigisOnPlaneT.Reset()
        canv.cd(4).Clear()
            

u=raw_input('done?')