import ROOT,os,sys,getopt
import rootUtils as ut
import shipunit as u
import RecoSettings
def SetGrid(pad):
  pad.SetGridx()
  pad.SetGridy()
  
def SetAxis(histo):
  histo.GetXaxis().SetTitleSize(.06)
  histo.GetYaxis().SetTitleSize(.06)
  histo.GetXaxis().SetTitleOffset(1.3)
  histo.GetYaxis().SetTitleOffset(1.3)
  histo.GetXaxis().SetNdivisions(5)
  histo.GetYaxis().SetNdivisions(5)


###################################################################################
###################################################################################
###################################################################################
def bookMFieldHistos(h):
  #ax0 = ';#chi^{2}/ndf;(P_{MCthuth}-P_{RECO})/P_{MCtruth}'
  #ax  = ';P_{MCthuth} [GeV];(P_{MCthuth}-P_{RECO})/P_{MCtruth}'
  colour = {"fit":2, "sim":1}
  for fall in ["fit", "sim"] :
    ut.bookHist(h,'magZ'  +fall, 'Field.Mag() vs Z at (0,0)'+fall+';z [cm];|B| [T]', 3000, 1000, 4000,200,0,2)
    h['magZ'+fall].SetLineColor(colour[fall])
    h['magZ'+fall].SetMarkerColor(colour[fall])
    ut.bookHist(h,'magXY1'+fall, 'Field.Mag() vs (X,Y) at Z=2500 cm, '+fall+';x [cm];y [cm]', 60, -300.,300.,60,-300,300)
    ut.bookHist(h,'magXY2'+fall, 'Field.Mag() vs (X,Y) at Z=2800 cm, '+fall+';x [cm];y [cm]', 60, -300.,300.,60,-300,300)
    ut.bookHist(h,'magXY3'+fall, 'Field.Mag() vs (X,Y) at Z=3000 cm, '+fall+';x [cm];y [cm]', 60, -300.,300.,60,-300,300)
  for histo in h:
    SetAxis(h[histo])
    h[histo].SetLineColor(1)
    h[histo].SetMarkerStyle(20)
    
###################################################################################
def plotMFieldHistos(h):
   ROOT.gStyle.SetStatY(0.9) 
   ROOT.gStyle.SetStatX(0.9) 
   ROOT.gStyle.SetStatW(0.3) 
   ROOT.gStyle.SetStatH(0.2) 
   ROOT.gStyle.SetPadBottomMargin(0.2)
   ROOT.gStyle.SetPadLeftMargin(0.2)  
   ROOT.gStyle.SetPadRightMargin(0.2)
   
   ROOT.gStyle.SetOptStat(0)
   ut.bookCanvas(h,"mfield1","mfield1",nx=600,ny=600,cx=1,cy=1)
   cv = h['mfield1'].cd(1)  
   h['magZsim'].Draw("")
   h['magZfit'].Draw("sames")
   h['mfield1'].Print('mfield1.gif')
   
   ut.bookCanvas(h,"mfield2","mfield2",nx=1200,ny=800,cx=3,cy=2)
   pad = 0
   for fall in ["fit", "sim"] :
    for ci in range (1,4):
      pad+=1
      cv  = h['mfield2'].cd(pad)
      h['magXY'+str(ci)+fall].Draw("colz")
    h['mfield2'].Print('mfield2.gif')
    

###################################################################################
###################################################################################
###################################################################################
def bookMomentumHistos(h,prefix=""):
  ax0 = ';#chi^{2}/ndf;(P_{MCthuth}-P_{RECO})/P_{MCtruth}'
  ax  = ';P_{MCthuth} [GeV];(P_{MCthuth}-P_{RECO})/P_{MCtruth}'
  ut.bookHist(h,'delPOverP0'+prefix, 'dP/P vs #chi^{2}/ndf'+prefix+ax0, 500,0.,50.,100,-0.5,0.5)
  ut.bookHist(h,'delPOverP1'+prefix, 'dP/P'+prefix+ax,100,0.,50.,100,-0.5,0.5)
  ut.bookHist(h,'delPOverP2'+prefix, 'dP/P ndf>'+str(RecoSettings.trackMinNofHits)+prefix+ax,100,0.,50.,100,-0.5,0.5)
  ut.bookHist(h,'delPOverP3'+prefix, 'dP/P cleaned #chi^{2}/ndf<'+str(RecoSettings.chi2CutOff)+prefix+ax,100,0.,50.,100,-0.5,0.5)
  ut.bookHist(h,'MomRes'+prefix, 'MomRes'+prefix+';P_{MCthuth} [GeV];#sigma_{Gauss}(#DeltaP/P)',5,0.,50.)
  for i in range (1,6):
    hhname='delPOverP3_'+str(i)+prefix
    cond  ='P_{MCtruth} < '+str(i*10)+" GeV"
    ut.bookHist(h, hhname,  cond+';(P_{MCthuth}-P_{RECO})/P_{MCtruth}',   200, -.2, .2)
    h[hhname].StatOverflows(True) # to include overflow/underflow in mean/sigma (range independence)
  
  for histo in h:
    h[histo].Sumw2()
    h[histo].SetLineColor(1)
    h[histo].SetMarkerStyle(20)
  
###################################################################################
def plotMomentumHistos(h,prefix=""):
   ROOT.gStyle.SetStatY(0.9) 
   ROOT.gStyle.SetStatX(0.9) 
   ROOT.gStyle.SetStatW(0.3) 
   ROOT.gStyle.SetStatH(0.2) 
   ROOT.gStyle.SetPadBottomMargin(0.2)
   ROOT.gStyle.SetPadLeftMargin(0.2)  
   
   ROOT.gStyle.SetOptStat(10)
   ut.bookCanvas(h,key='momentum1'+prefix,title='momentum1'+prefix,nx=1000,ny=800,cx=2,cy=2)
   for cid in range (1,5):
    cv = h['momentum1'+prefix].cd(cid)
    hname = 'delPOverP'+str(cid-1)+prefix
    h[hname].Draw("box")
    SetAxis(h[hname])
    if(cid==1):
      cv.SetLogx()
      h[hname].GetXaxis().SetRangeUser(.01,50)    
   h['momentum1'+prefix].Print('momentum1'+prefix+'.gif')
   
   ROOT.gStyle.SetStatY(0.93) 
   ROOT.gStyle.SetStatX(0.99) 
   ROOT.gStyle.SetStatW(0.4) 
   ROOT.gStyle.SetStatH(0.35) 
   f={}
   key='momentum2'+prefix
   ut.bookCanvas(h,key,title=key,nx=1200,ny=1200,cx=3,cy=2)
   rname = 'MomRes'+prefix
   for cvn in range (1,6):
    hname='delPOverP3_'+str(cvn)+prefix
    cv = h[key].cd(cvn) 
    h[hname].Draw("E")
    SetAxis(h[hname])
    ROOT.gStyle.SetOptFit(11111)
    f[hname] = ROOT.TF1(hname,'gaus',-.05, .05)
    f[hname].SetParameter(1,0.)
    f[hname].SetNpx(1000)
    h[hname].Fit(f[hname],"R")
    print "fit", f[hname].GetParameter(2), f[hname].GetParError(2)
    h[rname].SetBinContent(cvn, f[hname].GetParameter(2))
    h[rname].SetBinError  (cvn, f[hname].GetParError(2))
    h[key].Print(key+'.gif')
   
   ROOT.gStyle.SetOptStat(0)
   ROOT.gStyle.SetOptTitle(0)
   ut.bookCanvas(h,'momentum3'+prefix,title='momentum3'+prefix,nx=600,ny=600,cx=1,cy=1)
   h[rname].Draw('E')
   SetAxis(h[rname])
   h[hname].GetYaxis().SetTitleOffset(1.5)
   h['momentum3'+prefix].Print('momentum3'+prefix+'.gif')









###################################################################################
###################################################################################
###################################################################################   
def bookCompareFitHistos(h):
  ut.bookHist(h,'dNFoldFNnew',  'Number of fitted tracks;N(old);N(new)',   
	      10, -1, 9, 10, -1, 9)
  h['dNFoldFNnew'].SetMarkerSize(1.8)
  ut.bookHist(h,'FitCorrespondence',  'CorrespondenceBins: 0,5-ok',   
	      10, -1, 9)
  ut.bookHist(h,'NDFold',  'NDF old Fit;NDF', 140, 0, 70)
  ut.bookHist(h,'NDFnew',  'NDF new Fit;NDF', 140, 0, 70)
  ut.bookHist(h,'Chi2NDFold',  '#chi^{2}/NDF old Fit;#chi^{2}/NDF', 200, 0, 20)
  ut.bookHist(h,'Chi2NDFnew',  '#chi^{2}.NDF old Fit;#chi^{2}/NDF', 200, 0, 20)
  style={'old':0,'new':1}
  for name in ('NDF','Chi2NDF'):
    for st in style:
      h[name+st].SetLineColor(1+style[st])
      h[name+st].SetMarkerColor(1+style[st])
      h[name+st].SetMarkerStyle(22+style[st])
  for histo in h:
    h[histo].Sumw2()
    h[histo].SetLineWidth(2)


###################################################################################
def plotCompareFitHistos(h):
  ROOT.gStyle.SetStatY(0.9) 
  ROOT.gStyle.SetStatX(0.9) 
  ROOT.gStyle.SetStatW(0.3) 
  ROOT.gStyle.SetStatH(0.2) 
  ROOT.gStyle.SetOptStat(11)
  ut.bookCanvas(h,key='FitCompare',title='FitCompare',nx=800,ny=800,cx=2,cy=2)
  cv = h['FitCompare'].cd(1)
  cv.SetLogz()
  h['dNFoldFNnew'].Draw('coltext')  
  cv = h['FitCompare'].cd(2)
  cv.SetLogy()
  h['FitCorrespondence'].Draw()
  cv = h['FitCompare'].cd(3)
  cv.SetLogy()
  h['NDFold'].Draw()
  h['NDFnew'].Draw("sames")
  cv = h['FitCompare'].cd(4)
  cv.SetLogx()
  cv.SetLogy()
  h['Chi2NDFold'].Draw()
  h['Chi2NDFnew'].Draw("sames")
  h['FitCompare'].Print('FitCompare.gif')









###################################################################################
###################################################################################
###################################################################################   
def bookTrackingHistos(h):
  ut.bookHist(h,'dxdyALLTrHm',  'dxdy(MCtr,firstHit)@firstHitZ - all-;dx [cm];dy [cm]',   
	      2000, -100, 100, 2000, -100, 100)
  ut.bookHist(h,'dxdyALLTrHp',  'dxdy(MCtr,firstHit)@firstHitZ - all+;dx [cm];dy [cm]',   
	      2000, -100, 100, 2000, -100, 100)

  ut.bookHist(h,'dxdyHNLTrHm',  'dxdy(MCtr,firstHit)@firstHitZ - HNL-;dx [cm];dy [cm]',   
	      2000, -100, 100, 2000, -100, 100)
  ut.bookHist(h,'dxdyHNLTrHp',  'dxdy(MCtr,firstHit)@firstHitZ - HNL+;dx [cm];dy [cm]',   
	      2000, -100, 100, 2000, -100, 100)

  ut.bookHist(h,'dRvsPz',       'dR(MCtr,firstHit)@firstHitZ vs Pz;Pz_{MC} [GeV];dR [cm]',
	      1000, 0, 50, 1000, 0, 100)
  ut.bookHist(h,'dRvsZhit',     'dR(MCtr,firstHit)@firstHitZ vs firstHitZ;Z_{Hit0} [cm];dR [cm]',
	      2400, 2500, 3700, 1000, 0, 100)
  
  ut.bookHist(h,'YMC2dZMCVer',   'abs(YMCprop)/dZMCVer - HNL',   2000, 0,10)

  for histo in h:
    h[histo].Sumw2()
    h[histo].SetLineColor(1)

###################################################################################  
def plotMCTrackHitCorrelation(h):
   ROOT.gStyle.SetStatY(0.9) 
   ROOT.gStyle.SetStatX(0.9) 
   ROOT.gStyle.SetStatW(0.3) 
   ROOT.gStyle.SetStatH(0.2) 
   ROOT.gStyle.SetOptStat(11)

   ut.bookCanvas(h,key='dRCanvas',title='dRCanvas',nx=1200,ny=600,cx=2,cy=1)
   ROOT.gStyle.SetOptStat(1)
   cv = h['dRCanvas'].cd(1)
   h['dRvsPz'].Draw('colzbox')
   cv = h['dRCanvas'].cd(2)
   h['dRvsZhit'].Draw('colzbox')   
   h['dRCanvas'].Print('dRCanvas.gif')

   ROOT.gStyle.SetOptStat(11)
   ut.bookCanvas(h,key='ratioYZ',title='ratioYZ',nx=600,ny=600,cx=1,cy=1)
   ROOT.gStyle.SetOptStat(1)
   cv = h['ratioYZ'].cd(1)
   h['YMC2dZMCVer'].Draw()
   h['ratioYZ'].Print('ratioYZ.gif')
   
   for x in ['ALL','HNL']:
    ut.bookCanvas(h,key='MCTrackHit'+x,title='MCTrackHit'+x,nx=800,ny=800,cx=2,cy=2)
    ROOT.gStyle.SetOptStat(1)
    cv = h['MCTrackHit'+x].cd(1)
    h['dxdy'+x+'TrHp'].Draw('colz')
    cv = h['MCTrackHit'+x].cd(2)
    h['dxdy'+x+'TrHm'].Draw('colz')
    cv = h['MCTrackHit'+x].cd(3)
    SetGrid(cv)
    hh=h['dxdy'+x+'TrHp'].DrawCopy('colz')
    for ax in [hh.GetXaxis(), hh.GetYaxis()]:
      ax.SetRangeUser(-5,5)
    cv = h['MCTrackHit'+x].cd(4)
    SetGrid(cv)
    hh=h['dxdy'+x+'TrHm'].DrawCopy('colz')
    for ax in [hh.GetXaxis(), hh.GetYaxis()]:
      ax.SetRangeUser(-5,5)
    h['MCTrackHit'+x].Print('MCTrackHit'+x+'.gif')








###################################################################################
###################################################################################
###################################################################################   
def bookVertexHistos(h, prefix=""):
  if prefix=="NEW":
    ut.bookHist(h,'docaRel'+prefix,  'docaRel'+prefix+";step;doca/doca0",   5, 0, 5,  2000,  0., 2.)
    ut.bookHist(h,'doca'   +prefix,  'doca'   +prefix+";doca [cm]",   500,  0., 100.)
    h['doca'+prefix].SetLineColor(2)
    ut.bookHist(h,'doca0'  +prefix,  'doca0'  +prefix+";doca0 [cm]",   500,  0., 100.)
  ut.bookHist(h,'vdx2'+prefix,  'vdx2'+prefix,   500, 0, 5000, 400,  -20., 20.)
  ut.bookHist(h,'vdy2'+prefix,  'vdy2'+prefix,   500, 0, 5000, 400,  -20., 20.)
  ut.bookHist(h,'vdz2'+prefix,  'vdz2'+prefix,   500, 0, 5000, 4000, -400., 400.)
  ut.bookHist(h,'vdx2ndf'+prefix,  'vdx2ndf'+prefix,   500, 0, 5000, 400,  -20., 20.)
  ut.bookHist(h,'vdy2ndf'+prefix,  'vdy2ndf'+prefix,   500, 0, 5000, 400,  -20., 20.)
  ut.bookHist(h,'vdz2ndf'+prefix,  'vdz2ndf'+prefix,   500, 0, 5000, 4000, -400., 400.)
  if ( (prefix=="mUP") or (prefix=="mDN") ) : return
  for i in range (1,11):
    hhname='vdx2_'+str(i)+prefix
    ut.bookHist(h, hhname,  hhname,   100, -20., 20.)
    h[hhname].StatOverflows(True) # to include overflow/underflow in mean/sigma (range independence)
    hhname='vdy2_'+str(i)+prefix
    ut.bookHist(h, hhname,  hhname,   100, -20., 20.)
    h[hhname].StatOverflows(True)
    hhname='vdz2_'+str(i)+prefix
    ut.bookHist(h, hhname,  hhname,   200, -400., 400.)
    h[hhname].StatOverflows(True)
  for coord in ('x','y','z') :
    hname = "vres"+coord+prefix
    ut.bookHist(h, hname,  hname+"z_{0}-z_{truth} [cm];",   10, 0., 5000.)
  for histo in h:
    h[histo].Sumw2()
    h[histo].SetLineColor(1)


###################################################################################
def plotVertexHistos(h,prefix=""):
   ROOT.gStyle.SetStatY(0.9) 
   ROOT.gStyle.SetStatX(0.9) 
   ROOT.gStyle.SetStatW(0.3) 
   ROOT.gStyle.SetStatH(0.1) 
   ROOT.gStyle.SetOptStat(1)

   coordTitle = {1:'x',2:'y',3:'z'}
   cut        = {0:'',1:'ndf'}
   ut.bookCanvas(h,key='vertex2'+prefix,title='vertex2'+prefix,nx=600,ny=1200,cx=2,cy=3)
   for fall in cut: 
    for cid in coordTitle.keys():
      cv = h['vertex2'+prefix].cd(cid*2-1+fall)
      cv.SetBottomMargin(0.2)
      cv.SetLeftMargin(0.2)
      hname = 'vd'+coordTitle[cid]+'2'+cut[fall]+prefix
      h[hname].Draw("box")
      h[hname].GetXaxis().SetTitle("z_{0}-z_{truth} [cm]")
      ytitle = coordTitle[cid]+"_{truth}-v"+coordTitle[cid]+"_{reco} [cm]"
      h[hname].GetXaxis().SetTitleSize(.06)
      h[hname].GetYaxis().SetTitleSize(.06)
      h[hname].GetYaxis().SetTitleOffset(1.3)
      h[hname].GetYaxis().SetNdivisions(5)
   h['vertex2'+prefix].Print('vertex2'+prefix+'.gif')
   if ( (prefix=="mUP") or (prefix=="mDN") ) : return
 
   f={}
   for ct in coordTitle.values():
      key =ct+prefix
      ut.bookCanvas(h,key,title=key,nx=1600,ny=1200,cx=5,cy=2)
      f[key]={}
      for cvn in range (1,11):
	hname='vd'+ct+'2_'+str(cvn)+prefix
	cv = h[key].cd(cvn) 
	cv.SetBottomMargin(0.2)
	h[hname].Draw("E")
	xtitle = "v"+ct+"_{truth}-v"+ct+"_{reco} [cm]"
	h[hname].GetXaxis().SetTitle(xtitle)
	h[hname].GetXaxis().SetTitleSize(.06)    
	ROOT.gStyle.SetOptFit(11111)
	f[key][hname] = ROOT.TF1(hname,'gaus')
	f[key][hname].SetNpx(1000)
	f[key][hname].SetParameter(1,0.)
	f[key][hname].SetParameter(2,0.2*cvn)
	h[hname].Fit(f[key][hname])
	rmsht='vres'+key
	#h[rmsht].SetBinContent(cvn, h[hname].GetRMS())
	#h[rmsht].SetBinError  (cvn, h[hname].GetRMSError())
	h[rmsht].SetBinContent(cvn, f[key][hname].GetParameter(2))
	h[rmsht].SetBinError  (cvn, f[key][hname].GetParError(2))
      h[key].Print(ct+prefix+'.gif')
  
   key='vRes'+prefix
   ut.bookCanvas(h,key,title='vRes'+prefix,nx=800,ny=800,cx=2,cy=2)
   for cti in coordTitle :
    cv = h[key].cd(cti)
    rmsht ="vres"+coordTitle[cti]+prefix
    h[rmsht].SetTitle("vertex "+coordTitle[cti]+" resolution;z_{0}-z_{truth} [cm];#sigma_{Gauss}(#Delta"+coordTitle[cti]+") [cm]")
    h[rmsht].Draw("E")
   h[key].Print(key+'.gif')
  
   if ( prefix=='NEW' ) :
    key='docaPlots'+prefix
    ut.bookCanvas(h,key,key,nx=1200,ny=600,cx=2,cy=1)
    cv = h[key].cd(1)
    h['doca'+prefix].Draw()
    h['doca0'+prefix].Draw("sames")
    cv = h[key].cd(2)
    h['docaRel'+prefix].Draw()
   h[key].Print(key+'.gif')