########################################################################
# -*- python -*-
# File Name .....: new.py
# Original Author: Alexandros Attikis
# e-mail ........: attikis@cern.ch
# Institute .....: UCY
# Created .......: Friday, 16 Dec 2010 
# Last Edited ...: Thursday, 6 Jan 2011
# Description ...: Python script that takes as input a dat or text file
#                  generated by the hplusPrintCounters.py script and
#                  uses as input the number of events passing specific
#                  cuts. It returns the efficiency and purity of the QCD
#                  samples, based on Bayesian methods. The purity is not
#                  to be taken seriously as it is a rough calculation and
#                  not properly done. A more accurate method for estimating
#                  the purity of each QCD pT-bin sample is determined by 
#                  the use of the root macro "pseudo_experiments.C".
########################################################################

#!/usr/bin env python
### to execute: |myLinux> python readfile.py testFile.dat

### System imports
import os, sys
from array import array
### ROOT imports
from ROOT import *
### HiggsAnalysis imports
from HiggsAnalysis.HeavyChHiggsToTauNu.tools.dataset import *
from HiggsAnalysis.HeavyChHiggsToTauNu.tools.histograms import *
from HiggsAnalysis.HeavyChHiggsToTauNu.tools.tdrstyle import *
import HiggsAnalysis.HeavyChHiggsToTauNu.tools.styles as styles
### User imports
from bayes import * 
from myArrays import *

### Provide the name of the file to be opened 
if len(sys.argv) > 1:
    file = sys.argv[1]
### Otherwise ose as default the name file.dat
else:
    file = 'file.dat'

### If the file already exists remove it since the script will automatically generate one.
#if os.path.exists(file):
#    os.system('rm %s'%(file))
#os.system('hplusPrintCounters3.py --mode xsect --mainCounterOnly > %s ; echo "EOF" >> %s'%(file,file))

######################################################################################
### Options
verbosity = False ### False to suppress most "print" commands

### Notes
# For some QCD samples the purity appears to be quite hight (almost 100%) which doesnt make sense! The reason is the very high Normalisation
# factors that these samples have which are much bigger than the rest. So, because the statistics are so low, if you try to set an upper limit
# on the purity say, that event when weighted by the Normalization factor could be either 1000 or 0, anything in between. This overwhelms the
# rest of the events considered (EWK = TTbar + WJets) and that's why you have so high purity. This is okay, BUT, you have to make sure that
# PURITY ERROR bar is also huge (spans from max to 0!)

######################################################################################
### Open file for processing
file = open('%s'%(file), 'r')

### Define lists to be used
DataSetName = []
CutNamesAndValues = []
DataSetAndXSectAndNormFact = []

######################################################################################
### Loop over lines and save the cut-flow info (Cross Section for MC, Events for data)
lim = True
while file:
    line = file.readline()
    # split the line and puts elements in list s as strings
    s = line.split()
    if 'EOF' in s:
        break
    n = len(s)
    # 10 is the number of columns
    ### Investigate lists which have three elements (i.e. 3 is the number of columns of each list)
    if n == 3:
            ### Save the names of the datasets, cross section and normalisation factor
            DataSetAndXSectAndNormFact.append([s[0], s[1], s[2]])
    if n == 10:
        if lim == True:
            for i in range(n):
                DataSetName.append(s[i])
            lim = False
        else:
            #print "CutNamesAndValues\tCut\tValue" 
            for i in range(1,n):
                CutNamesAndValues.append([DataSetName[i],s[0],float(s[i])])
                #print "%s\t%s\t%.0f"%(DataSetName[i],s[0],float(s[i]))

if verbosity == True:
    ### Print the DataSet Name, Cross Section and MC Normalisation Factor (in this order)
    print
    print "********************************************************************"
    print "* Printing DataSet Name, Cross Section and MC Normalisation Factor *" 
    print "********************************************************************"
    print "DataSetAndXSectAndNormFact = ", DataSetAndXSectAndNormFact
    print "For Wjets  we have ..."
    print "DataSetAndXSectAndNormFact[0][0] = ", DataSetAndXSectAndNormFact[0][0]
    print "DataSetAndXSectAndNormFact[0][1] = ", DataSetAndXSectAndNormFact[0][1]
    print "DataSetAndXSectAndNormFact[0][2] = ", DataSetAndXSectAndNormFact[0][2]
    print "For QCD_Pt30to50 we have ..."
    print "DataSetAndXSectAndNormFact[1][0] = ", DataSetAndXSectAndNormFact[1][0]
    print "DataSetAndXSectAndNormFact[1][1] = ", DataSetAndXSectAndNormFact[1][1]
    print "DataSetAndXSectAndNormFact[1][2] = ", DataSetAndXSectAndNormFact[1][2]

### Declare the Normalisation Factors for each DataSet
WJets_NormFactor        = float(DataSetAndXSectAndNormFact[0][2])
qcd30to50_NormFactor    = float(DataSetAndXSectAndNormFact[1][2])
qcd50to80_NormFactor    = float(DataSetAndXSectAndNormFact[2][2])
qcd80to120_NormFactor   = float(DataSetAndXSectAndNormFact[3][2])
qcd120to170_NormFactor  = float(DataSetAndXSectAndNormFact[4][2])
qcd170to230_NormFactor  = float(DataSetAndXSectAndNormFact[5][2])
qcd230to300_NormFactor  = float(DataSetAndXSectAndNormFact[6][2])
TTbarJets_NormFactor    = float(DataSetAndXSectAndNormFact[7][2])
TTbar_NormFactor        = float(DataSetAndXSectAndNormFact[8][2])
### Define lists that will hold the information (Cut Name and Xsection/Events surviving Cut)
WJets        = []
qcd30to50    = []
qcd50to80    = []
qcd80to120   = []
qcd120to170  = []
qcd170to230  = []
qcd230to300  = []
TTbarJets    = []
TTbar        = []
######################################################################################
### Loop over all cut names and create new lists according to the DataSet Name
for item in CutNamesAndValues:
    name = item[0]
    tmp = [item[1],item[2]]
    if name == 'WJets':
        WJets.append(tmp)
    if name == 'QCD_Pt30to50':
        qcd30to50.append(tmp)
    if name == 'QCD_Pt50to80':
        qcd50to80.append(tmp)
    if name == 'QCD_Pt80to120':
        qcd80to120.append(tmp)
    if name == 'QCD_Pt120to170':
        qcd120to170.append(tmp)
    if name == 'QCD_Pt170to230':
        qcd170to230.append(tmp)
    if name == 'QCD_Pt230to300':
        qcd230to300.append(tmp)
    if name == 'TTbarJets':
        TTbarJets.append(tmp)
    if name == 'TTbar':
        TTbar.append(tmp)
######################################################################################
### Print lists 
if verbosity == True:
    print
    print "*******************************************"
    print "* Printing cut-flow for each DataSet Name *"
    print "*******************************************"
    print 'qcd30to50', qcd30to50
    print
    print 'qcd50to80', qcd50to80
    print
    print 'qcd80to120', qcd80to120
    print
    print 'qcd120to170', qcd120to170
    print
    print 'qcd170to230', qcd170to230
    print
    print 'qcd230to300', qcd230to300
    print
    print 'TTbarJets', TTbarJets
    print
    print 'TTbar', TTbar
    print
    print 'W+Jets', WJets
    print "***************************************************************************************************************"
######################################################################################
### Calculate cross sections before and after cuts
# ElectroWeak (TTbar + WJets)
XSectTotalTTbar  = TTbar[4][1]
XSectTotalWJets  = WJets[4][1]
XSectPassedTTbar = TTbar[-1][1]
XSectPassedWJets = WJets[-1][1]
XSectTotalEWK    = XSectTotalTTbar + XSectTotalWJets
XSectPassedEWK   = XSectPassedTTbar + XSectPassedWJets
### QCD 
XSectTotalQCD  = qcd30to50[4][1]+qcd50to80[4][1]+qcd80to120[4][1]+qcd120to170[4][1]+qcd170to230[4][1]+qcd230to300[4][1]
XSectPassedQCD = qcd30to50[-1][1]+qcd50to80[-1][1]+qcd80to120[-1][1]+qcd120to170[-1][1]+qcd170to230[-1][1]+qcd230to300[-1][1]
XSectPassedQCD_UpLimit = qcd30to50_NormFactor+qcd50to80_NormFactor+qcd80to120[-1][1]+qcd120to170[-1][1]+qcd170to230[-1][1]+qcd230to300[-1][1]
### Save Events before and after cuts
# ElectroWeak (TTbar + WJets)
EvtsTotalTTbar  = XSectTotalTTbar/TTbar_NormFactor
EvtsTotalWJets  = XSectTotalWJets/WJets_NormFactor
EvtsTotalEWK    = EvtsTotalTTbar + EvtsTotalWJets
#
EvtsPassedTTbar = XSectPassedTTbar/TTbar_NormFactor
EvtsPassedWJets = XSectPassedWJets/WJets_NormFactor
EvtsPassedEWK   = EvtsPassedTTbar + EvtsPassedWJets
### QCD 
EvtsTotalQCD  = qcd30to50[4][1]/qcd30to50_NormFactor + qcd50to80[4][1]/qcd50to80_NormFactor + qcd80to120[4][1]/qcd80to120_NormFactor + qcd120to170[4][1]/qcd120to170_NormFactor + qcd170to230[4][1]/qcd170to230_NormFactor + qcd230to300[4][1]/qcd230to300_NormFactor
EvtsPassedQCD = qcd30to50[-1][1]/qcd30to50_NormFactor + qcd50to80[-1][1]/qcd50to80_NormFactor + qcd80to120[-1][1]/qcd80to120_NormFactor + qcd120to170[-1][1]/qcd120to170_NormFactor + qcd170to230[-1][1]/qcd170to230_NormFactor + qcd230to300[-1][1]/qcd230to300_NormFactor

### Calculate Efficiencies
# QCD combined: Using Bayes
qcdObj_eff = bayesDivide(EvtsPassedQCD, EvtsTotalQCD)
QCD_efficiency_bayes = qcdObj_eff.GetY()[0]
QCD_eff_errUp = qcdObj_eff.GetEYhigh()[0]
QCD_eff_errDown = qcdObj_eff.GetEYlow()[0]
# QCD combined: Using Standard Efficiency Definition
QCD_efficiency_estimate = XSectPassedQCD/XSectTotalQCD

# TTbar: Using Bayes
ttbarObj_eff = bayesDivide(EvtsPassedTTbar, EvtsTotalTTbar)
TTbar_efficiency_bayes = ttbarObj_eff.GetY()[0]
TTbar_eff_errUp = ttbarObj_eff.GetEYhigh()[0]
TTbar_eff_errDown = ttbarObj_eff.GetEYlow()[0]
# TTbar: Using Standard Efficiency Definition
TTbar_efficiency_estimate = XSectPassedTTbar/XSectTotalTTbar

# WJets: Using Bayes
wjetsObj_eff = bayesDivide(EvtsPassedWJets, EvtsTotalWJets)
WJets_efficiency_bayes = wjetsObj_eff.GetY()[0]
WJets_eff_errUp = wjetsObj_eff.GetEYhigh()[0]
WJets_eff_errDown = wjetsObj_eff.GetEYlow()[0]
# WJets: Using Standard Efficiency Definition
WJets_efficiency_estimate = XSectPassedWJets/XSectTotalWJets

### Calculate Purity of QCD
# Using Bayes
qcdObj_pur = bayesDivide(EvtsPassedQCD, EvtsPassedEWK+EvtsPassedQCD)
QCD_purity_bayes = qcdObj_pur.GetY()[0]
QCD_pur_errUp = qcdObj_pur.GetEYhigh()[0]
QCD_pur_errDown = qcdObj_pur.GetEYlow()[0]
# Using Standard Purity Definition
QCD_purity_estimate = XSectPassedQCD/(XSectPassedEWK+XSectPassedQCD)
# Try to roughly place an Upper Limit
QCD_purity_UpLimit = XSectPassedQCD_UpLimit/(XSectPassedEWK+XSectPassedQCD_UpLimit)
######################################################################################
### Calculate Efficiency using the ROOT class "bayesDivide", for each QCD DataSet
# Declare (double) arrays to be filled with efficiencies and purities
effArray         = array("d",[ ])
effArray_errDown =  array("d",[ ])
effArray_errUp   =  array("d",[ ])

effObj = bayesDivide(qcd30to50[-1][1], qcd30to50[4][1])
eff = effObj.GetY()[0]
eff_errUp = effObj.GetEYhigh()[0]
eff_errDown = effObj.GetEYlow()[0]
effArray.append(eff)
effArray_errDown.append(eff_errDown)
effArray_errUp.append(eff_errUp)

effObj = bayesDivide(qcd50to80[-1][1], qcd50to80[4][1])
eff = effObj.GetY()[0]
eff_errUp = effObj.GetEYhigh()[0]
eff_errDown = effObj.GetEYlow()[0]
effArray.append(eff)
effArray_errDown.append(eff_errDown)
effArray_errUp.append(eff_errUp)

effObj = bayesDivide(qcd80to120[-1][1], qcd80to120[4][1])
eff = effObj.GetY()[0]
eff_errUp = effObj.GetEYhigh()[0]
eff_errDown = effObj.GetEYlow()[0]
effArray.append(eff)
effArray_errDown.append(eff_errDown)
effArray_errUp.append(eff_errUp)

effObj = bayesDivide(qcd120to170[-1][1], qcd120to170[4][1])
eff = effObj.GetY()[0]
eff_errUp = effObj.GetEYhigh()[0]
eff_errDown = effObj.GetEYlow()[0]
effArray.append(eff)
effArray_errDown.append(eff_errDown)
effArray_errUp.append(eff_errUp)

effObj = bayesDivide(qcd170to230[-1][1], qcd170to230[4][1])
eff = effObj.GetY()[0]
eff_errUp = effObj.GetEYhigh()[0]
eff_errDown = effObj.GetEYlow()[0]
effArray.append(eff)
effArray_errDown.append(eff_errDown)
effArray_errUp.append(eff_errUp)

effObj = bayesDivide(qcd230to300[-1][1], qcd230to300[4][1])
eff = effObj.GetY()[0]
eff_errUp = effObj.GetEYhigh()[0]
eff_errDown = effObj.GetEYlow()[0]
effArray.append(eff)
effArray_errDown.append(eff_errDown)
effArray_errUp.append(eff_errUp)
######################################################################################
print "effArray = ", effArray
if verbosity == True:
    print
    print "**************"
    print "* Efficiency *"
    print "**************"
    print "effArray = ", effArray
    print
    print "effArray_errDown = ", effArray_errDown
    print
    print "effArray_errUp = ", effArray_errUp
    print "***************************************************************************************************************"
######################################################################################
print "*******************"
print "* All QCD Samples *"
print "*******************"
print "--> Purity (Bayes) = ", QCD_purity_bayes
print "--> Purity (Estimate) = ", QCD_purity_estimate
print "--> Purity (Upper Limit) = ", QCD_purity_UpLimit
print
print "--> Efficiency (Bayes) = ", QCD_efficiency_bayes, " + ", QCD_eff_errUp, " - ", QCD_eff_errDown
print "--> Efficiency (Estimate) = ", QCD_efficiency_estimate
print "****************"
print "* TTbar Sample *"
print "****************"
print "--> Efficiency (Bayes) = ", TTbar_efficiency_bayes, " + ", TTbar_eff_errUp, " - ", TTbar_eff_errDown
print "--> Efficiency (Estimate) = ", TTbar_efficiency_estimate
print "****************"
print "* WJets Sample *"
print "****************"
print "--> Efficiency (Bayes) = ", WJets_efficiency_bayes, " + ", WJets_eff_errUp, " - ", WJets_eff_errDown
print "--> Efficiency (Estimate) = ", WJets_efficiency_estimate
print "***************************************************************************************************************"
######################################################################################

### Customise ROOT
gROOT.Reset()
ROOT.gROOT.SetBatch(False)
style = TDRStyle()
style.setPalettePretty()
style.setWide(True)

### Create TCanvas and TGraph with Asymmetric Error Bars using Bayesian Statistical Tools
canvas2 = TCanvas ("QCDMeasurementByIsolation", "QCDMeasurementByIsolation", 20, 20, 800, 800)
canvas2.cd()

qcdBin  = array("d",[50, 80, 120, 170, 230, 300, 300 ])
qcdBin_down = array("d",[0, 0, 0, 0, 0, 0, 0 ])
qcdBin_up  = array("d",[0, 0, 0, 0, 0, 0, 0 ])
graph = TGraphAsymmErrors(6, qcdBin, effArray, qcdBin_down,  qcdBin_up, effArray_errDown, effArray_errUp) ### works
# TGraphAsymmErrors(Int_t n, const Double_t* x, const Double_t* y, const Double_t* exl = 0, const Double_t* exh = 0, const Double_t* eyl = 0, const Double_t* eyh = 0)

graph.Draw("AP")
graph.SetMarkerStyle(kFullCircle)
graph.SetMarkerColor(kRed)
graph.SetMarkerSize(1)
gPad.SetLogy(0)
gPad.SetLogx(0)
#graph.GetXaxis().SetRangeUser( 0.6 , 0.8)
#graph.GetYaxis().SetRangeUser( 0.00001 , 0.2)
graph.GetYaxis().SetTitle("Efficiency (#epsilon)")
graph.GetXaxis().SetTitle("#hat{p}_{T}")
### Finally, re-draw graph and update canvas and gPad
graph.Draw("AP")
canvas2.Update()
#gPad.Update()

################
print
raw_input('Press \"Enter\" to exit to terminal ...')
