import sys,os
import ROOT as r
import numpy as np
import funcsByBarbara as tools

#r.gROOT.SetBatch(r.kTRUE)
r.gROOT.ProcessLine(".X lhcbStyle.C")
plots = {}
pdg = r.TDatabasePDG.Instance()

def addPlot(key, x_axis, tree, string, cuts=None):
    if not cuts:
        tree.Draw(string+'>>h1')
    else:
        tree.Draw(string+'>>h1', cuts)
    plots[key] = r.gDirectory.Get('h1')
    plots[key].GetXaxis().SetTitle(x_axis)
    plots[key].SetTitle(key)
    plots[key].SetName(key)
    c1 = r.TCanvas(); c1.cd()
    if "IP" in key or "Mass" in key or "DOCA" in key or "dChi2" in key:
        c1.SetLogy()
    plots[key].Draw()
    c1.Print('accPlots/%s.pdf'%key)
    c1.Print('accPlots/%s.png'%key)
    key += "-WEIGHTED"
    if 'bg' in key or ('nu' in key and key != 'mumunu'):
        if not cuts:
            tree.Draw(string+'>>h1', "NuWeight")
        else:
            cuts = "NuWeight*("+cuts+")"
            tree.Draw(string+'>>h1', cuts)
        plots[key] = r.gDirectory.Get('h1')
        plots[key].GetXaxis().SetTitle(x_axis)
        plots[key].SetTitle(key)
        plots[key].SetName(key)
        c1 = r.TCanvas(); c1.cd()
        if "IP" in key or "Mass" in key or "DOCA" in key or "dChi2" in key:
            c1.SetLogy()
        plots[key].Draw()
        c1.Print('accPlots/%s.pdf'%key)
        c1.Print('accPlots/%s.png'%key)
    elif 'cosmics' in key:
        if not cuts:
            tree.Draw(string+'>>h1',"HNLw*1.e6/200.")
        else:
            cuts = "HNLw*1.e6/200.*("+cuts+")"
            tree.Draw(string+'>>h1', cuts)
        plots[key] = r.gDirectory.Get('h1')
        plots[key].GetXaxis().SetTitle(x_axis)
        plots[key].SetTitle(key)
        plots[key].SetName(key)
        c1 = r.TCanvas(); c1.cd()
        if "IP" in key or "Mass" in key or "DOCA" in key or "dChi2" in key:
            c1.SetLogy()
        plots[key].Draw()
        c1.Print('accPlots/%s.pdf'%key)
        c1.Print('accPlots/%s.png'%key)
    else:
        if not cuts:
            tree.Draw(string+'>>h1',"HNLw")
        else:
            cuts = "HNLw*("+cuts+")"
            tree.Draw(string+'>>h1', cuts)
        plots[key] = r.gDirectory.Get('h1')
        plots[key].GetXaxis().SetTitle(x_axis)
        plots[key].SetTitle(key)
        plots[key].SetName(key)
        c1 = r.TCanvas(); c1.cd()
        if "IP" in key or "Mass" in key or "DOCA" in key or "dChi2" in key:
            c1.SetLogy()
        plots[key].Draw()
        c1.Print('accPlots/%s.pdf'%key)
        c1.Print('accPlots/%s.png'%key)




fpimu = r.TFile('../DATA/NewPIMU/ShipAna_newGen.root','read')#r.TFile('../DATA/MUPI/ShipAna.root','read')
pimu = fpimu.Get('t')
pimu_geo = tools.searchForNodes3_xyz_dict('../DATA/NewPIMU/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008.root')#'../DATA/MUPI/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008-2985556.root')

"""
fpimu = r.TFile('../DATA/signalFromThomas/ShipAna_newGen.root','read')#r.TFile('../DATA/MUPI/ShipAna.root','read')
pimu = fpimu.Get('t')
pimu_geo = tools.searchForNodes3_xyz_dict('../DATA/signalFromThomas/geofile_full.10.0.Pythia8-TGeant4_50k.root')#'../DATA/MUPI/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008-2985556.root')
"""

fmumunu = r.TFile('../DATA/NewMUMUNU/ShipAna_newGen.root','read')#r.TFile('../DATA/MUMUNU/ShipAna.root','read')
mumunu = fmumunu.Get('t')
mumunu_geo = tools.searchForNodes3_xyz_dict('../DATA/NewMUMUNU/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008-mumunu.root')#'../DATA/MUMUNU/geofile_full.10.0.Pythia8-TGeant4-1.0GeV-008-TOTAL.root')

ftiny = r.TFile('../DATA/SmallHNLs/ShipAna_newGen.root', 'read')
tiny = ftiny.Get('t')
tiny_geo = tools.searchForNodes3_xyz_dict('../DATA/SmallHNLs/geofile_full.10.0.Pythia8-TGeant4-0.1GeV-011.root')

#fcosmics = r.TFile('../DATA/Cosmics/ShipAna.root','read')
#cosmics = fcosmics.Get('ShipAna')
#cosmics_geo = tools.searchForNodes2_xyz_dict('../DATA/Cosmics/geofile_full.Signal.10.0.Cosmics3.25.1-8.root')
#
#fnubar = r.TFile('../DATA/NewKLONG/ShipAna_nubar.root','read')#'../DATA/NewKLONG/ShipAna74-75_nubar.root','read')
#nubar = fnubar.Get('ShipAna')
#nubar_geo = tools.searchForNodes3_xyz_dict('../DATA/NewKLONG/geofile_full.10.0.Genie-TGeant4_nubar.root')#neutrino74-75.root')
#
#fnu = r.TFile('../DATA/NewKLONG/ShipAna_nu.root','read')#'../DATA/NewKLONG/ShipAna77_nu.root','read')
#nu = fnu.Get('ShipAna')
#nu_geo = tools.searchForNodes3_xyz_dict('../DATA/NewKLONG/geofile_full.10.0.Genie-TGeant4_nu.root')#neutrino77.root')

# Neutrinos from yandex
fnu = r.TFile('../DATA/nuFromYandex/newGen_15340000_nuBg.root')
nu = fnu.Get('t')
nu_geo = tools.searchForNodes3_xyz_dict('../DATA/nuFromYandex/geofile_full.10.0.Genie-TGeant4.root')#neutrino77.root')

# For a correct normalisation, we subtract the events that have an HNL vertex produced outside of the 60m fiducial volume

studies = {
    'pimu': {'data':pimu, 'geo':pimu_geo, 'seen':[], 'ntot':20000-125},#136},
    #'pimu': {'data':pimu, 'geo':pimu_geo, 'seen':[], 'ntot':50000-125},#136},
    'mumunu': {'data':mumunu, 'geo':mumunu_geo, 'seen':[], 'ntot':20000-280},#289},
    'tiny': {'data': tiny, 'geo':tiny_geo, 'seen':[], 'ntot':20000-43},#66},
    'nu': {'data': nu, 'geo':nu_geo, 'seen':[], 'ntot':1.534e7},
    #'cosmics': {'data':cosmics, 'geo':cosmics_geo, 'seen':[], 'ntot':cosmics.GetEntriesFast()},
    #'nubar': {'data': nubar, 'geo':nubar_geo, 'seen':[], 'ntot':99999*2},#99999*2},
    #'nu': {'data': nu, 'geo':nu_geo, 'seen':[], 'ntot':99999*2}#99999}
}

#print 'Found %s entries in the cosmics file'%cosmics.GetEntriesFast()

# Cuts:
# - z_vertex in [vol[0]+5m, straw[0]]
# - has track converged
# - chi2/ndof < 10
# - DOCA < 10 cm
# - has no veto?

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
#flux_nu = (1.09e+12 * 2.e+20)/ 5.e+13
#flux_anti = 6.98e+11 * 2.e+20 / 5.e+13
entries_nu = 1.534e7#200000-2
entries_anti =175*10e+3#entries_anti = 200000-2
weight_nu = "(weight * %s / %s)"%(flux_nu,entries_nu)
weight_anti ="(weight * %s / %s)"%(flux_anti,entries_anti)

class MyParticle(object):
    """ Simulating an ntuple filled by particle from one filled by event :) """
    def __init__(self, event, index, noB='', eff=''):
        self.EventNumber = event.__getattr__('event')
        self.DaughtersFitConverged = [event.__getattr__('DaughtersFitConverged')[index], event.__getattr__('DaughtersFitConverged')[index]]
        chi2red = [event.__getattr__('%sMaxDaughtersRedChi2'%noB)[index], event.__getattr__('%sMaxDaughtersRedChi2'%noB)[index]]
        ndf = [event.__getattr__('%sMinDaughtersNdf'%noB)[index], event.__getattr__('%sMinDaughtersNdf'%noB)[index]]
        self.DaughtersChi2 = [chi2red[i]*ndf[i] for i in xrange(len(ndf))]
        self.DaughtersNPoints = ndf
        self.DOCA = event.__getattr__('%sDOCA'%noB)[index]
        self.vtxz = event.__getattr__('%svtxz'%noB)[index]
        try:
            self.vtxy = event.__getattr__('%svtxy'%noB)[index]
            self.vtxx = event.__getattr__('%svtxx'%noB)[index]
        except:
            self.vtxy = r.TMath.Sqrt(event.__getattr__('%svtxySqr'%noB)[index])
            self.vtxx = r.TMath.Sqrt(event.__getattr__('%svtxxSqr'%noB)[index])
        self.IP0 = event.__getattr__('%sIP0'%noB)[index]
        self.BadTruthVtx = event.__getattr__('BadTruthVtx')
        self.NParticles = event.__getattr__('nRecoed')
        self.DaughtersAlwaysIn = event.__getattr__('DaughtersAlwaysIn')[index]
        self.Has1Muon1 = event.__getattr__('Has1Muon1')[index]
        self.Has2Muon1 = event.__getattr__('Has2Muon1')[index]
        self.Has1Muon2 = event.__getattr__('Has1Muon2')[index]
        self.Has2Muon2 = event.__getattr__('Has2Muon2')[index]
        self.HasEcal = event.__getattr__('HasEcal')[index]
        self.HNLw = event.__getattr__('weight')
        self.NuWeight = event.__getattr__('weight')*flux_nu/entries_nu
        self.numberOfHitLSSegments = event.__getattr__('numberOfHitLSSegments')
        self.strawVetoAny = event.__getattr__('strawVetoAny%s'%eff)
        self.upstreamVetoAny = event.__getattr__('upstreamVetoAny%s'%eff)
        self.RPCany = event.__getattr__('RPCany%s'%eff)




class Cuts(object):
    def __init__(self, cutConfigList, study, byEvent=True): 
        self.byEvent = byEvent
        self.cutList = []      
        try:        
            for (index,cutConfig) in enumerate(cutConfigList):
                self.cutList.append( self.Cut( self, index, cutConfig['name'], study, cutConfig['params'] ) )               
        except BaseException,e:
            print "cutConfigList badly set: %r" %e
    def process(self, eventList, noB, eff=''):
        if self.byEvent:
            for event in eventList:
                for index in xrange(event.nRecoed):
                    for cut in self.cutList:
                        if not cut.count(MyParticle(event, index, noB, eff)): break
        else:
            for event in eventList:
                for cut in self.cutList:
                    if  not cut.count(event): break
        return self
    def select(self, event, imax):
        for i in xrange(imax):
            if not self.cutList[i].selection(event, self.cutList[i].study, self.cutList[i].params):
                return False
        return True

    class Cut(object):
        def __init__(self, cutCollection, index, cutName, study, params):
            self.cutCollection = cutCollection
            self.index = index
            self.name = cutName
            self.params = params
            self.study = study
            self.total = 0 
            self.weight = 0
            self.eventNumbers = []
            self.duplicates = []
            self.selection = self.__get(cutName)
        def __get(self, cutName):        
            name = "cut_"+cutName
            if hasattr(self,name):
                try:
                    return getattr(self,name)
                except AttributeError,e:
                    return None
            else:
                raise BaseException ( "bad cut name: %r" %cutName )
                return None
        def count(self, event):
            try:
                if (self.selection(event,self.study,self.params)):
                    self.total += 1
                    if (event.EventNumber not in self.eventNumbers):
                        self.weight += self.compute_weight(event, self.study)
                        self.eventNumbers.append(event.EventNumber)
                    else:
                        self.duplicates.append(event.EventNumber)
                    return True
                else :
                    return False
            except BaseException,e:
                print "Error performing %r: %r" %(self.name,e)
        def cut_fit_converged(self, event, study, params):
            return fit_converged(event)
        def cut_chi2(self, event, study, params):
            return chi2_cut(event,params['chi2max'])
        def cut_has_muons(self, event, study, params):
            for station in params['stations']:
                if not has_muon_station(event, study, station):
                    return False
            return True
        def cut_doca(self, event, study, params):
            return doca_cut(event, params['docamax'])
        def cut_z(self, event, study, params):
            return z_cut(event, study)
        def cut_normalization(self, event, study, params):
            return goodTruthVtx(event, study)
        def cut_good_tracks(self, event, study, params):
            return goodTrack(event)
        def cut_ip(self, event, study, params):
            return ip_cut(event, params['ipmax'])
        def cut_vtx_in_ellipse(self, event, study, params):
            return vtxInAcceptance(event, study)
        def cut_vtxSq_in_ellipse(self, event, study, params):
            return vtxSqInAcceptance(event, study)
        def cut_n_candidates(self, event, study, params):
            #print self.cutCollection.cutList[self.index-1].name
            if (event.EventNumber in self.cutCollection.cutList[self.index-1].duplicates):
                #if self.cutCollection.select(event, self.index):
                return nParticles_cut(event, params['nmax'])
            return True
        def cut_has_ecal(self, event, study, params):
            return has_ecalDeposit(event)
        def cut_has_n_muons_in(self, event, study, params):
            return has_n_muons_in_station(event, params['n'], params['station'])
        def cut_not_vetoed(self, event, study, params):
            return was_not_vetoed(event, study)
        def compute_weight(self, event, study):
            return weight(event, study)


def was_not_vetoed(event, study):
    if study in ['pimu', 'mumunu', 'tiny']:
        return True
    flag = (event.numberOfHitLSSegments == 0) and (event.strawVetoAny==0) and (event.upstreamVetoAny==0) and (event.RPCany==0)
    return flag

def fit_converged(event):
    flags = event.DaughtersFitConverged
    if np.product(flags):
        return True
    return False

def chi2_cut(event, chi2max):
    chi2 = event.DaughtersChi2
    n = event.DaughtersNPoints
    for i in xrange(len(n)):
        if not n[i]:
            print event.EventNumber
            return False
        if chi2[i]/n[i] > chi2max:
            return False
    return True

def has_muon_station(event, study, station):
    zIn = studies[study]['geo']['muondet%s_1'%(station-1)]['z']['pos'] - studies[study]['geo']['muondet%s_1'%(station-1)]['z']['dim']
    zOut = studies[study]['geo']['muondet%s_1'%(station-1)]['z']['pos'] + studies[study]['geo']['muondet%s_1'%(station-1)]['z']['dim']
    for z in event.muon_z:
        if zIn <= z <= zOut:
            return True
    return False

def has_n_muons_in_station(event, n, station):
    if n>2:
        print "has_n_muons_in_station ERROR: we have maximum 2 tracks!!"
        sys.exit()
    if station>2:
        print "has_n_muons_in_station ERROR: we looked only at the first and second stations!!"
        sys.exit()
    if n==1 and station==1:
        return event.Has1Muon1
    if n==2 and station==1:
        return event.Has2Muon1
    if n==1 and station==2:
        return event.Has1Muon2
    if n==2 and station==2:
        return event.Has2Muon2

def has_ecalDeposit(event):
    return event.HasEcal


def doca_cut(event, doca_max):
    if event.DOCA < doca_max:
        return True
    return False

def z_cut(event, study):
    z = event.vtxz
    #try:
    #   minz = studies[study]['geo']['lidT1I_1']['z']['pos'] + studies[study]['geo']['lidT1I_1']['z']['dim'] + 500.
    #   maxz = studies[study]['geo']['lidT6I_1']['z']['pos'] - studies[study]['geo']['lidT6I_1']['z']['dim']
    #except:
    minz = studies[study]['geo']['Veto_5']['z']['pos'] + studies[study]['geo']['Veto_5']['z']['dim'] + 0.01
    maxz = studies[study]['geo']['Tr1_1']['z']['pos'] - studies[study]['geo']['Tr1_1']['z']['dim'] - 0.01
    if minz < z < maxz:
        return True
    return False

def goodTruthVtx(event, study):
    if study in ['pimu', 'mumunu', 'tiny']:
        if not event.BadTruthVtx:
            return True
    return True

def nParticles_cut(event, nmax):
    if event.NParticles > nmax:
        return False
    return True

# This checks that:
# - there is at least one tracking station before and one after the magnet
# - hits in these tracking stations are within a 5x10m ellipse
def goodTrack(event):
    if event.DaughtersAlwaysIn:
        return True
    return False

def ip_cut(event, maxip):
    if event.IP0 < maxip:
        return True
    return False

def nothing_in_veto5(event):
    if len(event.veto5_x) > 0:
        return False
    return True

def nothing_in_liquidscint(event):
    if len(event.liquidscint_x) > 0:
        return False
    return True

def something_in_muon(event):
    if len(event.muon_x) > 0:
        return True
    return False

def something_in_ecal(event):
    if len(event.ecal_x) > 0:
        return True
    return False

def vtxInAcceptance(event, study):
    x = event.vtxx
    y = event.vtxy
    #if 'cosmics' in study or 'bg' in study:
    #   print study, x, y
    return tools.inEllipse(x,y,500./2,1000./2)

def vtxSqInAcceptance(event, study):
    xSq = event.vtxxSqr
    ySq = event.vtxySqr
    #if 'cosmics' in study or 'bg' in study:
    #   print study, x, y
    return tools.inEllipse(r.TMath.Sqrt(xSq), r.TMath.Sqrt(ySq), 500./2, 1000./2)

def weight(event, study):
    #if event.EventNumber not in studies[study]['seen']:
        #studies[study]['seen'].append(event.EventNumber)
    if ('bg' in study) or (study == 'nu' or study == 'nubar'):
        w = event.NuWeight
    elif 'cosmics' in study:
        w = event.HNLw * 1.e6 / 200.
    elif study in ['pimu', 'mumunu', 'tiny']:
        w = event.HNLw / studies[study]['ntot']
    else:
        print 'WEIGHTING ERROR: ', study
        sys.exit()
    #else:
    #   w = 0.
        #if study in ['pimu', 'mumunu', 'tiny']:
        #   print study, event.EventNumber, event.NParticles
    return w

def computeResidualWeight(filename=None):
    if not filename:
        filename = "forElena_nu_notVetoed.txt"
    fScan = file(filename)
    lines = fScan.readlines()
    sample = nu
    s = 'nu'
    if ('bar' in filename) or ('anti' in filename):
        sample = nubar
        s = 'nubar'
    lines = lines[3:-1]
    events = []
    for line in lines:
        events.append(int(line.rsplit("*")[1].replace(" ","")))
    count_reco = 0
    count = 0
    weighted_count = 0.
    print "unvetoed: ", len(events)
    print "total Reconstructed: ", sample.GetEntriesFast()
    for event in sample:
        if event.EventNumber in events:
            count_reco +=1
            if (fit_converged(event) and chi2_cut(event, 5.) and z_cut(event, s) and vtxInAcceptance(event, s)
                and goodTrack(event) and doca_cut(event, 30) and has_muon_station(event, s, 1)
                and has_muon_station(event, s, 2) and ip_cut(event, 250.)):
            #if True:#fit_converged(event): #and doca_cut(event, 50.) and ip_cut(event, 250.) and z_cut(event, 'nu') and vtxInAcceptance(event, 'nu'):
                #print event.EventNumber, 'mothers: ', event.DaughtersTruthMotherPDG[0], event.DaughtersTruthMotherPDG[1]
                print event.EventNumber
                print 'mothers: ', event.DaughtersTruthMotherPDG[0], event.DaughtersTruthMotherPDG[1]
                #print event.DOCA, event.IP0, event.DaughtersChi2[0]/event.DaughtersNPoints[0], event.DaughtersChi2[1]/event.DaughtersNPoints[1], event.DaughtersTruthPDG[0], event.DaughtersTruthPDG[1] 
                print "tracks: ", event.DaughtersTruthPDG[0], event.DaughtersTruthPDG[1] 
                print "vtx: ", event.vtxx, event.vtxy, event.vtxz
                print 'DOCA, IP: ', event.DOCA, event.IP0
                #print tools.interactionElement(event.DaughtersTruthProdX[0], event.DaughtersTruthProdY[0], event.DaughtersTruthProdZ[0], nu_geo), tools.interactionElement(event.DaughtersTruthProdX[1], event.DaughtersTruthProdY[1], event.DaughtersTruthProdZ[1],nu_geo)
                count +=1
                weighted_count += event.NuWeight
    print "Reconstructed: %s"%count_reco
    return count, weighted_count

def latex_float(f):
    float_str = "{0:.3g}".format(f)
    if "e" in float_str:
        base, exponent = float_str.split("e")
        return "${0} \\times 10^{{{1}}}$".format(base, int(exponent))
    else:
        return float_str


class cutsWithDraw(object):
    def __init__(self):
        self.noB = 'NoB_'#''
        self.opt = '' # '_eff' # 'Ethr'
        self.nu = False
    def setNoB(self, noB):
        # to ignore or take into account Katerina's modifications
        self.noB = noB
    def recoed(self):
        cut = " recoed==1 && Sum$(BadTruthVtx==0) "
        if self.nu:
            cut = " recoed==1 "
        string = "Event reconstructed"
        return cut, string
    def nRecoed(self, num):
        cut = " nRecoed==%s "%num
        if num == 1:
            strn = "1 HNL candidate"
        elif num>1:
            strn = "%s HNL candidates"%num
        string = strn
        return cut, string
    def doca(self, th, noB=None):
        if not noB:
            noB = self.noB
        cut = " Sum$(%sDOCA < %s) "%(noB, th)
        if th == int(round(th)):
            th = int(round(th))
        string = "DOCA $<$ %s~cm"%th
        return cut, string
    def ip(self, th, noB=None):
        if not noB:
            noB = self.noB
        cut =  " Sum$(%sIP0 < %s) "%(noB, th)
        string = "IP $<$ %.1f m"%(th/100.)
        return cut, string
    def z(self, zmin, zmax, noB=None):
        if not noB:
            noB = self.noB
        cut = " Sum$(%svtxz > %s && %svtxz < %s) "%(noB, zmin, noB, zmax)
        string = "$z_{straw\, veto} < z_{vtx} < z_{trk1}$"
        return cut, string
    def ellipse(self, a, b, noB=None):
        if not noB:
            noB = self.noB
        cut = " Sum$(( (%svtxx*%svtxx)/(%s*%s) + (%svtxy*%svtxy)/(%s*%s) ) < 1.) "%(noB,noB, a,a, noB,noB, b,b)
        string = "Vertex $\\in $ ellipse"
        return cut, string
    def fiducial(self, zmin, zmax, a, b, noB=None):
        cutz = self.z(zmin, zmax, noB)[0]
        cute = self.ellipseSq(a, b, noB)[0]
        string = "Vtx in fiducial vol."
        cut = cutz+"&&"+cute
        return cut, string
    def ellipseSq(self, a, b, noB=None):
        if not noB:
            noB = self.noB
        cut = " Sum$(( (%svtxxSqr)/(%s*%s) + (%svtxySqr)/(%s*%s) ) < 1.) "%(noB, a,a, noB, b,b)
        string = "Vertex $\\in $ ellipse"
        return cut, string
    def converged(self):
        cut = " Sum$(DaughtersFitConverged==1) "
        string = "Track fit converged"
        return cut, string
    def goodtracks(self):
        cut = " Sum$(DaughtersAlwaysIn==1) "
        string = "Tracks in fiducial vol."# $\\in$ ellipse"
        return cut, string
    def redChi2(self, th, noB=None):
        if not noB:
            noB = self.noB
        cut = " Sum$(%sMaxDaughtersRedChi2<%s) "%(noB, th)
        if th == int(round(th)):
            th = int(round(th))
        string = "$\\chi^2/\\text{N.d.f.} < %s $"%th
        return cut, string
    def ndf(self, th, noB=None):
        if not noB:
            noB = self.noB
        cut = " Sum$(%sMinDaughtersNdf>%s) "%(noB, th)
        string = "$\\text{N.d.f.} > %s$"%th
        return cut, string
    def muon1(self, num):
        cut = " Sum$(Has%sMuon1==1) "%num
        if num==1:
            strn = "1 muon"
        elif num==2:
            strn = "2 muons"
        string = strn+" in 1$^{\\text{st}}$ muon station"
        return cut, string
    def muon2(self, num):
        cut = " Sum$(Has%sMuon2==1) "%num
        if num==1:
            strn = "1 muon"
        elif num==2:
            strn = "2 muons"
        string = strn+" in 2$^{\\text{nd}}$ muon station"
        return cut, string
    def ecal(self):
        cut = " Sum$(HasEcal==1) "
        string = "150~MeV in Ecal"
        return cut, string
    def notVetoed(self, th, noB=None, opt=None):
        if not noB:
            noB = self.noB
        if not opt:
            opt = self.opt
        cut = " Sum$(%sIP0 < %s) &&"%(noB, th)+" Sum$(numberOfHitLSSegments == 0) && Sum$(strawVetoAny%s==0) && Sum$(upstreamVetoAny%s==0) && Sum$(RPCany%s==0) "%(opt, opt, opt)
        string = "Event not vetoed"
        return cut, string



def countWithDraw(study):
    t = studies[study]['data']
    geo = studies[study]['geo']
    ntot = studies[study]['ntot']
    tc = cutsWithDraw()
    if study == 'nu':
        tc.setNoB('NoB_')
        tc.nu = True
    zmin = geo['Veto_5']['z']['pos']+geo['Veto_5']['z']['dim']
    zmax = geo['Tr1_1']['z']['pos']-geo['Tr1_1']['z']['dim']
    if (study == 'pimu'):
        cuts =     [ tc.recoed()[0], tc.notVetoed(1000.)[0], tc.redChi2(5.)[0], tc.ndf(25)[0], tc.fiducial(zmin, zmax, 250., 500.)[0], tc.goodtracks()[0], 
                     tc.ecal()[0], tc.muon1(1)[0], tc.muon2(1)[0], tc.doca(30.)[0], tc.ip(250.)[0] ]
        cutnames = [ tc.recoed()[1], tc.notVetoed(1000.)[1], tc.redChi2(5.)[1], tc.ndf(25)[1], tc.fiducial(zmin, zmax, 250., 500.)[1], tc.goodtracks()[1], 
                     tc.ecal()[1], tc.muon1(1)[1], tc.muon2(1)[1], tc.doca(30.)[1], tc.ip(250.)[1] ]
    elif (study == 'nu'):
        cuts =     [ tc.recoed()[0], tc.notVetoed(1000.)[0], tc.redChi2(5.)[0], tc.ndf(25)[0], tc.fiducial(zmin, zmax, 250., 500.)[0], tc.goodtracks()[0], 
                     tc.ecal()[0], tc.muon1(1)[0], tc.muon2(1)[0], tc.doca(30.)[0], tc.ip(250.)[0] ]
        cutnames = [ tc.recoed()[1], tc.notVetoed(1000.)[1], tc.redChi2(5.)[1], tc.ndf(25)[1], tc.fiducial(zmin, zmax, 250., 500.)[1], tc.goodtracks()[1], 
                     tc.ecal()[1], tc.muon1(1)[1], tc.muon2(1)[1], tc.doca(30.)[1], tc.ip(250.)[1] ]
    elif (study == 'mumunu'):
        cuts =     [ tc.recoed()[0], tc.notVetoed(1000.)[0], tc.redChi2(5.)[0], tc.ndf(25)[0], tc.fiducial(zmin, zmax, 250., 500.)[0], tc.goodtracks()[0], 
                     tc.muon1(2)[0], tc.muon2(2)[0], tc.doca(30.)[0], tc.ip(250.)[0] ]
        cutnames = [ tc.recoed()[1], tc.notVetoed(1000.)[1], tc.redChi2(5.)[1], tc.ndf(25)[1], tc.fiducial(zmin, zmax, 250., 500.)[1], tc.goodtracks()[1], 
                     tc.muon1(2)[1], tc.muon2(2)[1], tc.doca(30.)[1], tc.ip(250.)[1] ]
    elif (study == 'tiny'):
        cuts =     [ tc.recoed()[0], tc.notVetoed(1000.)[0], tc.redChi2(5.)[0], tc.ndf(25)[0], tc.fiducial(zmin, zmax, 250., 500.)[0], tc.goodtracks()[0], 
                     tc.ecal()[0], tc.doca(30.)[0], tc.ip(250.)[0] ]
        cutnames = [ tc.recoed()[1], tc.notVetoed(1000.)[1], tc.redChi2(5.)[1], tc.ndf(25)[1], tc.fiducial(zmin, zmax, 250., 500.)[1], tc.goodtracks()[1], 
                     tc.ecal()[1], tc.doca(30.)[1], tc.ip(250.)[1] ]
    # if not cutnames:
    #    cutnames = [ "converged", "not vetoed", "chi2/n<5", "ndf>15", "z", "ellipse", "good tracks", "ecal", "muon1", "muon2", "doca<30", "IP<250" ]
    # cutStrings = [  "Track fit converged",
    #                 "Event not vetoed",
    #                 "$\\xi^2/ndf < 5 $",
    #                 "$ndf > 15$",
    #                 "$z_{straw veto} < z_{vtx} < z_{Trk1}$",
    #                 "Vertex $\\in $ ellipse",
    #                 "Tracks $\\in$ ellipse",
    #                 "150~MeV in Ecal",
    #                 "1$^{\\text{st}}$ muon station",
    #                 "2$^{\\text{nd}}$ muon station",
    #                 "DOCA $<$ 30~cm",
    #                 "IP $<$ 2.5~m",
    #                 ]
    entries = []
    weights = []
    print
    r.gROOT.SetBatch(r.kTRUE)
    for i,cut in enumerate(cuts):
        n = t.Draw("event", "&&".join([c for c in cuts[:i+1]]))
        #if study == 'nu':
        #    t.Draw("1.>>hw","(weight*%s/%s)*("%(flux_nu, entries_nu) + str("&&".join([c for c in cuts[:i+1]])) + ")")
        #elif study == 'nubar':
        #    t.Draw("1.>>hw","(weight*%s/%s)*("%(flux_anti, entries_anti) + str("&&".join([c for c in cuts[:i+1]])) + ")")
        #else:
        #    t.Draw("1.>>hw","(weight/%s)*("%studies[study]['ntot'] + str("&&".join([c for c in cuts[:i+1]])) + ")")
        if study == 'nu':
            t.Draw("1.>>hw","(weight*%s/%s)*("%(flux_nu, entries_nu) + str("&&".join([c for c in cuts[:i+1]])) + ")")
        elif study == 'nubar':
            t.Draw("1.>>hw","(weight*%s/%s)*("%(flux_anti, entries_anti) + str("&&".join([c for c in cuts[:i+1]])) + ")")
        else:
            t.Draw("1.>>hw","(weight/%s)*("%ntot + str("&&".join([c for c in cuts[:i+1]])) + ")")
        w = t.GetHistogram().GetSumOfWeights()#Integral()#Integral('width')#GetSum()
        #w = r.gDirectory.Get('hw').Integral()
        entries.append(n)
        weights.append(w)
        print cutnames[i], '\t\t\t', n, '\t', w, cuts[i]
    print
    r.gROOT.SetBatch(r.kFALSE)
    return entries, weights, cutnames

def tableWithDraw():
    results = {}
    for study in ['nu', 'pimu', 'mumunu', 'tiny']:
        results[study] = {}
        numbers, numbersWithWeights, cuts = countWithDraw(study)
        results[study]['numbers'] = numbers
        results[study]['numbersWithWeights'] = numbersWithWeights
        results[study]['cuts'] = cuts
    return results


def plotSigBg(cSaver, name, sig_tree, sig_ntot, bg_tree, what, cuts, xaxis, legSize, log, sig_cuts=None, bg_cuts=None):
    if not sig_cuts:
        sig_cuts = cuts
    if not bg_cuts:
        bg_cuts = cuts.replace("&& Sum$(BadTruthVtx==0) ", "")
    # colors
    bgcolor = r.kRed
    sigcolor = r.kBlack
    sigfill = r.kGray
    bgstyle = 1#9
    r.gStyle.SetOptTitle(r.kFALSE)
    # draw normalized
    c1 = r.TCanvas(name, name)
    whatbg = what
    #if what in ["DOCA", "vtxz"]:
    #    whatbg = "NoB_"+what
    if "z" in name:
        sig_tree.Draw(what+'>>h_sig_%s(100,-15500.,4500.)'%name, "(weight/%s)*("%sig_ntot + sig_cuts + ")", 'norm')
    elif "pimu-IP0-lin" in name:
        sig_tree.Draw(what+'>>h_sig_%s(100,0.,20.)'%name, "(weight/%s)*("%sig_ntot + sig_cuts + ")", 'norm')
    elif "IP0-lin" in name:
        sig_tree.Draw(what+'>>h_sig_%s(100,0.,250.)'%name, "(weight/%s)*("%sig_ntot + sig_cuts + ")", 'norm')
    else:
        sig_tree.Draw(what+'>>h_sig_%s'%name, "(weight/%s)*("%sig_ntot + sig_cuts + ")", 'norm')
    hsig = r.gDirectory.Get('h_sig_%s'%name)
    hsig.SetMarkerColor(sigcolor)
    hsig.SetLineColor(sigcolor)
    hsig.SetFillColor(sigfill)
    hsig.GetXaxis().SetTitle(xaxis)
    hsig.GetYaxis().SetTitle("a. u.")
    hsig.GetYaxis().SetDecimals()
    if (("mumunu-mass" in name) or ("mumunu-IP0" in name) or ("tiny-mass" in name) or ("tiny-IP0" in name)) and not ("lin" in name):
        hsig.GetYaxis().SetRangeUser(0.0001,5.)
    #if "z" in name:
    #    print name
    #    hsig.GetXaxis().SetLimits(-17000.,4500.)
    #    hsig.Draw()
    hbg = None
    if not "no-bg-mass" in name:
        bg_tree.Draw(whatbg+'>>h_bg_%s'%name, "(weight*%s/%s)*("%(flux_nu, entries_nu) + bg_cuts + ")", 'norm same')
        hbg = r.gDirectory.Get('h_bg_%s'%name)
        hbg.SetMarkerColor(bgcolor)
        hbg.SetLineColor(bgcolor)
        hbg.SetLineStyle(bgstyle)
    if "y" in log:
        c1.SetLogy()
    if "x" in log:
        c1.SetLogx()
    newLegSize = list(legSize)
    if ("tiny" in name) and not ("2" in name):
        newLegSize[2] += 0.02
        newLegSize[3] += 0.02
    leg = r.TLegend(*newLegSize)
    leg.SetTextFont(42)
    leg.SetLineColor(r.kWhite)
    leg.SetBorderSize(0)
    sigstring = "HNL #rightarrow #pi#mu"
    if ('eenu' in name) or ('tiny' in name):
        sigstring = sigstring.replace("#pi#mu", "ee#nu")
    if ('mumu' in name):
        sigstring = sigstring.replace("#pi#mu", "#mu#mu#nu")
    leg.AddEntry(hsig, sigstring, 'lf' )
    if not "no-bg-mass" in name:
        leg.AddEntry(hbg, "Neutrino BG", 'l' )
    leg.Draw("same")
    c1.Modified(); c1.Update()
    print "\n", name, hsig.GetMean(), hsig.GetRMS()
    cSaver[name] = (c1, leg, hsig, hbg)
    return c1, leg, hsig, hbg

cSaver = []
def accPlotsTP(study):
    if not os.path.exists('accPlotsTP'):
        os.system('mkdir accPlotsTP')
    r.gROOT.ProcessLine(".x newGen/mystyle.C")
    t = studies[study]['data']
    geo = studies[study]['geo']
    ntot = studies[study]['ntot']
    tc = cutsWithDraw()
    zmin = geo['Veto_5']['z']['pos']+geo['Veto_5']['z']['dim']
    zmax = geo['Tr1_1']['z']['pos']-geo['Tr1_1']['z']['dim']
    # BG stuff
    bg_t = studies['nu']['data']
    bg_geo = studies['nu']['geo']
    bg_ntot = studies['nu']['ntot']
    bg_tc = cutsWithDraw()
    bg_tc.setNoB('NoB_')
    bg_zmin = geo['Veto_5']['z']['pos']+geo['Veto_5']['z']['dim']
    bg_zmax = geo['Tr1_1']['z']['pos']-geo['Tr1_1']['z']['dim']
    plot_ip = plotSigBg(study+"-IP0", t, ntot, bg_t, "IP0", tc.notVetoed(1000.)[0], "Impact parameter to target [cm]")[0]
    cSaver.append(plot_ip)
    raw_input('')




# per particle
cutConfigList = [   {'name': 'normalization', 'params':{}},
                    {'name': 'fit_converged', 'params':{}},
                    {'name': 'chi2',          'params':{'chi2max':5.}},
                    {'name': 'z',             'params':{}},
                    {'name': 'vtx_in_ellipse','params':{}},
                    {'name': 'good_tracks',   'params':{}},
                    {'name': 'doca',          'params':{'docamax':50.}},
                    {'name': 'ip',            'params':{'ipmax':500.}},
                    {'name': 'has_muons',     'params':{'stations':[1,2]}},
                    {'name': 'doca',          'params':{'docamax':30.}},
                    {'name': 'ip',            'params':{'ipmax':250.}},
                    {'name': 'n_candidates',  'params':{'nmax':1}},
]

# new analysis sw - per event
cutConfigList = [   {'name': 'normalization',    'params':{}},
                    {'name': 'fit_converged',    'params':{}},
                    {'name': 'chi2',             'params':{'chi2max':5.}},
                    {'name': 'z',                'params':{}},
                    {'name': 'vtx_in_ellipse',   'params':{}},
                    {'name': 'good_tracks',      'params':{}},
                    {'name': 'doca',             'params':{'docamax':50.}},
                    {'name': 'ip',               'params':{'ipmax':500.}},
                    {'name': 'has_n_muons_in',   'params':{'n':1,'station':1}},
                    {'name': 'has_n_muons_in',   'params':{'n':1,'station':2}},
                    {'name': 'has_ecal',         'params':{}},
                    {'name': 'doca',             'params':{'docamax':30.}},
                    {'name': 'ip',               'params':{'ipmax':250.}},
]

cutConfigDict = { 'sig': [{'name': 'normalization',    'params':{}},
                          {'name': 'fit_converged',    'params':{}},
                          {'name': 'chi2',             'params':{'chi2max':5.}},
                          {'name': 'z',                'params':{}},
                          {'name': 'vtx_in_ellipse',   'params':{}},
                          {'name': 'good_tracks',      'params':{}},
                          {'name': 'doca',             'params':{'docamax':50.}},
                          {'name': 'ip',               'params':{'ipmax':500.}},
                          {'name': 'has_n_muons_in',   'params':{'n':1,'station':1}},
                          {'name': 'has_n_muons_in',   'params':{'n':1,'station':2}},
                          {'name': 'has_ecal',         'params':{}},
                          {'name': 'doca',             'params':{'docamax':30.}},
                          {'name': 'ip',               'params':{'ipmax':250.}}],

                  'bg':  [{'name': 'not_vetoed',       'params':{}},
                          {'name': 'fit_converged',    'params':{}},
                          {'name': 'chi2',             'params':{'chi2max':5.}},
                          {'name': 'z',                'params':{}},
                          {'name': 'vtx_in_ellipse',   'params':{}},
                          {'name': 'good_tracks',      'params':{}},
                          {'name': 'doca',             'params':{'docamax':50.}},
                          {'name': 'ip',               'params':{'ipmax':500.}},
                          {'name': 'has_n_muons_in',   'params':{'n':1,'station':1}},
                          {'name': 'has_n_muons_in',   'params':{'n':1,'station':2}},
                          {'name': 'has_ecal',         'params':{}},
                          {'name': 'doca',             'params':{'docamax':30.}},
                          {'name': 'ip',               'params':{'ipmax':250.}}]

}


def numbers_v2():
    mycuts, results = {}, {}
    for study in studies:
        print 'Reading %s...'%study
        if study in ['tiny', 'pimu', 'mumunu']:
            cutConfigList = cutConfigDict['sig']
            noB = ''
        else:
            cutConfigList = cutConfigDict['bg']
            noB = 'NoB_'
        mycuts[study] = Cuts(cutConfigList, study).process(studies[study]['data'], noB)
        for cut in mycuts[study].cutList:
            assert( len(cut.eventNumbers) == len(list(set(cut.eventNumbers))) )
            print '%s\t%s\t%s\t%s'%(cut.name, cut.total, len(cut.eventNumbers), cut.weight)
        results[study] = {}
        results[study]['numbers']=[ cut.total for cut in mycuts[study].cutList[1:] ]
        results[study]['numbersWithWeights']=[ cut.weight for cut in mycuts[study].cutList[1:] ]
        print
        #print 'Events with more than one particle:'
        #for entry in mycuts[study].cutList[-2].duplicates: #[ x for x in mycuts[study].cutList[-1].eventNumbers if mycuts[study].cutList[-1].eventNumbers.count(x) > 1 ]:
        #    for tr in studies[study]['data']:
        #        if tr.EventNumber == entry:
        #            if mycuts[study].select(tr, len(mycuts[study].cutList)-1):
        #                print tr.EventNumber, tr.vtxz, tr.IP0, [pdg.GetParticle(code).GetName() for code in tr.DaughtersTruthPDG]
        #print
        print
    return results

header = {  'pimu': ("$HNL \\to \\pi \\mu$", "tab:recpimu"),
            'mumunu': ("$HNL \\to \\mu \\mu \\nu$", "tab:recmmnu"),
            'tiny': ("$HNL \\to e e \\nu$", "tab:receenu"),
            'nu': ("neutrino-induced background", "tab:recnubkg")
         }

def makeTex(results):
    for study in results.keys():
        weighted_string = "Acceptance"
        if study == 'nu':
            weighted_string = "Events / 5 years"
        results[study]['efficiencies'] = [ results[study]['numbersWithWeights'][i]/results[study]['numbersWithWeights'][i-1] for i in xrange(1, len(results[study]['numbersWithWeights']))]
        results[study]['effStrings'] = ['-'] + [ "{:.1f} \\%".format(float(e) * 100) for e in results[study]['efficiencies'] ]
        #for s in results[study]['effStrings']:
        #    print s
        print '\n\n\n'
        print '\t\\begin{table}[!!!!h]'
        print '\t\\begin{center}'
        print '\t\t\\caption{ Effect of the offline selection on %s \\label{%s}}'%(header[study][0],header[study][1])
        print '\t\t\\begin{tabular}{|l|c|c|c|}'
        print '\t\t\t\\hline'
        print '\t\t\t\\textbf{Selection} & \\textbf{Entries} & \\textbf{%s} & \\textbf{Selection efficiency} \\\\'%weighted_string
        print '\t\t\t\\hline'
        assert(len(results[study]['cuts']) == len(results[study]['numbers']) == len(results[study]['numbersWithWeights']))
        for i in xrange(len(results[study]['cuts'])):
            print '\t\t\t', results[study]['cuts'][i], ' & ', '%s'%int(results[study]['numbers'][i]), ' & ', '%s'%latex_float(results[study]['numbersWithWeights'][i]), ' & ', results[study]['effStrings'][i], '\\\\'
            print '\t\t\t\\hline'
        print '\t\t\\end{tabular}'
        print '\t\\end{center}'
        print '\t\\end{table}'
        print '\n\n'


def makeTexOld(results, cuts=None):
    print
    print 'Background:'
    print '\t\\begin{table}'
    print '\t\\centering Backgrounds: \\tiny'
    print '\t\t\\begin{tabular}{|l|c|c|c|c|c|c|}'
    print '\t\t\t\\hline'
    print '\t\t\tCut & Weighted $\\nu$ & $\\nu$ entries & Weighted $\\bar{\\nu}$ & $\\bar{\\nu}$ entries & Weighted cosmics & cosmics entries \\\\'
    print '\t\t\t\\hline'
    if not cuts:
        cuts = ['Track fit converged', '$\chi^2/n < 5$', '$z_{in}+5~\\text{m} < z_{vtx} < z_{out}$', '$z_{vtx} \in $ ellipse', 'Tracks $\in$ ellipse', 'DOCA $<$ 50 cm',
                'IP $<$ 5 m', 'Muon1 \& Muon2', 'DOCA $<$ 30 cm', 'IP $<$ 2.5 m', '$N_{candidates} = 1$']
    for i in xrange(len(results['nu']['numbers'])):
        print   '\t\t\t', cuts[i], ' & ', '%s'%latex_float(results['nu']['numbersWithWeights'][i]), ' & ', '%s'%results['nu']['numbers'][i], ' & ',\
                '%s'%latex_float(results['nubar']['numbersWithWeights'][i]), ' & ', '%s'%results['nubar']['numbers'][i], ' & ',\
                '%s'%latex_float(results['cosmics']['numbersWithWeights'][i]), ' & ', '%s'%results['cosmics']['numbers'][i], '\\\\ '
        print '\t\t\t\\hline'
    print '\t\t\\end{tabular}'
    print '\t\\end{table}'
    print
    print 'Signal:'
    print '\t\\begin{table}'
    print '\t\\centering Signal: \\tiny'
    print '\t\t\\begin{tabular}{|l|c|c|c|c|c|c|}'
    print '\t\t\t\\hline'
    print '\t\t\tCut & $\\mathcal{A}(\\pi\\mu)$ & $\\pi\\mu$ entries & $\\mathcal{A}(\\mu\\mu\\nu)$ & $\\mu\\mu\\nu$ entries & $\\mathcal{A}(ee\\nu)$ & $ee\\nu$ entries \\\\'
    print '\t\t\t\\hline'
    #cuts = ['Track fit converged', '$\chi^2/n < 5$', '$z_{in}+5~\\text{m} < z_{vtx} < z_{out}$', '$z_{vtx} \in $ ellipse', 'Tracks $\in$ ellipse', 'DOCA $<$ 50 cm',
    #        'IP $<$ 5 m', 'Muon1 \& Muon2', 'DOCA $<$ 30 cm', 'IP $<$ 2.5 m', '$N_{candidates} = 1$']
    for i in xrange(len(results['nu']['numbers'])):
        print   '\t\t\t', cuts[i], ' & ', '%s'%latex_float(results['pimu']['numbersWithWeights'][i]), ' & ', '%s'%results['pimu']['numbers'][i], ' & ',\
                '%s'%latex_float(results['mumunu']['numbersWithWeights'][i]), ' & ', '%s'%results['mumunu']['numbers'][i], ' & ',\
                '%s'%latex_float(results['tiny']['numbersWithWeights'][i]), ' & ', '%s'%results['tiny']['numbers'][i], '\\\\ '
        print '\t\t\t\\hline'
    print '\t\t\\end{tabular}'
    print '\t\\end{table}'
    print   



def numbers_for_TP():
    results = {}
    for s in ['pimu', 'mumunu', 'tiny', 'nu', 'nubar', 'cosmics']:#studies:
        print 'Reading %s...'%s
        t = studies[s]['data']
        ntot = t.GetEntriesFast()
        nAfterCuts = [0]*10
        nAfterCuts_weighted = [0.]*10
        initialEventNums = []
        eventNums = []
        for event in t:
            w = weight(event,s) # e' importante fare questa cosa a monte, x come e' definita la funzione!
            if fit_converged(event) and goodTruthVtx(event,s):
                nAfterCuts[0] +=1
                nAfterCuts_weighted[0] += w
                initialEventNums.append(int(event.EventNumber))
                if chi2_cut(event, 5.):
                    nAfterCuts[1]+=1
                    nAfterCuts_weighted[1]+=w
                    if z_cut(event, s):#chi2_cut(event, 5):
                        nAfterCuts[1+1] += 1
                        nAfterCuts_weighted[1+1] += w
                        if vtxInAcceptance(event, s):#goodTrack(event):
                            nAfterCuts[2+1] +=1
                            nAfterCuts_weighted[2+1] +=w
                            #if 'cosmics' in s:
                            #   print 'mothers: ', event.DaughtersTruthMotherPDG[0], event.DaughtersTruthMotherPDG[1]
                            #   print "tracks: ", event.DaughtersTruthPDG[0], event.DaughtersTruthPDG[1] 
                            #   print "p candidate: ", event.P, "\t IP: ", event.IP0
                            #   print "pt candidate: ", event.Pt
                            if goodTrack(event):#doca_cut(event, 50):
                                nAfterCuts[3+1] +=1
                                nAfterCuts_weighted[3+1] +=w
                                if doca_cut(event, 50):#vtxInAcceptance(event, s):
                                    nAfterCuts[4+1]+=1
                                    nAfterCuts_weighted[4+1]+=w
                                    if has_muon_station(event, s, 1) and has_muon_station(event, s, 2):#ip_cut(event,500):#something_in_muon(event):
                                        nAfterCuts[5+1] +=1
                                        nAfterCuts_weighted[5+1] +=w
                                        if doca_cut(event, 30.):#something_in_ecal(event):
                                            nAfterCuts[6+1] +=1
                                            nAfterCuts_weighted[6+1] +=w
                                            if ip_cut(event, 500.):
                                                nAfterCuts[7+1] +=1
                                                nAfterCuts_weighted[7+1] +=w
                                                if ip_cut(event, 250.):
                                                    nAfterCuts[8+1] +=1
                                                    nAfterCuts_weighted[8+1] +=w
                                                    eventNums.append(int(event.EventNumber))
                                                    studies[s]['seen'].append(event.EventNumber)

        print '%s \t survived particles: '%s, nAfterCuts, ' of %s total'%ntot
        print '%s \t weighted survived particles: '%s, nAfterCuts_weighted, ' of %s total'%ntot
        results[s] = {}
        results[s]['ntot']=studies[s]['ntot']
        results[s]['numbers']=nAfterCuts
        results[s]['numbersWithWeights']=nAfterCuts_weighted
        results[s]['listOfSurvivors']=list(set(eventNums))
        studies[s]['listOfSurvivors']=list(set(eventNums))
        results[s]['listOfRecoed']=list(set(initialEventNums))
        studies[s]['listOfRecoed']=list(set(initialEventNums))
    print
    for s in ['pimu', 'mumunu', 'tiny']:#studies.keys():#['pimu', 'mumunu', 'bg']:
        for lst in results[s]:
            if not ('Survivors' in lst) and not ('Recoed' in lst):
                print s,list,results[s][lst]
    print
    print 'Events with surviving nu: %s \t with nubar: %s'%(len(results['nu']['listOfSurvivors']), len(results['nubar']['listOfSurvivors']))
    print 
    print 'Background:'
    print '\t\\begin{table}'
    print '\t\\centering Backgrounds: \\tiny'
    print '\t\t\\begin{tabular}{|l|c|c|c|c|c|c|}'
    print '\t\t\t\\hline'
    print '\t\t\tCut & Weighted $\\nu$ & $\\nu$ entries & Weighted $\\bar{\\nu}$ & $\\bar{\\nu}$ entries & Weighted cosmics & cosmics entries \\\\'
    print '\t\t\t\\hline'
    cuts = ['Track fit converged', '$\chi^2/n < 5$', '$z_{in}+5~\\text{m} < z_{vtx} < z_{out}$', '$z_{vtx} \in $ ellipse', 'Tracks $\in$ ellipse', 'DOCA $<$ 50 cm',
            'Muon1 \& Muon2', 'DOCA $<$ 30 cm', 'IP $<$ 5 m', 'IP $<$ 2.5 m']
    for i in xrange(len(results['nu']['numbers'])):
        print   '\t\t\t', cuts[i], ' & ', '%s'%latex_float(results['nu']['numbersWithWeights'][i]), ' & ', '%s'%results['nu']['numbers'][i], ' & ',\
                '%s'%latex_float(results['nubar']['numbersWithWeights'][i]), ' & ', '%s'%results['nubar']['numbers'][i], ' & ',\
                '%s'%latex_float(results['cosmics']['numbersWithWeights'][i]), ' & ', '%s'%results['cosmics']['numbers'][i], '\\\\ '
        print '\t\t\t\\hline'
    print '\t\t\\end{tabular}'
    print '\t\\end{table}'
    print 'Signal:'
    print '\t\\begin{table}'
    print '\t\\centering Signal: \\tiny'
    print '\t\t\\begin{tabular}{|l|c|c|c|c|c|c|}'
    print '\t\t\t\\hline'
    print '\t\t\tCut & $\\mathcal{A}(\\pi\\mu)$ & $\\pi\\mu$ entries & $\\mathcal{A}(\\mu\\mu\\nu)$ & $\\mu\\mu\\nu$ entries & $\\mathcal{A}(ee\\nu)$ & $ee\\nu$ entries \\\\'
    print '\t\t\t\\hline'
    cuts = ['Track fit converged', '$\chi^2/n < 5$', '$z_{in}+5~\\text{m} < z_{vtx} < z_{out}$', '$z_{vtx} \in $ ellipse', 'Tracks $\in$ ellipse', 'DOCA $<$ 50 cm',
            'Muon1 \& Muon2', 'DOCA $<$ 30 cm', 'IP $<$ 5 m', 'IP $<$ 2.5 m']
    for i in xrange(len(results['nu']['numbers'])):
        print   '\t\t\t', cuts[i], ' & ', '%s'%latex_float(results['pimu']['numbersWithWeights'][i]), ' & ', '%s'%results['pimu']['numbers'][i], ' & ',\
                '%s'%latex_float(results['mumunu']['numbersWithWeights'][i]), ' & ', '%s'%results['mumunu']['numbers'][i], ' & ',\
                '%s'%latex_float(results['tiny']['numbersWithWeights'][i]), ' & ', '%s'%results['tiny']['numbers'][i], '\\\\ '
        print '\t\t\t\\hline'
    print '\t\t\\end{tabular}'
    print '\t\\end{table}'
    print
    out = {'pimu':[], 'mumunu':[], 'tiny':[]}
    for s in ['pimu', 'mumunu', 'tiny']:
        for event in studies[s]['data']:
            if event.EventNumber in studies[s]['listOfSurvivors']:
                if event.NParticles > 1:
                    out[s].append(event.EventNumber)
                    #print s, event.EventNumber, event.NParticles, r.gGeoManager.FindNode(event.vtxx, event.vtxy, event.vtxz).GetName()
        print s, list(set(out[s]))

        
def badEvents():
    numbers_for_TP()
    hnpnu_recoed = r.TH1F('np nu recoed', 'np nu recoed', 60, 0, 59)
    hnpnu_survivors = r.TH1F('np nu survived', 'np nu survived', 60, 0, 59)
    seen = []
    for part in studies['nu']['data']:
        if not (part.EventNumber in seen):
            if part.EventNumber in studies['nu']['listOfSurvivors']:
                hnpnu_survivors.Fill(part.NParticles)
            if part.EventNumber in studies['nu']['listOfRecoed']:
                hnpnu_recoed.Fill(part.NParticles)
            seen.append(part.EventNumber)
    #c1 = r.TCanvas(); c1.cd()
    #hnpnu.Draw()
    #c1.Print('badEvents_nu_nparticles.png')
    hnpnubar_recoed = r.TH1F('np nubar recoed', 'np nubar recoed', 60, 0, 59)
    hnpnubar_survivors = r.TH1F('np nubar survived', 'np nubar survived', 60, 0, 59)
    seen = []
    for part in studies['nubar']['data']:
        if not (part.EventNumber in seen):
            if part.EventNumber in studies['nubar']['listOfSurvivors']:
                hnpnubar_survivors.Fill(part.NParticles)
            if part.EventNumber in studies['nubar']['listOfRecoed']:
                hnpnubar_recoed.Fill(part.NParticles)
            seen.append(part.EventNumber)
    #c1 = r.TCanvas(); c1.cd()
    #hnpnubar.Draw()
    #c1.Print('badEvents_nubar_nparticles.png')
    hnpnu_recoed.SetLineColor(r.kBlack); hnpnu_recoed.SetMarkerColor(r.kBlack)
    hnpnu_survivors.SetLineColor(r.kMagenta); hnpnu_survivors.SetMarkerColor(r.kMagenta)
    c1 = r.TCanvas(); c1.cd()
    hnpnu_recoed.DrawNormalized('p'); hnpnu_survivors.DrawNormalized('p same')
    c1.Print('nuAndNubar/nu_recoed_and_survived.png')
    c1.Print('nuAndNubar/nu_recoed_and_survived.pdf')
    c1.Close()
    hnpnubar_recoed.SetLineColor(r.kBlack); hnpnubar_recoed.SetMarkerColor(r.kBlack)
    hnpnubar_survivors.SetLineColor(r.kMagenta); hnpnubar_survivors.SetMarkerColor(r.kMagenta)
    c1 = r.TCanvas(); c1.cd()
    hnpnubar_recoed.DrawNormalized('p'); hnpnubar_survivors.DrawNormalized('p same')
    c1.Print('nuAndNubar/nubar_recoed_and_survived.png')
    c1.Print('nuAndNubar/nubar_recoed_and_survived.pdf')
    c1.Close()
    hnpnu_recoed.SetLineColor(r.kBlue); hnpnu_recoed.SetMarkerColor(r.kBlue)
    hnpnubar_recoed.SetLineColor(r.kGreen); hnpnubar_recoed.SetMarkerColor(r.kGreen)
    c1 = r.TCanvas(); c1.cd()
    hnpnu_recoed.DrawNormalized('p'); hnpnubar_recoed.DrawNormalized('p same')
    c1.Print('nuAndNubar/nu_and_nubar_recoed.png')
    c1.Print('nuAndNubar/nu_and_nubar_recoed.pdf')
    c1.Close()
    hnpnu_survivors.SetLineColor(r.kBlue); hnpnu_survivors.SetMarkerColor(r.kBlue)
    hnpnubar_survivors.SetLineColor(r.kGreen); hnpnubar_survivors.SetMarkerColor(r.kGreen)
    c1 = r.TCanvas(); c1.cd()
    hnpnu_survivors.DrawNormalized('p'); hnpnubar_survivors.DrawNormalized('p same')
    c1.Print('nuAndNubar/nu_and_nubar_survived.png')
    c1.Print('nuAndNubar/nu_and_nubar_survived.pdf')
    c1.Close()



def do_my_studies():
    results = {}
    for s in studies:
        print 'Reading %s...'%s
        t = studies[s]['data']
        ntot = t.GetEntriesFast()
        nAfterCuts = [0]*10
        nAfterCuts_weighted = [0.]*10
        for event in t:
            w = weight(event,s)
            if fit_converged(event): 
                nAfterCuts[0] += 1
                nAfterCuts_weighted[0] += w
                if chi2_cut(event,25):
                    nAfterCuts[1] +=1
                    nAfterCuts_weighted[1] +=w
                    #if chi2_cut(event,15):
                    if chi2_cut(event,5):
                        nAfterCuts[2]+=1
                        nAfterCuts_weighted[2]+=w
                        if doca_cut(event, 50.):
                            nAfterCuts[2+1] += 1
                            nAfterCuts_weighted[2+1] += w
                            if doca_cut(event, 30.):
                                nAfterCuts[3+1]+=1
                                nAfterCuts_weighted[3+1]+=w
                                if z_cut(event, s):
                                    nAfterCuts[4+1] += 1
                                    nAfterCuts_weighted[4+1] += w
                                    if ip_cut(event,500.):
                                        nAfterCuts[5+1]+=1
                                        nAfterCuts_weighted[5+1]+=w
                                        if ip_cut(event,250.):
                                            nAfterCuts[6+1]+=1
                                            nAfterCuts_weighted[6+1]+=w
                                            if nothing_in_veto5(event):
                                                nAfterCuts[7+1] += 1
                                                nAfterCuts_weighted[7+1] += w
                                                if nothing_in_liquidscint(event):
                                                    nAfterCuts[8+1] += 1
                                                    nAfterCuts_weighted[8+1] += w
        print '%s \t survived particles: '%s, nAfterCuts, ' of %s total'%ntot
        print '%s \t weighted survived particles: '%s, nAfterCuts_weighted, ' of %s total'%ntot
        results[s] = {}
        results[s]['ntot']=studies[s]['ntot']
        results[s]['numbers']=nAfterCuts
        results[s]['numbersWithWeights']=nAfterCuts_weighted
        addPlot(s+'-0cuts-IP', 'Impact parameter to target [cm]', t, 'IP0')
        addPlot(s+'-0cuts-Mass', 'Reco inv. mass', t, 'Mass'              )
        addPlot(s+'-0cuts-dPt', 'Tracks Pt', t, 'DaughtersPt'             )
        addPlot(s+'-0cuts-P','Reconstructed Momentum', t, 'P')
        addPlot(s+'-0cuts-Pt', 'Reconstructed Pt', t, 'Pt')
        addPlot(s+'-0cuts-PtOverP', 'Reconstructed Pt/P', t, 'Pt/P')
        addPlot(s+'-0cuts-dMinP', 'Min Tracks P', t, 'DaughterMinP'             )
        addPlot(s+'-0cuts-dMinPt', 'Min Tracks Pt', t, 'DaughterMinPt'             )
        addPlot(s+'-0cuts-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' )
        addPlot(s+'-0cuts-DOCA', 'Daughter tracks DOCA [cm]', t, 'DOCA' )
        addPlot(s+'-0cuts-dChi2', 'Daughter tracks #chi^{2}/n', t, 'DaughtersChi2/DaughtersNPoints' )
        #cut = "DaughtersChi2<25*DaughtersNPoints"
        cut = "DaughtersChi2<5*DaughtersNPoints"
        addPlot(s+'-chi2cut-IP', 'Impact parameter to target [cm]', t, 'IP0', cut)
        addPlot(s+'-chi2cut-Mass', 'Reco inv. mass', t, 'Mass'              , cut)
        addPlot(s+'-chi2cut-dPt', 'Tracks Pt', t, 'DaughtersPt'             , cut)
        addPlot(s+'-chi2cut-P', 'Reconstructed Momentum', t, 'P', cut)
        addPlot(s+'-chi2cut-Pt','Reconstructed Pt', t, 'Pt', cut)
        addPlot(s+'-chi2cut-PtOverP', 'Reconstructed Pt/P', t, 'Pt/P', cut)
        addPlot(s+'-chi2cut-dMinP', 'Min Tracks P', t, 'DaughterMinP'        , cut     )
        addPlot(s+'-chi2cut-dMinPt', 'Min Tracks Pt', t, 'DaughterMinPt'           , cut   )
        addPlot(s+'-chi2cut-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' , cut)
        addPlot(s+'-chi2cut-DOCA', 'Daughter tracks DOCA [cm]', t, 'DOCA'  , cut)
        cut += " && DOCA < 25"
        addPlot(s+'-chi2andDOCAcuts-IP', 'Impact parameter to target [cm]', t, 'IP0', cut)
        addPlot(s+'-chi2andDOCAcuts-Mass', 'Reco inv. mass', t, 'Mass'              , cut)
        addPlot(s+'-chi2andDOCAcuts-P','Reconstructed Momentum', t, 'P', cut)
        addPlot(s+'-chi2andDOCAcuts-Pt','Reconstructed Pt', t, 'Pt', cut)
        addPlot(s+'-chi2andDOCAcuts-PtOverP', 'Reconstructed Pt/P', t, 'Pt/P', cut)
        addPlot(s+'-chi2andDOCAcuts-dPt', 'Tracks Pt', t, 'DaughtersPt'             , cut)
        addPlot(s+'-chi2andDOCAcuts-dMinP', 'Min Tracks P', t, 'DaughterMinP'        , cut      )
        addPlot(s+'-chi2andDOCAcuts-dMinPt', 'Min Tracks Pt', t, 'DaughterMinPt'     , cut         )
        addPlot(s+'-chi2andDOCAcuts-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' , cut)
        #try:
        #   minz = studies[s]['geo']['lidT1I_1']['z']['pos'] + studies[s]['geo']['lidT1I_1']['z']['dim'] + 500.
        #   maxz = studies[s]['geo']['lidT6I_1']['z']['pos'] - studies[s]['geo']['lidT6I_1']['z']['dim']
        #except:
        minz = studies[s]['geo']['Veto_5']['z']['pos'] + studies[s]['geo']['Veto_5']['z']['dim']
        maxz = studies[s]['geo']['Tr1_1']['z']['pos'] + studies[s]['geo']['Tr1_1']['z']['dim']
        cut += " && vtxz > %s && vtxz < %s"%(minz, maxz)
        addPlot(s+'-chi2andDOCAandZcuts-IP', 'Impact parameter to target [cm]', t, 'IP0', cut)
        addPlot(s+'-chi2andDOCAandZcuts-Mass', 'Reco inv. mass', t, 'Mass'              , cut)
        addPlot(s+'-chi2andDOCAandZcuts-P','Reconstructed Momentum', t, 'P', cut)
        addPlot(s+'-chi2andDOCAandZcuts-Pt','Reconstructed Pt', t, 'Pt', cut)
        addPlot(s+'-chi2andDOCAandZcuts-PtOverP', 'Reconstructed Pt/P', t, 'Pt/P', cut)
        addPlot(s+'-chi2andDOCAandZcuts-dPt', 'Tracks Pt', t, 'DaughtersPt'             , cut)
        addPlot(s+'-chi2andDOCAandZcuts-dMinP', 'Min Tracks P', t, 'DaughterMinP'         , cut     )
        addPlot(s+'-chi2andDOCAandZcuts-dMinPt', 'Min Tracks Pt', t, 'DaughterMinPt'      , cut        )
        addPlot(s+'-chi2andDOCAandZcuts-VtxZ', 'Reco vertex z position [cm]', t, 'vtxz' , cut)
    
    
    for s in studies.keys():#['pimu', 'mumunu', 'bg']:
        for list in results[s]:
            print s,list,results[s][list]
