##################################################################################################
#                                                                                                #    
#                                      ```...`                                                   #    
#                                   `......-::::.`...---.`                                       #    
#                                  ``     `.-::-.....---:::-`                                    #    
#                                           ..`      `.--:://-                                   #    
#                                     `//:`            .-:::/+:                                  #    
#                                    +m/`sy            `.-:://+`                                 #    
#                              `    :mMmdN/        .::.`.-:://+-                                 #    
#                   ``-//-----------odMMMs`      -yh:/m+.-:///+-                                 #    
#               `..--------/oyy/---::/+sh.      -mNd/omh-:://++.                                 #    
#             `.-------------:------::///:.     yNMMMMNs-:///+:                                  #    
#           `------------------------:::://:`   omNNNmy-:///+:`                                  #    
#          .--------........---------:::::/+/.  `+sys/-::/+++:`                                  #    
#         .---------.---------------:::::///++:..------::://+o+.                                 #    
#        .---------...-------------:::::://++/----------:::///+o-                                #    
#        -------------------------::::::////:---.-------::::///+o.                               #    
#       `:----------------------:::::::///:---..-------:::::///+o/                               #    
#       .:--------------------:::::://///:-------------::::////+o+             `.-               #    
#       `::---------------:::::::://////:------------::::::///+++/            ..-:               #    
#       `::::--------:::::::////////////-----------::::://///++++-           .--:/`  ``..        #    
#        -/:::::::::::://///////////////--------::::://////+++++/`          .--:/+...--/.        #    
#         :///::::://////++++++++++++++/-----::::///////++++++//.         `.---::----:/:         #    
#          :+//////++++++++++++++++++o+/-:::::///////++++++++/:`         `..-------:/+/`````.`   #    
#           ./++++ooooooooooooooooooo/:::::///+++++++++++++//-`         `.---------///-..--:-    #    
#            `-+osssso/+sssssssssso+/::://+++ooooooooo+++/:-`          `----------------://-     #    
#               .:+:.` `-+ossyyyysooo++++oosssssoooo++//+:.            ---------------:/+/.      #    
#                         `-/+osyyyyyysssyysssssoo++++++oso/.----.````.-------------://+/`       #    
#                             `..-://+++osysssssooosyso+ooossydhs+----------------:://++:        #    
#                                       `+ssssssssssyyyo+++++/:--------------::::////////-`      #    
#                   `                  -+osooooooooosshhs/:----..--------::://////++//:----.     #    
#           ..`    ...`               :ssooooossssssyyy+:------------::/++ooooooooooo+/:----.`   #    
#           .--..``-.--`              .oysssssssssshdo:----------::/++ossyhddhhdo....--------`   #    
#     ``    `:-.---:----`             .----------:+mh:------:://++oohhhdddddddddms`              #    
#    `.......-:----------.`          .-...```````..smy++//++ooooooooyddddddddddddmy`             #    
#      `.-------------------.`````..-/:--.....`````.:shhhdddho/::/+oohdddddhhhhhddms             #    
#         `-------....--------------/:-...````````````..----:/.` ``-/+hddddhhhddddmd:            #    
#           `.-----.--..-----------:/:-....````````````````..::``  `.:ommddhdddmmmmms            #    
#             `---------------::://++:--......``````````....--/:`````.:sdddmmmmmdddmh`           #    
#             .--------:://+++++//:-//:--.....................-/-````.-/yydhyyyhhhhdd`           #    
#           `..----:://///:--..``   :+::--------..............-//.```.-:sdhyyyyyhhhhd`           #    
#          ..--::://::-..-://+///:-`.o/::----.................-:+:.``..:+hyyyyyyhhhhs            #    
#         .-::-..`   `:+ooooooooooso/o+:::---------------------:++-....:/yyyyyyyyhhh:            #    
#          `        .syso+++++++oooosso+/::::::-----------::::::+s/-..--:shyyyyyhhys`            #    
#                `.../yysoooooo+oooossyso///::------------::://++yo:---:/ohyyyyhhys.             #    
#                :o+:`.oysoooooooooossyhhoo+:----------:::::://++dh+::::/+dhhhhhys.              #    
#                .ooo:``ohyssoooooosssyhhoooo+:------::::::////++ymo//::/+dhhddh/`               #    
#                 /sso:``ohyyssssssssyyhdysooooo/:::://++///++++shds+////+mddho-                 #    
#                  :sss/.`odhyyyyyyyyyyhdddyssooso///+oysooooosshNhso++//+hyo-                   #    
#                   -syy+-.sdhhhyyyyyyyyyyhhhyyssssoosyyysssyhdmNdysso++/ossso+/-.               #    
#                    `/syo:-sddhhhyyyyysssyyhhhhyyyysosssssossyyhyyssso+ossoo++//-               #    
#                      `+hy/-oddhhyyyyssssyyyhddddhh+  `/oosssssssssoo+ydhho/:-.`                #    
#                       :osso:+ddhhyyyssssyyyhddhhyo.   `/oosssssssssssyhhhhhhhhy+.              #    
#                        -sso+::yddhyyyyyyyyyhho `       `:+osssssssssyhsoshdhhhhddo-`           #    
#                         .oys/--odddhhhhyyhhhs.           .:osyyyyyyhsoooyhdhyssyhds/           #    
#                          `/yy+:-/dmdddhhhhyo.             `yhhhhhysoosyhddhsoosyyds-           #    
#                            .+ys+/+dmdddhy+.               `shyssoosyhdddhysooosyhdo-`          #    
#                              .+so:-/::-.                   `-+oshhhhhhyyssoooosyhd/-           #    
#                                ``                        .:+oooo++oooooooooooosyhh-.           #    
#                                                        `+oooo++++++++++++++ooosyho`            #    
#                                                        +ssooooo++++++++++++oosyhd/.            #    
#                                                       `yyysssoooooooo+++oooosyyhh-.            #    
#                                                        sdhhyyyssssssssssssyyyhhd+-.            #    
#                                                        .hmddddhhhhhhhhhhhhhddddo:-             #    
#                                                         `:ohddmmddddddddddhyso/::              #    
#                                                            -:/++ooooooo++//:--.`               #    
##################################################################################################
#                           ____       _                                                         #
#                          / ___| __ _| | __ _ _ __   __ _  __ _  ___                            #
#                         | |  _ / _` | |/ _` | '_ \ / _` |/ _` |/ _ \                           #
#                         | |_| | (_| | | (_| | |_) | (_| | (_| | (_) |                          # 
#                     _____\____|\__,_|_|\__,_| .__/ \__,_|\__, |\___/ _                         #  
#                    |  ___| __ __ _ _ __ ___ |_|____      |___/  _ __| | _                      #_
#                    | |_ | '__/ _` | '_ ` _ \ / _ \ \ /\ / / _ \| '__| |/ /                     #
#                    |  _|| | | (_| | | | | | |  __/\ V  V / (_) | |  |   <                      # 
#                    |_|  |_|  \__,_|_| |_| |_|\___| \_/\_/ \___/|_|  |_|\_\                     #
#                                                                                                # 
##################################################################################################

import ROOT as r
from   ROOT import gROOT, TCanvas, TFile, TGraphErrors, SetOwnership
import math, sys, optparse, array, copy, os
import gc, inspect, __main__
import numpy as np
import time

import include.Sample as Sample
import include.Launcher as Launcher
import include.helper as helper
import include.Canvas as Canvas
import include.CutManager as CutManager
from include.Utils import *

####################################### CLASS DEFINITION #########################################

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'



################################# GLOBAL VARIABLES DEFINITION ####################################

runningfile = os.path.abspath(__file__)
WORKPATH = ''
for level in runningfile.split('/')[:-1]: 
    WORKPATH += level
    WORKPATH += '/'

if __name__ == "__main__":

    print bcolors.HEADER
    print "                                 ```.....`                              "        
    print "                               `.````.-:::-.-------.                    "        
    print "                              ``      .--.`````.--://-                  "
    print "                                   `   `        `-::://`                "
    print "                                `os/y:           .-::/+/                "
    print "                                yNdod+        ````-:://+`               "
    print "                 ``-/-..------.-hNMMs`     `oy/sy.-:://+.               "
    print "              `..------:+ss/--::/+sh.     `dNh/sN/-://+/`               "
    print "           `.-------------------::///.    /NMMMNm::///+-                "
    print "         `.---------------------:::://:`  -dmNNh/:///+-                 "
    print "        `--------.---.----------:::::/+/.``-++/:-://+o+.                "
    print "       `--------..--------------:::://++/---------:://++:               "
    print "       .----------------------::::::/++:----------::://+o:              "
    print "       ---------------------::::::://:---..------:::://++o`             "
    print "      `:-------------------::::::///:----.------::::://++o.             "
    print "      `::--------------::::::://///:-----------:::::///+++.             "
    print "       ::::------:::::::://////////----------:::::////+++/`             "
    print "       ./:::::::::::///////////////-------:::://////++++/-              "
    print "        -///////////++++++++++++++/----:::///////++++++/-               "
    print "         ./+///+++++ooooooooo+ooo+::::::///++++++++++//.                "
    print "          `-+ooosssooooosssssso+/:::///+++++oooo+++//-`                 "
    print "            `-+so/-.:+sssyyysso++///+oooooooooo++//-`                   "
    print "               ``    `.:+osyyyyysssssssssssoo++++oo+:`..```             "
    print "                         `.-:/+++oooyyysssooosso+ooosoyhyo/-.           "
    print "                                  ``+sssssssssyyso++++/:------::        "
    print "                                  -+oooooooooosyhy+:----.-------::      "
    print '########################################################################' 
    print '                  Starting IFCA-LLP analysis...                         '
    print '########################################################################' + bcolors.ENDC

    outputFolder = WORKPATH + 'plots'
    print "outputfolder = ",  outputFolder


    parser = optparse.OptionParser(usage='usage: %prog [opts] FilenameWithSamples', version='%prog 1.0')
    parser.add_option('-o', '--out', action='store', type=str, dest='out', default='', help='Output tag')
    parser.add_option('-q', '--qenv', action='store', type=str, dest='qenv', default='', help='Select if you want to send the job a condor queue')
    (opts, args) = parser.parse_args()

    if opts.qenv == '': opts.qenv = False

    ############# Set the TDR plot style
    r.gROOT.LoadMacro(WORKPATH + 'include/tdrstyle.C+')
    r.gROOT.SetBatch(1)
    r.setTDRStyle()



    ############# DEFINICION DE BACKGROUND
    ###
    ### Lista de fondos que queremos incluir (tercera columna del dat):
    Backgrounds = []
    Backgrounds.append('DYJetsToLL_M-50') 



    ############# DEFINICION DE SENYAL
    ###
    ### Lista de senyales que queremos incluir (tercera columna del dat):
    Signals = []
    Signals.append('HXX_400_50_4mm')



    ############# DEFINICION DE PARAMETROS
    ###
    ### Este es el valor de la luminosidad al que se normalizan los plots e.g. 35.9 fb correspondiente a 2016
    lumi = 35.9 # luminosity



    ############# CREACION DE OBJETOS TREE DE GALAPAGO
    ### 
    ### Toman como uno de los inputs la .dat file que queremos meter (en este caso dat/Samples_cern.dat)
    ### y una lista con los nombres de las muestras que queremos incluir de ese dat.
    ### Es importante indicar (cuando proceda) que isdata = 0 para que se aplique la normalizacion
    treeMC = Sample.Tree( fileName = helper.selectSamples(WORKPATH + 'dat/Samples_cern.dat', Backgrounds, 'MC'), name = 'MC', isdata = 0 )
    treeSI = Sample.Tree( fileName = helper.selectSamples(WORKPATH + 'dat/Samples_cern.dat', Signals, 'SI'), name = 'SI', isdata = 0 )


    ############# DEFINICION DE CORTES
    ### Como ejemplo pongo algunos de los cortes basicos de los displaced global muon pais (DMDM)
    ### Hay algunos definidos en include/CutManager, pero desde aqui podeis definir los que querais
    ### La sintaxis es la que se utilizaria con el metodo .Draw de un TTree normal

    cuts = CutManager.CutManager() # Clase CutManager
    haveDMDM = cuts.brackets('nDMDMBase > 0') # Corte 1
    lowPhi = cuts.brackets('fabs(DMDMBase_dPhi[DMDMBase_maxIxy])< 3.14/2.0') # Corte 2
    allcuts = cuts.AddListB([haveDMDM, lowPhi]) # Combinamos ambos en un unico string


    ############# EJEMPLOS:

    ### --> Ejemplo de coger un TH1F con getTH1F():
    ###      Nombre del histograma: 'hBCK' (fondo con treeMC) y 'hSig' (senyal con treeSI)
    ###      Normalizado a la lumi
    ###      Variable que ploteamos: Distancia transversa del vertice del primer DMDM candidate
    ###      En 20 bins, de 0 a 10 (nos pintara overflow tambien)
    ###      options: Nunca lo uso
    ###      xlabel: El label del eje x
    fOut = TFile(outputFolder + "/output.root", "RECREATE")
    hBCK = treeMC.getTH1F(lumi = lumi, name = 'hBCK', var = 'DMDMBase_Lxy[0]', nbin = 20, xmin = 0, xmax = 10, cut = allcuts, options = '', xlabel = 'Dimuon #mu#mu L_{xy} (cm)')
    hSI = treeSI.getTH1F(lumi = lumi, name = 'hSI', var = 'DMDMBase_Lxy[0]', nbin = 20, xmin = 0, xmax = 10, cut = allcuts, options = '', xlabel = 'Dimuon #mu#mu L_{xy} (cm)')
    

    plot = Canvas.Canvas('hist_DMDM_Lxy', 'png', 0.55, 0.8, 0.9, 0.9, 1)
    plot.addHisto(hBCK, 'hist', '', 'l', '', 1, 0)
    plot.addHisto(hSI, 'hist, same', '', 'l', '', 1, 1)
    plot.save(legend = True, isData = False, log = True, lumi = lumi, labelx = '', outputDir = outputFolder)


    ### --> Ejemplo de sacar la sensitivity y las cumulativas
    Sensitivity, cumSig, cumBkg = computeSensitivity(hBCK, hSI, r.kGreen)

    plot_sensitivity = Canvas.Canvas('hist_sensitivity', 'png', 0.55, 0.8, 0.9, 0.9, 1)
    plot_sensitivity.addHisto(Sensitivity, 'hist', '', 'l', '', 1, 0)
    plot_sensitivity.save(legend = False, isData = False, log = False, lumi = lumi, labelx = '', outputDir = outputFolder)

    maxValue = max([cumBkg.GetMaximum(), cumSig.GetMaximum()])
    cumBkg.SetMaximum(10.0*maxValue)
    cumBkg.SetMinimum(0.1)
    plot_yields = Canvas.Canvas('hist_yields', 'png', 0.55, 0.8, 0.9, 0.9, 1)
    plot_yields.addHisto(cumBkg, 'hist', 'Background', 'l', '', 1, 0)
    plot_yields.addHisto(cumSig, 'hist, same', 'Signal', 'l', '', 1, 1)
    plot_yields.save(legend = True, isData = False, log = True, lumi = lumi, labelx = '', outputDir = outputFolder)

    fOut.cd()
    hBCK.Write()
    hSI.Write()

    Sensitivity.Write()
    cumSig.Write()
    cumBkg.Write()
