#!/usr/bin/env pythono
import os, shlex
import subprocess
import glob
import shutil
import ROOT
from collections import OrderedDict
import CMS_lumi, tdrstyle
from optparse import OptionParser
parser = OptionParser()


parser.add_option("--inputPaths", type="string", dest="inputPaths", help="Full path of where combine results are seperated by :")
parser.add_option("--outPath", type="string", dest="outPath", help="Full path to store plot ")
parser.add_option("--spinCase", type="string", dest="spinCase", help="spin case")
parser.add_option("--ylabel", type="string", dest="ylabel", help="plot label")
parser.add_option("--scanlabels", type="string", dest="scanlabels", help="scan labels seperated by :")
parser.add_option("--log", action='store_true',dest="logy", default=False)
parser.add_option("--min", type="float", dest="miny", help="min y", default = 0.0)
parser.add_option("--max", type="float", dest="maxy", help="max y", default = 0.0)
parser.add_option("--maxX", type="float", dest="maxX", help="max y", default = -1.0)
parser.add_option("--minX", type="float", dest="minX", help="max y", default = -1.0)
parser.add_option("--unblind", action='store_true',dest="unblind", default=False)
parser.add_option("--private", action='store_true',dest="private", default=False)
parser.add_option("--preliminary", action='store_true',dest="preliminary", default=False)
(options, args) = parser.parse_args()
inputPaths = options.inputPaths
scanlabels = options.scanlabels
outPath = options.outPath
spinCase = options.spinCase
ylabel = options.ylabel
logy = options.logy
miny = options.miny
maxy = options.maxy
maxX = options.maxX
minX = options.minX
unblind = options.unblind
private=options.private
preliminary=options.preliminary
def getLimits(file_name): 
    file = ROOT.TFile(file_name)
    tree = file.Get("limit")
    limits = [ ]
    for quantile in tree:
        limits.append(tree.limit)
    return limits[:6]

def plotUpperLimits(resultdicts, scanlabels, xlabel, ylabel,outpath):
    # see CMS plot guidelines: https://ghm.web.cern.ch/ghm/plots/
    tgraphs = []
    tgraphs_observed = []
    up2s = [ ]
    # colors =  [ROOT.kBlack, # combined
    #            ROOT.kBlue, # 4tau
    #            ROOT.kRed, #1l3tau
    #            41, #2lss
    #            ROOT.kGray+1, # 2l2tau
    #            ROOT.kMagenta+2, #3l
    #            ROOT.kOrange+7, #3l1tau
    #            ROOT.kCyan+2 #4l
    # ]

    colors =  [214,
               628,
               418,
               921,
               222,
               7,
               797,
               413,
               798
    ]
    colors += 10*[921]
    markers = [20,
               21,
               22,
               23,
               24,
               25,
               26,
               32,
               27,
               33,
               28,
               34,
               29,
               30
    ]
    # markers = [20,
    #            21,
    #            22,
    #            23,
    #            24,
    #            25,
    #            26,
    #            32
    # ]
    if len(resultdicts)==2:
        markers=[23,24]
    for k,r in enumerate(resultdicts):
        values = []
        for key in r.keys():
            values.append(int(key))
        N = len(values)
        values = sorted(values)
        median = ROOT.TGraph(N)      # median line
        observed = ROOT.TGraph(N)      # observed line
        i = 0
        for v in values:
            limit = r[str(v)]
            print(scanlabels.split(":")[k], v, limit[5], limit[2], limit[0], limit[1], limit[3], limit[4])
            if unblind and len(limit)<6:
                raise Exception("Not enougth entries in limit to unblind")
            up2s.append(limit[2])
            if unblind:
                observed.SetPoint(    i,    values[i], limit[5] ) # observed
            median.SetPoint(    i,    values[i], limit[2] ) # median
            i = i + 1
        tgraphs.append(median)
        if unblind: tgraphs_observed.append(observed)
    W = 800
    H  = 600
    T = 0.08*H
    B = 0.12*H
    L = 0.12*W
    R = 0.04*W
    c = ROOT.TCanvas("c","c",100,100,W,H)
    c.SetFillColor(0)
    c.SetBorderMode(0)
    c.SetFrameFillStyle(0)
    c.SetFrameBorderMode(0)
    c.SetLeftMargin( 1.05*L/W )
    c.SetRightMargin( 1.05*R/W )
    c.SetTopMargin( T/H )
    c.SetBottomMargin( B/H )
    c.SetTickx(0)
    c.SetTicky(0)
    #if not private: c.SetGrid()
    c.cd()
    frame = c.DrawFrame(1.4,0.001, 4.1, 10)
    frame.GetYaxis().CenterTitle()
    frame.GetYaxis().SetTitleSize(0.05)
    frame.GetXaxis().SetTitleSize(0.05)
    frame.GetXaxis().SetLabelFont(43)
    frame.GetYaxis().SetLabelFont(43)
    frame.GetYaxis().SetLabelSize(30)
    frame.GetXaxis().SetLabelSize(30)
    #frame.GetXaxis().SetLabelSize(0.055)
    #frame.GetYaxis().SetLabelSize(0.055)
    frame.GetYaxis().SetTitleOffset(1.15)
    frame.GetXaxis().SetTitleOffset(1.15)
    frame.GetXaxis().SetNdivisions(508)
    frame.GetYaxis().CenterTitle(True)
    frame.GetYaxis().SetTitle("#bf{95% CL upper limit #sigma(pp#rightarrow#Chi#rightarrow HH) [pb]}")
    #frame.GetYaxis().SetTitle("95% upper limit on #sigma (#Chi#rightarrow HH) [pb]")
    frame.GetXaxis().SetTitle(xlabel)
    frame.SetMinimum(miny)
    frame.SetMaximum(max(up2s)*1.05)
    if logy:
        frame.SetMaximum(max(up2s)*15)
    frame.GetXaxis().SetLimits(min(values),max(values))
    if logy:
        c.SetLogy()
    if(maxy>0.): frame.SetMaximum(maxy)
    tgraphs[0].Draw()
    for i, g in enumerate(tgraphs):
        g.SetLineColor(colors[i])
        print scanlabels.split(":")[i], markers[i]
        g.SetMarkerColor(colors[i])
        g.SetMarkerStyle(markers[i])
        g.SetMarkerSize(1.4)
        if not unblind:
            g.SetLineWidth(3)
            g.Draw("PLSame")
        else:
            g.SetLineStyle(2) 
            g.SetLineWidth(2)
            g.Draw("LSame")
    for i, g in enumerate(tgraphs_observed):
        g.SetLineColor(colors[i])
        print scanlabels.split(":")[i], markers[i]
        g.SetMarkerColor(colors[i])
        g.SetMarkerStyle(markers[i])
        g.SetMarkerSize(1.4)
        g.SetLineWidth(3)
        g.Draw("PLSame")
        
    CMS_lumi.lumi_sqrtS = ylabel
    if not (preliminary or private):
        CMS_lumi.lumiTextSize = 1.3*CMS_lumi.lumiTextSize
    CMS_lumi.CMS_lumi(c,0,11)
    #multilep label
    mlep_label_x = 0.44
    mlep_label_y = 0.975
    if preliminary:
        mlep_label_x = 0.54
        mlep_label_y = 0.97
    if private:
        mlep_label_x = 0.58
        mlep_label_y = 0.97
        #mlep_label_x = 0.565
        #mlep_label_y = 0.975
    #mlep_label = "#bf{HH #rightarrow Multilepton (%s)}"%spinCase
    mlep_label = "#bf{HH #rightarrow Multilepton}"
    if private:
        mlep_label = "3\\ell\\text{ +1}\\tau_{h}}"
    #if private:
    #    mlep_label = "3\\ell\\text{ +1}\\tau_{h}\\text{ (%s)}"%spinCase
    mlep = ROOT.TLatex()
    mlep.SetNDC()
    mlep.SetTextAngle(0)
    mlep.SetTextAlign(23)
    mlep.SetTextSize(0.055)
    if preliminary:
        mlep.SetTextSize(0.045)
    if private:
        mlep.SetTextSize(0.045)
    mlep.DrawLatex(mlep_label_x, mlep_label_y, mlep_label)
    spinlab = ROOT.TLatex()
    spinlab.SetNDC()
    spinlab.SetTextAngle(0)
    spinlab.SetTextAlign(23)
    spinlab.SetTextSize(0.055)
    #if not private:
    spinlab.DrawLatex(0.25, 0.25, spinCase)
    #if private:
    #    spinlab.DrawLatex(0.75, 0.5, spinCase)

    obsexp_label_x = 0.55
    obsexp_label_y = 0.69
    obsexp_label ="#bf{Solid (dashed) lines: observed (median expected) limits}"
    if not unblind:
        obsexp_label ="#bf{Solid lines: median expected limits}"
    obsexp = ROOT.TLatex()
    obsexp.SetNDC()
    obsexp.SetTextAngle(0)
    obsexp.SetTextAlign(23)
    obsexp.SetTextSize(0.045)
    obsexp.DrawLatex(obsexp_label_x, obsexp_label_y, obsexp_label)
    ROOT.gPad.SetTicks(1,1)
    frame.Draw('sameaxis')
    if maxX>=0:
        if minX>=0:
            frame.GetXaxis().SetRangeUser(minX,maxX)
        else:
            frame.GetXaxis().SetRangeUser(0,maxX)
    elif minX>=0:
        frame.GetXaxis().SetRangeUser(minX,frame.GetXaxis().GetXmax())
 
    #x1 = 0.459
    x1 = 0.15
    x2 = 0.95
    y2 = 0.9175
    y1 = 0.7
    # if private:
    #     y2 = 0.89
    #     y1 = 0.65
    #     x1 = 0.35
    #     x2 = 0.959
    legend = ROOT.TLegend(x1,y1,x2,y2)
    if private: legend.SetFillStyle(0)
    #legend.SetFillColor(1)
    legend.SetNColumns(3)
    legend.SetBorderSize(0)
    legend.SetTextSize(0.055)
    legend.SetTextFont(42)
    if unblind:
        for i,g in enumerate(tgraphs_observed):
            legend.AddEntry(g, scanlabels.split(":")[i],'PL')
        #legend.AddEntry(tgraphs[0], 'expected','L')
    else:
        for i,g in enumerate(tgraphs):
            legend.AddEntry(g, scanlabels.split(":")[i],'PL')
    legend.Draw()
    print " "
    c.SaveAs(outpath)
    c.SaveAs(outpath.replace('.png','.root'))
    c.SaveAs(outpath.replace('.png','.pdf'))
    c.SaveAs(outpath.replace('.png','.eps'))
    c.Close()
 

# CMS style
CMS_lumi.cmsText = "CMS"
CMS_lumi.extraText = ""
if preliminary:
    CMS_lumi.extraText = "Preliminary"
if private:
    #CMS_lumi.extraText = "Private Work"
    CMS_lumi.extraText = "Work in Progress"
CMS_lumi.cmsTextSize = 0.65
if not (preliminary or private):
    CMS_lumi.cmsTextSize = 1.3*0.65
CMS_lumi.outOfFrame = True
tdrstyle.setTDRStyle()
resultdicts = []
for path in inputPaths.split(":"):
    listproc = glob.glob( "%s/*.root" % path)
    resultdict = {}
    for result in listproc:
        mass = result.split('_')[-1].split('.')[0]
        resultdict[mass]=getLimits(result)
    resultdicts.append(resultdict)
#plotUpperLimits(resultdicts,scanlabels, '#bf{m_{#Chi} [GeV]}', ylabel, outPath)
plotUpperLimits(resultdicts,scanlabels, 'm_{#Chi} [GeV]', ylabel, outPath)
#plotUpperLimits(resultdicts, scanlabels, 'm#Chi [GeV] (%s)'%spinCase, ylabel, outPath)
