##################################################################################################
#                                                                                                #    
#                                      ```...`                                                   #    
#                                   `......-::::.`...---.`                                       #    
#                                  ``     `.-::-.....---:::-`                                    #    
#                                           ..`      `.--:://-                                   #    
#                                     `//:`            .-:::/+:                                  #    
#                                    +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

import include.Sample as Sample
import include.helper as helper
import include.Canvas as Canvas


####################################### 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 ####################################

WORKPATH = os.path.abspath('./') + '/'

##################################### FUNCTION DEFINITION ########################################

    
def makePlot(lumi, var, name, xlabel, logx, treeMC, treeSI = False, treeDATA = False, LLlabel = False, normed = False):

    hMCS = treeMC.getLoopStack('histMC_'+name, var, xlabel)

    luminosity = lumi

    ### MC total histogram
    f = 0
    for _i, _h in enumerate(hMCS.GetHists()):
        if not f: hMC = copy.deepcopy(_h)
        else: hMC.Add(_h, 1.)
        f = 1
    hMC.SetMarkerStyle(20) # Auxiliar to save the ratio correctly 

    ### Signal histograms
    if treeSI:
        hSIS = treeSI.getLoopStack('histSI_'+name, var, xlabel)

        s_histos = []
        for _i, _h in enumerate(hSIS.GetHists()):
            s_histos.append(copy.deepcopy(_h))

    ### Data histogram
    if treeDATA:
        hDATA = treeDATA.getLoopTH1F('histDATA_'+name, var, xlabel)
        hDATA.SetMarkerStyle(20)
        hDATA.SetMarkerSize(0.8)
        hDATA.Scale(hMC.Integral()/hDATA.Integral())

    # Normalization
    if normed:
        hMC.Scale(1.0/hMC.Integral())
        for _h in s_histos: _h.Scale(1.0/_h.Integral())
        luminosity = 1.0

    ### Get maximum
    maxValMC = hMC.GetMaximum()
    maxValSI = 0 if not treeSI else max([s_histos[i].GetMaximum() for i in range(0, len(s_histos))])
    maxValDATA = 0 if not treeDATA else hDATA.GetMaximum()
    maxVal = max([maxValMC, maxValSI, maxValDATA])

    ### Set Maximum
    if not logx:
        hMCS.SetMaximum(1.3*maxVal)
        hMCS.SetMinimum(0.0)
        hMC.SetMaximum(1.3*maxVal)
        hMC.SetMinimum(0.0)
        if treeSI:
            for _h in s_histos: 
                _h.SetMaximum(1.3*maxVal)
                _h.SetMinimum(0.0)
        if treeDATA:
            hDATA.SetMaximum(1.3*maxVal)
            hDATA.SetMinimum(0.0)
    else:
        hMCS.SetMaximum(10.0*maxVal)
        hMCS.SetMinimum(0.1)
        hMC.SetMaximum(10.0*maxVal)
        hMC.SetMinimum(0.1)
        if treeSI:
            for _h in s_histos: 
                _h.SetMaximum(10.0*maxVal)
                _h.SetMinimum(0.1)
        if treeDATA:
            hDATA.SetMaximum(10.0*maxVal)
            hDATA.SetMinimum(0.1)


    ### Canvas object
    plot = Canvas.Canvas('hist_'+name, 'png', 0.34, 0.68, 0.9, 0.9)

    if normed: plot.addHisto(hMC, 'HIST', '', 'l', r.kBlue, 1, 0) # Background
    else: plot.addStack(hMCS, 'HIST', 1, 0) # Background

    if treeSI:
        for _h in s_histos:
            plot.addHisto(_h, 'HIST, SAME', '', 'l', _h.GetFillColor(), 1, 3) # Signal

    if treeDATA:
        plot.addHisto(hDATA, 'P, SAME', '', 'p', r.kBlack, 1, 4)


    ### Dilepton banner
    if LLlabel == 'EE':
        plot.addLatex(0.17, 0.86, 'e^{+}e^{-} channel')
    if LLlabel == 'MM':
        plot.addLatex(0.17, 0.86, '#mu^{+}#mu^{-} channel')

    ### Save it
    #plot.save(1, 0, logx, luminosity, '', outputDir = WORKPATH + 'plots/')
    plot.saveRatio(1, 0, logx, luminosity, hDATA, hMC, label="Data/MC", outputDir = 'newplots/')


def makeSensitivity(lumi, treeSI, treeMC, var, name, xlabel, logx, LLlabel = False, normed = False):

    hMCS = treeMC.getLoopTH1F('histMC_'+name, var, xlabel)
    hSIS = treeSI.getLoopStack('histSI_'+name, var, xlabel)

    luminosity = lumi

    ### Signal histograms
    s_histos = []
    for _i, _h in enumerate(hSIS.GetHists()):
        s_histos.append(copy.deepcopy(_h))

    ### Get significance values
    plot = Canvas.Canvas('sensitivity_'+name, 'png', 0.5, 0.7, 0.9, 0.9, 1)
    sens = []

    for _i, _h in enumerate(s_histos):
        gh = r.TH1F('gr', '', hMCS.GetNbinsX(), hMCS.GetXaxis().GetXmin(), hMCS.GetXaxis().GetXmax())
        for n in range(1, hMCS.GetNbinsX() + 1):
            sig_value = 0.0
            back_value = 0.0
            value = 0.0
            for j in range(n, hMCS.GetNbinsX() + 1):
                sig_value = sig_value + _h.GetBinContent(j)
                back_value = back_value + hMCS.GetBinContent(j)
            if sig_value + back_value == 0:
                value = 0.0
            else:
                value = sig_value/math.sqrt(sig_value + back_value)

            gh.SetBinContent(n, value)

        gh.SetTitle(_h.GetTitle())
        gh.SetLineWidth(2)
        gh.GetYaxis().SetTitle('S/#sqrt{S + B}')
        gh.GetXaxis().SetTitle(xlabel)
        gh.SetLineColor(_h.GetFillColor())
        gh1 = gh.Clone()
        gh2 = copy.deepcopy(gh1)
        sens.append(gh2)

    print([sens[i].GetMaximum() for i in range(0, len(sens))])
    maximum = max([sens[i].GetMaximum() for i in range(0, len(sens))])
    print('maximum: ', maximum)


    for _s,gh2 in enumerate(sens):

        gh2.SetMaximum(1.3*maximum)

        if _s == 0:
            plot.addHisto(gh2, 'l', gh2.GetTitle(), 'l', gh2.GetLineColor(), 1, 1)
        else:
            plot.addHisto(gh2, 'l, same', gh2.GetTitle(), 'l', gh2.GetLineColor(), 1, 1)
            
    ### Dilepton banner
    if LLlabel == 'EE':
        plot.addLatex(0.17, 0.86, 'e^{+}e^{-} channel')
    if LLlabel == 'MM':
        plot.addLatex(0.17, 0.86, '#mu^{+}#mu^{-} channel')

    ### Save it
    plot.save(1, 0, logx, luminosity, '', outputDir = WORKPATH + 'sensitivity/')



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

    parser = optparse.OptionParser(usage='usage: %prog [opts] FilenameWithSamples', version='%prog 1.0')
    parser.add_option('-o', '--output', action='store', type=str, dest='outputFile', default='outputHistos_prueba.root', help='the output file. default \'outputHistograms.root\'')
    parser.add_option('-l', '--avoidLoop', action='store_true', dest='avoidLoop', help='It determines if the histograms need to be done again')

    (opts, args) = parser.parse_args()


    ############# Initialize the output root file with the histos
    if not opts.avoidLoop:
        output = r.TFile( WORKPATH + opts.outputFile, "RECREATE" )
        output.Close()


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

    ############# Background definition
    Backgrounds = []
    Backgrounds.append('WJetsToLNu') 
    Backgrounds.append('WW') 
    Backgrounds.append('WZ') 
    Backgrounds.append('ZZ') 
    Backgrounds.append('TTJets_DiLept') 
    Backgrounds.append('DYJetsToLL_M-50') 
    Backgrounds.append('DYJetsToLL_M-10to50') 
    Backgrounds.append('QCD_Pt-30to40') 
    Backgrounds.append('QCD_Pt-40toInf') 


    ############# Signal definition
    Signal = []
#    Signal.append('ScalarBosons_1000_350_350')
#    Signal.append('ScalarBosons_1000_350_35')
#    Signal.append('ScalarBosons_1000_150_100')
#    Signal.append('ScalarBosons_1000_150_10')
#    Signal.append('ScalarBosons_400_150_400')
#    Signal.append('ScalarBosons_400_150_40')
    Signal.append('DisplacedSUSY_1500_494_160')
    Signal.append('DisplacedSUSY_1000_148_60')
    Signal.append('DisplacedSUSY_350_148_173')
    Signal.append('DisplacedSUSY_120_48_165')


    ############## Data definition
    Data = []
    Data.append('DoubleMuon-Run2016B')


    ############# Parameter definition
    lumi = 36.773 # luminosity

    ############# Tree creation
    treeMC = Sample.Tree(helper.selectSamples(WORKPATH + 'dat/MC.dat', Backgrounds, 'MC'), 'MC', 0, WORKPATH + opts.outputFile)
    treeSI = Sample.Tree(helper.selectSamples(WORKPATH + 'dat/SI.dat', Signal, 'SI'), 'SI', 0, WORKPATH + opts.outputFile)
    treeDATA = Sample.Tree(helper.selectSamples(WORKPATH + 'dat/DATA.dat', Data, 'DATA'), 'DATA', 1, WORKPATH + opts.outputFile)

    ############# Tree loop
    if not opts.avoidLoop:
        treeMC.Loop(lumi, False, False)
        treeSI.Loop(lumi, False, False)
        treeDATA.Loop(lumi, False, False)

    
    makeSensitivity(lumi, treeSI, treeMC, 'SR_EEsel_minIxy', 'SR_EEsel_minIxy', '|d_{0}|/#sigma_{d}', 0, 'EE')
    makeSensitivity(lumi, treeSI, treeMC, 'SR_MMsel_minIxy', 'SR_MMsel_minIxy', '|d_{0}|/#sigma_{d}', 0, 'MM')
    makeSensitivity(lumi, treeSI, treeMC, 'SR_EEsel_minIxy', 'SR_EEsel_minIxy', '|d_{0}|/#sigma_{d}', 0, 'EE')
    


    makePlot(lumi, 'nMM', 'nMM', 'Number of MM', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'SR_MMsel_minIxy', 'SR_MMsel_minIxy', '|d_{0}|/#sigma_{d}', 1, treeMC, treeSI, treeDATA, 'MM')
    
    makePlot(lumi, 'SR_MMsel_invMass', 'SR_MMsel_invMass', 'Mass (GeV/c^{2})', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'SR_MMsel_Chi2', 'SR_MMsel_Chi2', 'Vertex #Chi^{2}', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'SR_MMsel_dPhi', 'SR_MMsel_dPhi', 'Collinearity |#Delta#Phi|', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'SR_MMsel_leadingPt', 'SR_MMsel_leadingPt', 'Leading p_{T} (GeV/c)', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'SR_MMsel_subleadingPt', 'SR_MMsel_subleadingPt', 'Subleading p_{T} (GeV/c)', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'SR_MMsel_cosAlpha', 'SR_MMsel_cosAlpha', 'cos(#alpha)', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'SR_MMsel_Ptll', 'SR_MMsel_Ptll', 'p_{T}(ll)', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'CR_MMsel_minIxy', 'CR_MMsel_minIxy', '|d_{0}|/#sigma_{d}', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'CR_MMsel_invMass', 'CR_MMsel_invMass', 'Mass (GeV/c^{2})', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'CR_MMsel_Chi2', 'CR_MMsel_Chi2', 'Vertex #Chi^{2}', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'CR_MMsel_dPhi', 'CR_MMsel_dPhi', 'Collinearity |#Delta#Phi|', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'CR_MMsel_leadingPt', 'CR_MMsel_leadingPt', 'Leading p_{T} (GeV/c)', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'CR_MMsel_subleadingPt', 'CR_MMsel_subleadingPt', 'Subleading p_{T} (GeV/c)', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'CR_MMsel_cosAlpha', 'CR_MMsel_cosAlpha', 'cos(#alpha)', 1, treeMC, treeSI, treeDATA, 'MM')
    makePlot(lumi, 'CR_MMsel_Ptll', 'CR_MMsel_Ptll', 'p_{T}(ll)', 1, treeMC, treeSI, treeDATA, 'MM')
   


