"""
 **********************************************************************************
 * Project: HistFitter - A ROOT-based package for statistical data analysis       *
 * Package: HistFitter                                                            *
 *                                                                                *
 * Description:                                                                   *
 *      Example pull plot based on the pullPlotUtils module. Adapt to create      *
 *      your own style of pull plot. Illustrates all functions to redefine to     *
 *      change labels, colours, etc.                                              * 
 *                                                                                *
 * Authors:                                                                       *
 *      HistFitter group, CERN, Geneva                                            *
 *                                                                                *
 * Redistribution and use in source and binary forms, with or without             *
 * modification, are permitted according to the terms listed in the file          *
 * LICENSE.                                                                       *
 **********************************************************************************
"""

#!/usr/bin/env python

import ROOT
from ROOT import *
ROOT.PyConfig.IgnoreCommandLineOptions = True
gSystem.Load("libSusyFitter.so")
#gROOT.Reset()
ROOT.gROOT.SetBatch(True)

import os, pickle, subprocess

import pullPlotUtils
from pullPlotUtils import makePullPlot 

# Build a dictionary that remaps region names
def renameRegions():
    myRegionDict = {}

    # Remap region names using the old name as index, e.g.:
#    myRegionDict["SS_metmeff2Jet"] = "SR exclusion"
  #  myRegionDict["Top1L"] = "Top"
   # myRegionDict["Top2L"] = "Top"
    
    
    return myRegionDict

# Build a list with all the regions you want to use
def makeRegionList():
    regionList=[]
    regionList += ["tW1L_CRtt","tW1L_CRWp","tW1L_CRWm","tW2L_CRtt","tW2L_CRttZ","tW2L_CRWZ","tW1L_VRtt_VR1","tW1L_VRtt_VR2","tW1L_VRW_VR1p","tW1L_VRW_VR1m","tW1L_VRW_VR2p","tW1L_VRW_VR2m","tW2L_VRtt","tW2L_VR3L"]
    regionList += ["SR1LBin0","SR1LBin1","SR1LBin2","SR1LBin3","SR1LBin4","SR2L"]

    return regionList

# Define the colors for the pull bars
def getRegionColor(name):
    if name.find("tW2L_CRtt") != -1: return kBlue+3
    if name.find("tW2L_CRttZ") != -1: return kBlue+3
    if name.find("tW2L_CRWZ") != -1: return kBlue+3
    if name.find("tW2L_VRtt")  != -1: return kOrange
    if name.find("tW2L_VR3L")  != -1: return kOrange
    if name.find("tW1L_CRtt") != -1: return kBlue+3
    if name.find("tW1L_CRWp") != -1: return kBlue+3
    if name.find("tW1L_CRWm") != -1: return kBlue+3
    if name.find("tW1L_VRtt_VR1")  != -1: return kOrange
    if name.find("tW1L_VRtt_VR2")  != -1: return kOrange
    if name.find("tW1L_VRW_VR1p")  != -1: return kOrange
    if name.find("tW1L_VRW_VR1m")  != -1: return kOrange
    if name.find("tW1L_VRW_VR2p")  != -1: return kOrange
    if name.find("tW1L_VRW_VR2m")  != -1: return kOrange
    if name.find("SR1LBin0")  != -1: return kRed
    if name.find("SR1LBin1")  != -1: return kRed
    if name.find("SR1LBin2")  != -1: return kRed
    if name.find("SR1LBin3")  != -1: return kRed
    if name.find("SR1LBin4")  != -1: return kRed
    if name.find("SR2L")  != -1: return kRed
 
    return 1

# Define the colors for the stacked samples
def getSampleColor(sample):
    if sample == "Top1L":         return kOrange - 2
    if sample == "Top2L":         return kOrange - 2
    if sample == "SingleTop":       return kOrange + 2
    if sample == "Wjets":     return kCyan + 2
    if sample == "ttV1L":     return kAzure - 3
    if sample == "ttV2L":     return kAzure - 3
    if sample == "Zjets":     return kRed - 4
    if sample == "Diboson1L":     return kMagenta - 4
    if sample == "Diboson2L":     return kMagenta - 4
    if sample == "ttH":     return kYellow
    if sample == "tWZ":     return kViolet + 1
    if sample == "Fakes2L":     return kGray
    if sample == "Fakes3L":     return kGray
    else:
        print "cannot find color for sample (",sample,")"

    return 1

def main():
    # Override pullPlotUtils' default colours (which are all black)
    pullPlotUtils.getRegionColor = getRegionColor
    pullPlotUtils.getSampleColor = getSampleColor

    # Where's the workspace file? 
    wsfilename = "./BkgOnly_combined_NormalMeasurement_model_afterFit.root" # 

    # Where's the pickle file?
    pickleFilename = "MyYieldsTable.pickle"
    
    # Run blinded?
    doBlind = False

    # Used as plot title
    region = "tW_Combination"

    # Samples to stack on top of eachother in each region
    samples = "Top1L,Top2L,SingleTop,Wjets,Zjets,Diboson1L,Diboson2L,ttV1L,ttV2L,ttH,tWZ,Fakes2L,Fakes3L"
    
    # Which regions do we use? 
    regionList = makeRegionList()

    # Regions for which the label gets changed
    renamedRegions = renameRegions()

    if not os.path.exists(pickleFilename):
        print "pickle filename %s does not exist" % pickleFilename
        print "will proceed to run yieldstable again"
        
        # Run YieldsTable.py with all regions and samples requested
        cmd = "YieldsTable.py -c %s -s %s -w %s -o MyYieldsTable.tex" % (",".join(regionList), samples, wsfilename)
        print cmd
        subprocess.call(cmd, shell=True)

    if not os.path.exists(pickleFilename):
        print "pickle filename %s still does not exist" % pickleFilename
        return
    
    # Open the pickle and make the pull plot
    makePullPlot(pickleFilename, regionList, samples, renamedRegions, region, doBlind, plotSignificance='atlas')

if __name__ == "__main__":
    main()