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
from open_DataChain import PyChainMaker

def getGraphFromProf(hist):
    profile=hist.ProfileX('{0}_proj'.format(hist.GetTitle()),0,-1)
    graph=ROOT.TGraphErrors()
    step=1./profile.GetNbinsX()
    for i in range(0,profile.GetNbinsX()*2,2):
        if profile.GetBinContent(i)==0:
            continue
        graph.SetPoint(i,1-(i*step),profile.GetBinContent(i))
        graph.SetPoint(i+1,1+(i*step),profile.GetBinContent(i))
        #graph.SetPointError(i,step,profile.GetBinError(i))
    graph.SetTitle('Profile of {0}'.format(hist.GetTitle()))
    graph.SetMarkerStyle(20)
    return graph


parser=argparse.ArgumentParser(description="plots digi position projected on detector plane")
parser.add_argument('recofile',help='the reco file',type=str,default='./')
parser.add_argument('--parfile',help='the parameter file',type=str,default='')
parser.add_argument('--rfile',help='filename to save histos into',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)
parser.add_argument('--title',help='title',type=str,default='')
parser.add_argument('--noMC',help='dont use mc, but track refpos',action='store_true')
parser.add_argument('--runs',help='if data, the runs to use',type=str,default='3888')
parser.add_argument('--pattern',help='if data the file name patters',type=str,default='')
parser.add_argument('--verbose',help='verbose output',action='store_true')

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()

if os.path.isdir(args.recofile):
    pyChain=PyChainMaker()
    pyChain.setPath(args.recofile)
    if args.noMC:
        pyChain.setRuns(args.runs)
        pyChain.setPattern(args.pattern)
    else:
        pyChain.setMC()
    #print pyChain.getFileList()
    tree=pyChain.getChain()
else:
    RFile=ROOT.TFile(args.recofile,'read')
    tree=RFile.Get('cbmsim')
if tree==None:
    print 'no tree found'
    RFile.Close()
    exit()

nentries=tree.GetEntries()
print nentries,"entries in tree"
if args.events==-2:
    args.events=nentries-2

if args.parfile!='':
    infos={}
    parfile=ROOT.TFile.Open(args.parfile)
    digiPar=parfile.Get('TpcDigiPar')
    infos['Dl']        =digiPar.getGas().Dl()
    infos['Dt']        =digiPar.getGas().Dt()
    infos['vd']        =digiPar.getGas().VDrift()
    infos['t0']        =digiPar.getFrontend().t0()
    infos['tbin']      =1./(digiPar.getFrontend().samplingFrequency())*1000 # 1/MHz -> ns
    for inf in infos:
        print inf,'=',infos[inf]

    
    
tree.SetBranchStatus("*",0)
tree.SetBranchStatus("CutCosmics.*",1)

canv=ROOT.TCanvas('canv','digi MC res '+args.title,1000,1000)
canv.Divide(3,2)
canv2=ROOT.TCanvas('canv2','Residuals '+args.title,1000,1000)
canv2.Divide(4,2)
canv3=ROOT.TCanvas('canv3','AmpDists '+args.title,1000,1000)
canv3.Divide(3,2)
canv4=ROOT.TCanvas('canv4','AmpDistsPlane '+args.title,1000,1000)
canv4.Divide(3,2)
canv5=ROOT.TCanvas('canv5','DigiLen '+args.title,1000,1000)
canv5.Divide(3,2)
canv6=ROOT.TCanvas('canv6','zoff calib'+args.title,1500,700)
canv6.Divide(2,1)

hx=ROOT.TH1D('hx','hx',1000,-1,1)
hy=ROOT.TH1D('hy','hy',1000,-1,1)
hz=ROOT.TH1D('hz','hz',1000,-1,1)

hresx=ROOT.TH1D('hresx','Residuals X',1000,-1,1)
hresy=ROOT.TH1D('hresY','Residuals Y',1000,-1,1)
hresz=ROOT.TH1D('hresz','Residuals Z',1000,-1,1)
hresxvsz=ROOT.TH2D('hresxvsz','Residuals X vs Z',75,0,75,1000,-1,1)
hresyvsz=ROOT.TH2D('hresYvsz','Residuals Y vs Z',75,0,75,1000,-1,1)
hreszvsz=ROOT.TH2D('hreszvsz','Residuals Z vs Z',75,0,75,1000,-1,1)
hresmag=ROOT.TH1D('hresmag','Residuals Mag',1000,-1,1)
hresmagvsz=ROOT.TH2D('hresmagvsz','Residuals Mag vs Z',75,0,75,1000,-1,1)

hrestvst=ROOT.TH2D("hrestvst","Residual t;t (ns);t_{digi}-t_{MC} (ns)",255,0,511,500,-10,10)

gx=ROOT.TGraph()
gx.SetTitle('resX')
gx.SetMarkerStyle(20)
gy=ROOT.TGraph()
gy.SetTitle('resY')
gz=ROOT.TGraph()
gz.SetTitle('resZ')
hamp=ROOT.TH1D('hamp','amp',1500,0,15000)
hampVsZ=ROOT.TH2D('hampVsZ','Amp vs Z',75,0,75,1500,0,15000)
hampVsAbsDist=ROOT.TH2D('hampVsAbsDist','Amp Vs Absolute Distance to MC',1000,0,1,150,0,15000)
hampVsPerpDist=ROOT.TH2D('hampVsPerpDist','Amp Vs Perp Distance to MC',1000,0,1,150,0,15000)
hampVsZDist=ROOT.TH2D('hampVsZDist','Amp Vs Z Distance to MC',1000,0,1,150,0,15000)

hampVsUMCres=ROOT.TH2D('hampVsUMCres','Amp Vs U Distance to MC',1000,0,1,150,0,15000)
hampVsVMCres=ROOT.TH2D('hampVsVMCres','Amp Vs V Distance to MC',1000,0,1,150,0,15000)
hampVsUVMCres=ROOT.TH2D('hampVsUVMCres','Amp Vs UV Distance to MC',1000,0,1,150,0,15000)

hresZvsDigiLen=ROOT.TH2D('hresZvsDigiLen','Resiudal in Z vs Digi Length',1000,0,1,20,0,20)
hresPerpVsDigiLen=ROOT.TH2D('hresPerpvsDigiLen','Resiudal Perp vs Digi Length',1000,0,1,20,0,20)
hresMagVsDigiLen=ROOT.TH2D('hresMagvsDigiLen','Resiudal Perp vs Digi Length',1000,0,1,20,0,20)

ev=-1
u='w'
for e in tree:
    ev+=1
    if ev%100==0:
        print 'At Event: {0}'.format(ev)
    itr=-1
    if args.verbose:
        print e.CutCosmics.GetEntries(),'tracks to process'
    for tr in e.CutCosmics:
        itr+=1
        U=tr.GetMidPlaneU()
        V=tr.GetMidPlaneV()
        O=tr.GetMidPlaneO()
        
        if args.verbose:
            print tr.nhits(),'cluster to process'
        for icl in range(tr.nhits()):
            
            digiamp=tr.GetDigisOnPlaneAmp(icl)
            digisOnPlane=tr.GetDigisOnPlane(icl)
            digiLen=tr.GetDigiLen(icl)
            
            if args.noMC:
                digires=tr.GetDigiTrackRes(icl)
                mcPos=tr.GetProjectionPoint(icl)
            else:
                digires=tr.GetDigiMCRes(icl)    
                mcPos=tr.GetMCRefPos(icl)
                
            mcPos=mcPos-O[icl]
            mcPosUV=ROOT.TVector3(mcPos*U[icl]/U[icl].Mag(),mcPos*V[icl]/V[icl].Mag(),0)

            zpos=tr.GetHitPositionsZ(icl)
            
            for idi in range(len(digires)):
                digiposUV=digisOnPlane[idi]
                gx.SetPoint(idi,digires[idi].X(),abs(digiamp[idi]))
                hx.Fill(digires[idi].X(),abs(digiamp[idi]))
                hamp.Fill(abs(digiamp[idi]))
                hampVsZ.Fill(zpos,abs(digiamp[idi]))
                hampVsAbsDist.Fill(digires[idi].Mag(),abs(digiamp[idi]))
                hampVsPerpDist.Fill(digires[idi].Perp(),abs(digiamp[idi]))
                hampVsZDist.Fill(digires[idi].Z(),abs(digiamp[idi]))
                hresx.Fill(digires[idi].X())
                hresy.Fill(digires[idi].Y())
                hresz.Fill(digires[idi].Z())
                hresmag.Fill(digires[idi].Mag())
                hresxvsz.Fill(zpos,digires[idi].X())
                hresyvsz.Fill(zpos,digires[idi].Y())
                hreszvsz.Fill(zpos,digires[idi].Z())
                hresmagvsz.Fill(zpos,digires[idi].Mag())

                hampVsUMCres.Fill(digiposUV.X()-mcPosUV.X(),abs(digiamp[idi]))
                hampVsVMCres.Fill(digiposUV.Y()-mcPosUV.Y(),abs(digiamp[idi]))
                hampVsUVMCres.Fill((digiposUV-mcPosUV).Perp(),abs(digiamp[idi]))

                hresZvsDigiLen.Fill(digires[idi].Z(),digiLen[idi])
                hresPerpVsDigiLen.Fill(digires[idi].Perp(),digiLen[idi])
                hresMagVsDigiLen.Fill(digires[idi].Mag(),digiLen[idi])

                if(args.parfile!=''):
                    timePos=int((zpos/infos['vd']-infos['t0'])/infos['tbin'])
                    timeRes=(digires[idi].Z()/infos['vd']-infos['t0'])/infos['tbin']
                    #print timePos,timeRes
                    hrestvst.Fill(timePos,timeRes)
                    
    if args.events!=-1 and args.events!=ev:
        continue
    canv.cd(1)
    gx.Draw('AP')
    canv.cd(4)
    hx.Draw()
    canv.cd(2)
    hamp.Draw()
    canv.cd(3)
    hampVsZ.Draw('colz')
    canv.cd(5)

    canv.cd(6)
    #hampVsAbsDist_proj.Draw()
 
    canv.Update()

    canv2.cd(1)
    hresx.Draw()
    canv2.cd(2)
    hresy.Draw()
    canv2.cd(3)
    hresz.Draw()    
    canv2.cd(4)
    hresmag.Draw()
    canv2.cd(5)
    hresxvsz.Draw('colz')
    canv2.cd(6)
    hresyvsz.Draw('colz')
    canv2.cd(7)
    hreszvsz.Draw('colz')
    canv2.cd(8)
    hresmagvsz.Draw('colz')
    canv2.Update()

    canv6.cd(1)
    hrestvst.Draw('colz')
    canv6.cd(2)
    hrestvst.FitSlicesY(0,0,-1,0,"QNRG3S")
    hresTslices=ROOT.gDirectory.Get(hrestvst.GetName()+"_1")
    hresTslices.Draw()
    canv6.Update()

    landau=ROOT.TF1("landau","landau",0,14000)
    hampVsPerpDist.FitSlicesY(landau,0,-1,0,"QNRG3S")
    landau_mpv = ROOT.gDirectory.Get(hampVsPerpDist.GetName()+"_1")
    landau_sigma = ROOT.gDirectory.Get(hampVsPerpDist.GetName()+"_2")
    canv3.cd(1)
    hampVsAbsDist.Draw('colz')
    canv3.cd(4)
    landau_mpv.Draw()
    #hampVsAbsDist_proj_trans=getGraphFromProf(hampVsAbsDist)
    #hampVsAbsDist_proj_trans.Draw('AP')    
    canv3.cd(2)
    hampVsPerpDist.Draw('colz')
    canv3.cd(5)
    landau_sigma.Draw()
    #hampVsPerpDist_proj_trans=getGraphFromProf(hampVsPerpDist)
    #hampVsPerpDist_proj_trans.Draw('AP')
    canv3.cd(3)
    hampVsZDist.Draw('colz')
    canv3.cd(6)
    hampVsZDist_proj_trans=getGraphFromProf(hampVsZDist)
    hampVsZDist_proj_trans.Draw('AP')
    canv3.Update()

    canv4.cd(1)
    hampVsUMCres.Draw('colz')
    canv4.cd(2)
    hampVsVMCres.Draw('colz')
    canv4.cd(3)
    hampVsUVMCres.Draw('colz')
    canv4.Update()

    canv5.cd(1)
    hresZvsDigiLen.Draw('colz')
    canv5.cd(2)
    hresPerpVsDigiLen.Draw('colz')
    canv5.cd(3)
    hresMagVsDigiLen.Draw('colz')
    canv5.Update()

    u=raw_input('next track?')
    gx.Clear()
    hamp.Reset()
    args.events=-1
    if u=='q':
        break
    if u=='n':
        continue

if args.rfile!='':
    rfile=ROOT.TFile(args.rfile,"RECREATE")
    canv6.Write()
    hrestvst.Write()
    hresTslices.Write()
if u=='q':
    exit()