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
import ROOT, argparse,math

def setaxis(h):
    h.GetXaxis().SetTitleSize(0.05)
    h.GetXaxis().SetLabelSize(0.05)
    h.GetXaxis().SetTitle("Z (cm)")

    h.GetYaxis().SetTitleSize(0.05)
    h.GetYaxis().SetLabelSize(0.05)
    h.GetYaxis().SetTitleOffset(0.6)
    h.GetYaxis().SetTitle("R (cm)")

    h.GetZaxis().SetTitleSize(0.05)
    h.GetZaxis().SetLabelSize(0.05)
    h.GetZaxis().SetTitleOffset(0.6)
    h.GetZaxis().SetTitle("Deviation (cm)")
    return 0

def setRange(h,mini=0,maxi=0):
    for i in range(h.GetNbinsX()+1):
        for j in range(h.GetNbinsY()+1):
            if (maxi!=mini):
                cont=h.GetBinContent(i,j)
                if (cont>maxi):
                    h.SetBinContent(i,j,maxi)
                if (cont<mini):
                    h.SetBinContent(i,j,mini)
    if(mini!=maxi):
        h.GetZaxis().SetRangeUser(mini,maxi)
    return

def drawPrelim() :
    prelim=ROOT.TLatex()
    prelim.SetTextColor(ROOT.kGray)
    prelim.SetTextSize(0.1)
    prelim.SetTextAngle(20)
    prelim.SetNDC()
    prelim.DrawLatex(0.38,0.45,"preliminary")
    return

parser=argparse.ArgumentParser(description='plot deviation map')
parser.add_argument('filename',help='name of the devmap file',type=str)
parser.add_argument('--inv',help='create an inverse devmap',action='store_true')
parser.add_argument('--vd',help='drift velocity',type=float,default=0.0023764)
parser.add_argument('--range',help='range for the plot (min,max)',type=float,default=[-1,1],nargs=2)
parser.add_argument('--pdir',help='directory to save the pic in',type=str,default='NO')
parser.add_argument('--xdev',help='plot also the deviation in x-direction in xy projection',action='store_true')
parser.add_argument('--phidev',help='plot also the deviation in phi-direction in xy projection',action='store_true')
parser.add_argument('--out',help='plot the electrons which left the chamber')
parser.add_argument('--devd',help='plot the distribution of deviation',action='store_true')
parser.add_argument('--arrows',help='plot arrows of distortion',action="store_true")
parser.add_argument('--zmax',help='max z to bla',type=float, default=1000)
parser.add_argument('--zmin',help='min z to bla',type=float, default=1000)
parser.add_argument('--batch',help='batch mode. just produce the pictures and quit',action='store_true')
parser.add_argument('--prelim',help='draw preliminary',action='store_true')
parser.add_argument('--rout',help='name of the root output file',type=str,default='')
parser.add_argument('--nR',help='number of bins in R',type=int,default=-1)
parser.add_argument('--nZ',help='number of bins in Z',type=int,default=-1)
parser.add_argument('--phd',help='plot without colz for further use in phd',action='store_true')
parser.add_argument('--printBin',help='print only the value of a certain bin',type=int,default=[-1,-1],nargs=2)
args=parser.parse_args()

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

name=args.filename.split('/')

if args.pdir=='NO':
    c1=ROOT.TCanvas("bla",name[-1],700,1000)
    c2=ROOT.TCanvas("bla2",name[-1],700,1000)
    
fdevMap=ROOT.TpcDevmapCylWrapper(args.filename,args.vd,args.inv)



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

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

if args.nR!=-1:
    nRbin=args.nR
    rwidth=float(rmax-rmin)/float(nRbin)
else:
    rwidth = fdevMap.getRWidth()

if args.nZ!=-1:
    nZbin=args.nZ
    zwidth=float(zmax-zmin)/float(nZbin)
else:
    zwidth = fdevMap.getZWidth()

print "Deviation Map initialized"
print "r-Max:",rmax,"r-Min:",rmin
print "z-Max:",zmax,"z-Min:",zmin
print "r-Width:",rwidth,"z-Width:",zwidth
if args.zmax==1000:
    args.zmax=zmax
if args.zmin==1000:
    args.zmin=zmin

fdevMap.printc()

ROOT.gStyle.SetOptStat(0)
hrdev =ROOT.TH2D("hrdev","Radial deviation y=0 [cm]",nZbin,zmin,zmax,nRbin,rmin,rmax)
hrdev2=ROOT.TH2D("hrdev2","Radial deviation x=0 [cm]",nZbin,zmin,zmax,nRbin,rmin,rmax)
hpdev =ROOT.TH2D("hpdev","Perpendicular deviation y=0 [cm]",nZbin,zmin,zmax,nRbin,rmin,rmax)
hpdev2=ROOT.TH2D("hpdev2","Perpendicular deviation x=0 [cm]",nZbin,zmin,zmax,nRbin,rmin,rmax)
hzdev =ROOT.TH2D("hzdev","Z deviation y=0 [cm]",nZbin,zmin,zmax,nRbin,rmin,rmax)
hzdev2=ROOT.TH2D("hzdev2","Z deviation x=0 [cm]",nZbin,zmin,zmax,nRbin,rmin,rmax)

hdx=ROOT.TH1D("hdx","x dev distr",1000,0,110000)
for nr in range(nRbin):
    for nz in range(nZbin):
                
        zpos=(nz+0.5)*zwidth+zmin
        rpos=(nr+0.5)*rwidth+rmin
        point=ROOT.TVector3(rpos,0.,zpos)
        value=fdevMap.value(point)
        point2=ROOT.TVector3(0.,rpos,zpos)
        value2=fdevMap.value(point2)
        
        if args.printBin[0]!=-1 and args.printBin[1]!=-1:
            if args.printBin[0]==nr and args.printBin[1]==nz:
                print "the dev value at ({0},{1})/({2},{3})  is:".format(nr,nz,rpos,zpos)
                print "y=0"
                value.Print()
                print "x=0"
                value2.Print()
                exit()
                
        hrdev .SetBinContent(nz+1,nr+1,value.X())
        hpdev .SetBinContent(nz+1,nr+1,value.Y())
        hzdev .SetBinContent(nz+1,nr+1,value.Z())
        
        hrdev2.SetBinContent(nz+1,nr+1,value2.X())
        hpdev2.SetBinContent(nz+1,nr+1,value2.Y())
        hzdev2.SetBinContent(nz+1,nr+1,value2.Z())
#        print "*****",nr,nz,"******"\
        #print "zpos=",zpos, value2.Z()

set_palette("palette",256)

setaxis(hrdev)
setaxis(hpdev)
setaxis(hzdev)
setRange(hrdev ,args.range[0],args.range[1])
setRange(hpdev ,args.range[0],args.range[1])
setRange(hzdev ,args.range[0],args.range[1])
setRange(hrdev2,args.range[0],args.range[1])
setRange(hpdev2,args.range[0],args.range[1])
setRange(hzdev2,args.range[0],args.range[1])

if args.pdir=='NO':
    opt='colz'
    if args.phd:
        opt=''
    c1.Divide(1,3,0.001,0.001,0)
    c1.cd(1)
    hrdev.Draw(opt)
    c1.cd(2)
    hpdev.Draw(opt)
    c1.cd(3)
    hzdev.Draw(opt)
    c1.Update()
    
    c2.Divide(1,3,0.001,0.001,0)
    c2.cd(1)
    hrdev2.Draw(opt)
    c2.cd(2)
    hpdev2.Draw(opt)
    c2.cd(3)
    hzdev2.Draw(opt)
    c2.Update()

if args.xdev or args.phidev:
    arrows=[]
    hxybin=150
    hxy=ROOT.TH2D("hxy",'xy projection',hxybin,-15,15,hxybin,-15,15)
    hxyxdev=ROOT.TH2D('hxyxdev','xy projection of x deviation',hxybin,-15,15,hxybin,-15,15)
    hxyydev=ROOT.TH2D('hxyydev','xy projection of y deviation',hxybin,-15,15,hxybin,-15,15)
    hxypdev=ROOT.TH2D('hxypdev','xy projection of phi deviation',hxybin,-15,15,hxybin,-15,15)
    hxyrdev=ROOT.TH2D('hxyrdev','xy projection of r deviation',hxybin,-15,15,hxybin,-15,15)
    xstep=float(30./hxybin)
    ystep=float(30./hxybin)
    z=args.zmin
    while(z<=args.zmax):
        x=-15
        while(x<15):
            y=-15
            while(y<15):
                rad=math.sqrt(x**2+y**2)
                if rad<=5.5 or rad>=15:
                    y+=ystep
                    continue
                point=ROOT.TVector3(x,y,z)
                hxy.Fill(x,y)
                value=fdevMap.value(point)
                hxyxdev.Fill(x,y,value.X())
                hxyydev.Fill(x,y,value.Y())
                arrows.append(ROOT.TArrow())
                absl=math.sqrt(value.X()**2+value.Y()**2)
                arrows[-1].SetX1(x)
                arrows[-1].SetY1(y)
                arrows[-1].SetX2(x+(value.X()))#/absl*xstep*2))
                arrows[-1].SetY2(y+(value.Y()))#/absl*ystep*2))
                arrows[-1].SetOption("|>")
                arrows[-1].SetArrowSize(0.005)
                arrows[-1].SetFillColor(1)

                rotval=value
                rotval.RotateZ(-point.Phi())
                hxypdev.Fill(x,y,rotval.Y())
                hxyrdev.Fill(x,y,rotval.X())

                
                y+=ystep
            x+=xstep
        print "finished z-bin",z
        z+=1
        
    if args.xdev:
        c2=ROOT.TCanvas("xdev",args.filename+"xdev",500,1000)
        c2.Divide(1,2)
        c2.cd(1)
        hxyxdevnorm=Divideh(hxyxdev,hxy,'hxyxdevnorm','xy projection of deviation x norm')
        hxyydevnorm=Divideh(hxyydev,hxy,'hxyydevnorm','xy projection of deviation y norm')
        hxyxdevnorm.SetStats(0)
        setRange(hxyxdevnorm,args.range[0],args.range[1])
        hxyxdevnorm.Draw("colz")
        if args.arrows:
            for a in range(len(arrows)):
                if a%7==0:
                    arrows[a].Draw()
        c2.cd(2)
        hxyydevnorm.SetStats(0)
        setRange(hxyydevnorm,args.range[0],args.range[1])
        hxyydevnorm.Draw("colz")
    if args.phidev:
        c3=ROOT.TCanvas("rpdev",args.filename+"rpdev",500,1000)
        c3.Divide(1,2)
        c3.cd(1)
        hxypdevnorm=Divideh(hxypdev,hxy,'hxypdevnorm','xy projection of deviation phi norm')
        hxyrdevnorm=Divideh(hxyrdev,hxy,'hxyrdevnorm','xy projection of deviation rad norm')
        hxypdevnorm.SetStats(0)
        hxyrdevnorm.SetStats(0)
        setRange(hxyrdevnorm,args.range[0],args.range[1])
        setRange(hxypdevnorm,args.range[0],args.range[1])
        hxyrdevnorm.Draw("colz")
        c3.cd(2)
        hxypdevnorm.Draw("colz")
    if args.phidev or args.xdev:
        c4=ROOT.TCanvas("xy",args.filename+"xy",500,500)
        hxy.SetStats(0)
        hxy.Draw("colz")


if args.pdir!='NO':
    
    c2.cd(1).SetRightMargin(0.15)
    hrdev.SetTitle("")
    hrdev.Draw("colz")
    if args.prelim:
        drawPrelim()
    c2.SaveAs(args.pdir+name[-1]+'_hrdev.eps')

    hpdev.SetTitle("")
    hpdev.Draw("colz")
    if args.prelim:
        drawPrelim()
    c2.SaveAs(args.pdir+name[-1]+'_hpdev.eps')

    hzdev.SetTitle("")
    hzdev.Draw("colz")
    if args.prelim:
        drawPrelim()
    c2.SaveAs(args.pdir+name[-1]+'_hzdev.eps')

if args.rout!='':
    routfile=ROOT.TFile(args.rout,'recreate')
    hrdev.Write()
    hpdev.Write()
    hzdev.Write()
    if args.xdev:
        hxyxdev.Write()
        
    routfile.Close()
    
#point=ROOT.TVector3(-10,-8,0)
#value=fdevMap.value(point)
#print "some vla:",value.X(),value.Y(),value.Z()

if not args.batch:
    u=''
    while u!='q':
        u=raw_input("quit? (q)")