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


parser=argparse.ArgumentParser(description='check the inversion of the devmap by comparing inverted with non inverted')
parser.add_argument('devmapfile',help='the devmap file to use',type=str)
parser.add_argument('--vd',help='drift velocity',type=float,default=0.0023764)

args=parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 

fdevMap  =ROOT.TpcDevmapCylWrapper(args.devmapfile,args.vd,False)
fdevMap_i=ROOT.TpcDevmapCylWrapper(args.devmapfile,args.vd,True)

nRbin=fdevMap.getLengthR()
nZbin=fdevMap.getLengthZ()

rmax = fdevMap.getRMax()
rmin = fdevMap.getRMin()
zmax = fdevMap.getZMax()
zmin = fdevMap.getZMin()

hpos    =ROOT.TH3D("hpos","pos;X (cm);Y (cm);Z (cm)"         ,30,-15,15,30,-15,15,72,0,72)
hdevpos =ROOT.TH3D("hdevpos","dev pos;X (cm);Y (cm);Z (cm)"  ,30,-15,15,30,-15,15,72,0,72)
hcorrpos=ROOT.TH3D("hcorrpos","corr pos;X (cm);Y (cm);Z (cm)",30,-15,15,30,-15,15,72,0,72)

hposYZ=ROOT.TH2D("hposYZ","posYZ;Y(cm);Z(cm)",30,-15,15,728,0,728)
hdevYZ={}
hcorrYZ={}
hdiffYZ={}

for coord in ["X","Y","Z"]:
    hdevYZ [coord]=ROOT.TH2D("hdev" +coord+"_YZ","dev" +coord+"_YZ;Y(cm);Z(cm)",30,-15,15,72,0,72)
    hcorrYZ[coord]=ROOT.TH2D("hcorr"+coord+"_YZ","corr"+coord+"_YZ;Y(cm);Z(cm)",30,-15,15,72,0,72)
    hdiffYZ[coord]=ROOT.TH2D("hdiff"+coord+"_YZ","diff"+coord+"_YZ;Y(cm);Z(cm)",30,-15,15,72,0,72)

hdiff1=ROOT.TH1D("hdiff1","diff1",1000,-0.1,0.1)
pos    =ROOT.TVector3(0,0,0)
devpos =ROOT.TVector3(0,0,0)
corrpos=ROOT.TVector3(0,0,0)

for ix in range(-15,15):
    for iy in range(-15,15):
        for iz in range(0,72,1):
            if math.sqrt(ix**2+iy**2)<5:
                continue
            pos.SetXYZ(ix,iy,iz)
            print "x,y,z=", ix,iy,iz
            dev =fdevMap.value(pos)
            devpos=pos+dev
            #devpos+=dev
            corr=fdevMap_i.value(devpos)
            corrpos=devpos-corr
            hpos.Fill(pos.X(),pos.Y(),pos.Z())
            hdevpos.Fill(devpos.X(),devpos.Y(),devpos.Z())
            hcorrpos.Fill(devpos.X(),devpos.Y(),devpos.Z())
            hdiff1.Fill((pos-corrpos).Mag())
            
            if ix==0:
                hdevYZ["X"].Fill(pos.Y(),pos.Z(),dev.X())
                hdevYZ["Y"].Fill(pos.Y(),pos.Z(),dev.Y())
                hdevYZ["Z"].Fill(pos.Y(),pos.Z(),dev.Z())
                
                hcorrYZ["X"].Fill(pos.Y(),pos.Z(),corr.X())
                hcorrYZ["Y"].Fill(pos.Y(),pos.Z(),corr.Y())
                hcorrYZ["Z"].Fill(pos.Y(),pos.Z(),corr.Z())
                        
                hdiffYZ["X"].Fill(pos.Y(),pos.Z(),dev.X()-corr.X())
                hdiffYZ["Y"].Fill(pos.Y(),pos.Z(),dev.Y()-corr.Y())
                hdiffYZ["Z"].Fill(pos.Y(),pos.Z(),dev.Z()-corr.Z())

doublecountdev=0
doublecountcorr=0

#for ibinX in range(hpos.GetNbinsX()):
    #for ibinY in range(hpos.GetNbinsY()):
        #for ibinZ in range(hpos.GetNbinsZ()):

for ibin in range(hpos.GetSize()):
    if ibin%1000==0:
        print ibin, 'of' ,hpos.GetSize()
    devposbin=hdevpos.GetBinContent(ibin)
    corrposbin=hcorrpos.GetBinContent(ibin)
    if devposbin>1:
        doublecountdev+=1
    if corrposbin>1:
        doublecountcorr+=1

canv=ROOT.TCanvas("canv",'canvas',1000,1000)
canv.Divide(2,2)
canv.cd(1)
hdiff1.Draw()
#canv.cd(2)
#hdevpos.Draw("colz")
#canv.cd(3)
#hcorrpos.Draw("colz")

canvases=[]
counter=0
for coord in ["X","Y","Z"]:
    counter+=1
    canvases.append(ROOT.TCanvas("c"+coord,coord,1000,1000))
    canvases[-1].Divide(2,2)
    canvases[-1].cd(1)
    hdevYZ[coord].SetStats(0)
    hdevYZ[coord].Draw("colz")
    canvases[-1].cd(2)
    hcorrYZ[coord].SetStats(0)
    hcorrYZ[coord].Draw("colz")
    canvases[-1].cd(3)
    hdiffYZ[coord].SetStats(0)
    hdiffYZ[coord].Draw("colz")



print 'doublecounts: dev=', doublecountdev, ' corr=',doublecountcorr

thisIsTheEnd()