#!/usr/bin/env python

# contourPlotterExample.py #################
#
# Example for using contourPlotter.py
#
# See README for details
#
# By: Larry Lee - Dec 2017

import ROOT
import os

import contourPlotter

ROOT.gROOT.SetStyle( 'ATLAS' )

def setup_plot():

    plot = contourPlotter.contourPlotter("contourPlotterExample",800,600)
    plot.ATLASlabel = "ATLAS"
    plot.ATLASlabel2 = "Internal"
    plot.processLabel = ""
    plot.limitCaveatLabel = "All limits at 95% CL"
    plot.lumiLabel = "#sqrt{s}=13 TeV, 139 fb^{-1}"#, All limits at 95% CL"
    plot.drawAxes( [100,400,650,1500] )
    return plot

tW1L_nominal = ROOT.TFile( 'split_norm_backup/mamh_Unified.root' )
tW2L_nominal = ROOT.TFile( 'twoL/tw.root' )
#tWComb_nominal = ROOT.TFile( 'mamh_Comb.root' )
tWComb_nominal = ROOT.TFile( 'mamh_all.root' )
tW1L_dmtt = ROOT.TFile( 'split_norm_backup/mamh_dmtt.root' )
tW2L_dmtt = ROOT.TFile( 'twoL/dmtt.root' )
tWComb_dmtt = ROOT.TFile( 'mamh_Comb_DMtt.root' )

xsec_up = ROOT.TFile( 'xsec_up.root' )
xsec_down = ROOT.TFile( 'xsec_down.root' )

plot = setup_plot()
dodmtt = True
if not dodmtt:
    #plot.drawOneSigmaBand(  tWComb_nominal.Get("Band_1s_0")   )
    #plot.drawShadedRegion( tWComb_nominal.Get("Obs_0"), alpha = 0.1, title="tW1L + tW2L" )
    plot.drawOneSigmaBand(  tWComb_nominal.Get("Band_1s_0") )
    plot.drawTwoSigmaBand(  tWComb_nominal.Get("Band_2s_0") )    
    plot.drawExpected( tWComb_nominal.Get("Exp_0"), color=ROOT.kGreen - 1, lineWidth = 2, title = "DMt 1L+2L" )
    plot.drawObserved( tWComb_nominal.Get("Obs_0"), color=ROOT.kGreen - 1, title = "DMt 1L+2L" )

    plot.drawExpected( tW1L_nominal.Get("Exp_0"), color=ROOT.kViolet - 1, lineWidth = 2, title = "DMt 1L" )
    plot.drawObserved( tW1L_nominal.Get("Obs_0"), color=ROOT.kViolet - 1, title="DMt 1L" )
    #plot.drawShadedRegion( tW1L_nominal.Get("Obs_0"), color=ROOT.kAzure - 1, alpha = 0.1, title="tW1L" )

    plot.drawExpected( tW2L_nominal.Get("Exp_0"), color=ROOT.kPink + 9, lineWidth = 2, title = "DMt 2L" )
    plot.drawObserved( tW2L_nominal.Get("Obs_0"), color=ROOT.kPink + 9, title = "DMt 2L" )
    #plot.drawShadedRegion( tW2L_nominal.Get("Obs_0"), color=ROOT.kOrange - 2, alpha = 0.1, title="tW2L" )
if dodmtt:
    plot.drawExpected( tWComb_dmtt.Get("Exp_0"), color=ROOT.kGreen - 1, lineWidth = 2, title = "DMt+t#bar{t} 1L+2L" )
    plot.drawObserved( tWComb_dmtt.Get("Obs_0"), color=ROOT.kGreen - 1, title = "DMt+t#bar{t} 1L+2L" )
    plot.drawOneSigmaBand(  tWComb_dmtt.Get("Band_1s_0") )
    plot.drawTwoSigmaBand(  tWComb_dmtt.Get("Band_2s_0") )

    plot.drawExpected( tW1L_dmtt.Get("Exp_0"), color=ROOT.kViolet - 1, lineWidth = 2, title = "DMt+t#bar{t} 1L" )
    plot.drawObserved( tW1L_dmtt.Get("Obs_0"), color=ROOT.kViolet - 1, title = "DMt+t#bar{t} 1L" )

    plot.drawExpected( tW2L_dmtt.Get("Exp_0"), color=ROOT.kPink + 9, lineWidth = 2, title = "DMt+t#bar{t} 2L" )
    plot.drawObserved( tW2L_dmtt.Get("Obs_0"), color=ROOT.kPink + 9, title = "DMt+t#bar{t} 2L" )

## Axis Labels

plot.setXAxisLabel( "m_{a} [GeV]" )
plot.setYAxisLabel( "m_{H^{#pm}} [GeV]"  )

plot.createLegend( shape=(0.70,0.40,0.90,0.50) ).Draw()

	#if drawTheorySysts:
	#	plot.drawTheoryUncertaintyCurve( f.Get("Obs_0_Up") )
	#	plot.drawTheoryUncertaintyCurve( f.Get("Obs_0_Down") )
		# coordinate in NDC
	#	plot.drawTheoryLegendLines( xyCoord=(0.234,0.6625), length=0.057 )

plot.drawATLASlabel( 0.65, 0.83 )
plot.decorateCanvas(  )

expHist = ROOT.TH1D( 'exp', 'exp', 1, -1, 1 )
obsHist = ROOT.TH1D( 'obs', 'obs', 1, -1, 1 )
theoHist = ROOT.TH1D( 'theo', 'theo', 1, -1, 1 )
uncertHist = ROOT.TH1D( 'uncert', 'uncert', 1, -1, 1 )

expHist.SetLineStyle( 7 )
expHist.SetLineColor( 1 )
obsHist.SetLineColor( 1 )
theoHist.SetLineColor( ROOT.kGreen - 1 )
theoHist.SetLineStyle( 3 )
uncertHist.SetFillColorAlpha( ROOT.TColor.GetColor("#ffd700"), 0.5 )
uncertHist.SetFillStyle( 1001 )
uncertHist.SetLineColorAlpha( ROOT.kGray,0.5 )

expHist.SetFillColor( 0 )
obsHist.SetFillColor( 0 )
theoHist.SetFillColor( 0 )

expHist.Draw( "HIST SAME" )
obsHist.Draw( "HIST SAME" )
theoHist.Draw( "HIST SAME" )
uncertHist.Draw("HIST SAME" )

latexObject = ROOT.TLatex()
latexObject.SetTextFont(42)
latexObject.SetTextSize(0.035)
#latexObject.DrawLatexNDC(0.16,0.69,"SR(tW_{2L})")
latexObject.DrawLatexNDC(0.705,0.35,"2HDM+a, Dirac DM")
latexObject.SetTextSize(0.032)
latexObject.DrawLatexNDC(0.728,0.30,"m_{#chi} = 10 GeV, g_{#chi} = 1")
latexObject.DrawLatexNDC(0.724,0.25,"sin#theta = 1/#sqrt{2}, tan#beta = 1") #mamh grid
#latexObject.DrawLatexNDC(0.64,0.73,"sin#theta = 1/#sqrt{2}, m_{a} = 250 GeV") #mhtb grid
latexObject.DrawLatexNDC(0.777,0.20,"m_{H^{#pm}} = m_{H} = m_{A}")

doTheory = False

plot.expObsLegend( (0.70,0.52,0.90,0.69), expHist, obsHist, theoHist, uncertHist, doTheory )

if doTheory:

    plot.drawTheoryUncertaintyCurve( xsec_up.Get("Obs_0"), color=ROOT.kGreen - 1 )
    plot.drawTheoryUncertaintyCurve( xsec_down.Get("Obs_0"), color=ROOT.kGreen - 1 )

plot.writePlot( )