from ROOT import *

gROOT.SetBatch(kTRUE)

nbins = 20
gROOT.ProcessLine(".x mystyle.C")
#from ShipStyle import *
#lhcbstyleSetup()


# legTuple (title, h, hcolor)
def makeLegend(legTupl, pos=(0.66,0.78,0.90,0.90)):
    leg = TLegend(*pos);
    leg.SetLineColor(kWhite)
    for lt in legTuple:
        leg.AddEntry(lt[0],lt[1],"l");
        
    return leg


#f = TFile("/afs/cern.ch/work/a/austyuzh/public/4barbara/newGen_15340000_nuBg.root")#newGen/ShipAna_Barbara_nu_th45MeV_LSSegments45MeV_FINALTEST.root")
f = TFile("../../DATA/nuFromYandex/newGen_15340000_nuBg.root")
t = f.Get("t")



#h_veto = TH1F("veto","veto",nbins,0,nbins)
#h_effOff = TH1F("effOff","effOff",nbins,0,nbins)
#h_tot = TH1F("tot","tot",nbins,0,nbins)


#flux_nu = 1.09e+12 * 2.e+20/ 4.e+13
#flux_anti = 6.98e+11 * 2.e+20 / 4.e+13
#entries_nu = 1.345e+7#200000-2
#entries_anti = 200000-2

#1.61x10^11,
flux_nu = 1.05e+11 * 2.e+20/ 5.e+13#1.61e+11 * 2.e+20/ 4.e+13
flux_anti = 7.38e+10 * 2.e+20 / 5.e+13#1.02e+11 * 2.e+20 / 4.e+13
entries_nu = 1.534e+7#200000-2
entries_anti =175*10e+3
weight_nu = "(weight * %s / %s)"%(flux_nu,entries_nu)
'''
cutEllipse = "( Sum$(NoB_vtxxSqr/(245.*245.)) + Sum$(NoB_vtxySqr/(490.*490.)) ) < 1."

docaVtxIp = "Sum$(NoB_DOCA < 30.) && Sum$(NoB_IP0 < 250.) && Sum$(NoB_vtxz > -1958.) && Sum$(NoB_vtxz < 2568.)"

goodTracks = "Sum$(DaughtersFitConverged==1) && Sum$(DaughtersAlwaysIn==1)"

trackQ = "Sum$(NoB_MaxDaughtersRedChi2<5.)"

ndf = "Sum$(NoB_MinDaughtersNdf>10.)"

hasMuon = "Sum$(Has1Muon1==1) && Sum$(Has1Muon2==1)"

hasCalo = "Sum$(HasEcal==1)"

vtxz  = "recoed==1"#"Sum$(NoB_vtxz<2500)"
offline_cut = "%s && %s && %s && %s && %s && %s && %s && %s"%(docaVtxIp,goodTracks,trackQ,hasMuon,cutEllipse,hasCalo,ndf,vtxz)
'''

'''
trackQ = "Sum$(NoB_MaxDaughtersRedChi2<5.)"
cutEllipse = "Sum$(( (NoB_vtxxSqr)/(245.*245.) + (NoB_vtxySqr)/(490.*490.) ) < 1.)"
docaVtxIp = "Sum$(NoB_DOCA < 30.) && Sum$(NoB_IP0 < 250.) && Sum$(NoB_vtxz > -1958. && NoB_vtxz < 2568.)"
vtxz  = "fabs(Max$(NoB_vtxz)-Min$(NoB_vtxz))<30"
offline_cut = "%s && %s && %s && %s"%(docaVtxIp,vtxz,trackQ,cutEllipse)
'''

c1 = "Sum$(NoB_MaxDaughtersRedChi2<5.)"
c2 = "Sum$(( (NoB_vtxxSqr)/(245.*245.) + (NoB_vtxySqr)/(490.*490.) ) < 1.)"
c3 = "Sum$(NoB_DOCA < 30.)"
c4 = "Sum$(NoB_IP0 < 250.)" 
c5 = "Sum$(NoB_vtxz > -1958. && NoB_vtxz < 2568.)"
c6  = "TMath::Abs(Max$(NoB_vtxz)-Min$(NoB_vtxz))<30"
offline_cut = "%s && %s && %s && %s && %s && %s"%(c1,c2,c3,c4,c5,c6)

f_dummy = TFile("dummy.root","RECREATE")

cuts = {"not-vetoed":"(nRecoed>0 && numberOfHitLSSegments == 0  && strawVetoAny_eff==0 && upstreamVetoAny_eff==0 && RPCany_eff==0 && TrackSyst==1)",
        "selected" :"%s"%offline_cut,
        "vetoed":"!(nRecoed>0 && numberOfHitLSSegments == 0  && strawVetoAny_eff==0 && upstreamVetoAny_eff==0 && RPCany_eff==0 && TrackSyst==1 && Sum$(NoB_IP0 < 1000.))",
        "not-selected" :"!(%s)"%offline_cut,
        "recoed":"nRecoed>0"}

c = {"not-vetoed":kRed,
     "vetoed":kBlue,
     "selected":kRed,
     'killed':kBlue,
     "reco":kBlack}

n = {'not-vetoed':0,
     "vetoed":0,
     "selected":0,
     'not-selected':0,
     "recoed":0}

h = {}

s = 0
types = ["nChargedPart_fromNu","nNeutrPart_fromNu","nPart_fromNu"]
titles = ["charged", "neutral", ""]
canv = {}
for (ty,tt) in zip(types,titles): 

    h[ty] =  {"not-vetoed":TH1F("%s_not-vetoed"%ty,"not-vetoed",nbins,0,nbins),
              "selected" :TH1F("%s_selected"%ty,"selected",nbins,0,nbins),
              "vetoed":TH1F("%s_vetoed"%ty,"vetoed",nbins,0,nbins),
              "not-selected" :TH1F("%s_not-selected"%ty,"not-selected",nbins,0,nbins),
              "recoed":TH1F("%s_recoed"%ty,"recoed",nbins,0,nbins)}


    for k in n:
        t.Draw("%s>>%s"%(ty,h[ty][k].GetName()),"(%s)*%s"%(cuts[k],weight_nu))
    
    h[ty]['res'] = h[ty]['vetoed'].Clone()
    h[ty]['res'].Divide(h[ty]['recoed'])
    h[ty]['res'].SetLineColor(kBlack)
    h[ty]['res'].SetMarkerColor(kBlack)
    
    yaxis = h[ty]['res'].GetYaxis()
    xaxis = h[ty]['res'].GetXaxis()
    yaxis.SetRangeUser(0.96,1.003)
    yaxis.SetDecimals()
    yaxis.SetTitle("efficiency")
    yaxis.SetTitleOffset(1.1)
    xaxis.SetTitle("%s particles / event"%tt)
    h[ty]['res2'] = h[ty]['not-selected'].Clone()
    h[ty]['res2'].Divide(h[ty]['recoed'])
    h[ty]['res2'].SetLineColor(kBlue)
    h[ty]['res2'].SetMarkerColor(kBlue)
    h[ty]['res2'].SetLineStyle(2)
    
    
    canv[ty] = TCanvas("c_efficiency_%s"%ty,"c_efficiency_%s"%ty)
    h[ty]['res'].Draw('')
    h[ty]['res2'].Draw('same')
    
    legTuple = [(h[ty]['res'],'veto',kBlack),(h[ty]['res2'],'rejection','kBlue')]
    leg = makeLegend(legTuple, pos=(0.30,0.34,0.56,0.47))#(0.34,31,58,43))
    leg.Draw("same")

    canv[ty].Print("eff/"+canv[ty].GetName()+".pdf")
assert(False)



#h['res'].Multiply(h['veto'])
#h['res'].Draw()
#s = h['res'].GetSum()
#print s




'''
s = 0
for i in xrange(nbins):
    for k in n:
        print k, i
        print t.Draw("event>>h_%s_%s"%(k,i),"(%s && nChargedPart_fromNu==%s)*%s"%(cuts[k],i,weight_nu))
        #t.Draw("event","(%s && nChargedPart_fromNu==%s)*%s"%(cuts[k],i,weight_nu))
        #n[k] = float(t.GetHistogram().GetSum())
        print i, k, t.GetHistogram().GetSum()
        print t.GetHistogram().GetSum()
        print t.GetHistogram().GetSum()
        assert(False)
    h_veto.SetBinContent(i+1,n['veto'])
    if n['reco']>0:
        print n["off"]/n["reco"], n['veto']," -> ", n["off"]/n["reco"]*n["veto"]
        #       h_effOff.SetBinContent(i+1,n["off"]/n["reco"])
        #h_tot.SetBinContent(i+1,n["off"]/n["reco"]*n["veto"])
        s += n["off"]/n["reco"]*n["veto"]
#print h_tot.GetSum()
#h_tot.Draw()
print s
'''