"""
Author: Federica Lionetto
Date: November 16th, 2015

Description:
Define some useful functions to be used to accomplish common tasks with ROOT objects.
- LHCbStyle()
- InitHist(hist,title,x,y)
- InitHist2(hist,title,x,y,z)
- InitGraph(graph,title,x,y)
- SetStyleObj(obj,markerColor,lineColor,fillColor,markerSize,markerStyle,lineWidth,lineStyle)
- AdjustOffsets(hist,axis,labelOffset,titleOffset)
- DrawObj(canvas,obj,leg,option,folder)
- DrawOjbCompare(canvas,objList,colorDict,optionDict,leg,folder)
- CreateLegend(objList,labelDict,optionDict,x1,y1,x2,y2)
- CreateStats(hist,option,x1,y1,x2,y2)
- TestStyle()

How to run it:
Include it in the other python scripts.
"""

from array import array

import numpy as np

import ROOT

def LHCbStyle() :
    lhcbFont = 132
    lhcbWidth = 2
    lhcbTSize = 0.06

    ROOT.gROOT.SetStyle('Plain')
    lhcbStyle = ROOT.TStyle('lhcbStyle','LHCb plots style')

    lhcbStyle.SetFillColor(0)
    lhcbStyle.SetFillStyle(1001)
    lhcbStyle.SetFrameFillColor(0)
    lhcbStyle.SetFrameBorderMode(0)
    lhcbStyle.SetPadBorderMode(0)
    lhcbStyle.SetPadColor(0)
    lhcbStyle.SetCanvasBorderMode(0)
    lhcbStyle.SetCanvasColor(0)
    lhcbStyle.SetStatColor(0)
    lhcbStyle.SetLegendBorderSize(0)
    lhcbStyle.SetLegendFont(132)

    lhcbStyle.SetPalette(1)
    colors=array('i',[0,5,7,3,6,2,4,1])
    # colors=np.array([0,5,7,3,6,2,4,1],dtype=np.int_)
    print colors
    print len(colors)
    lhcbStyle.SetPalette(len(colors),colors)

    lhcbStyle.SetPaperSize(20,26)
    lhcbStyle.SetPadTopMargin(0.05)
    lhcbStyle.SetPadRightMargin(0.05)
    lhcbStyle.SetPadBottomMargin(0.16)
    lhcbStyle.SetPadLeftMargin(0.14)

    lhcbStyle.SetTextFont(lhcbFont)
    lhcbStyle.SetTextSize(lhcbTSize)
    lhcbStyle.SetLabelFont(lhcbFont,"x")
    lhcbStyle.SetLabelFont(lhcbFont,"y")
    lhcbStyle.SetLabelFont(lhcbFont,"z")
    lhcbStyle.SetLabelSize(lhcbTSize,"x")
    lhcbStyle.SetLabelSize(lhcbTSize,"y")
    lhcbStyle.SetLabelSize(lhcbTSize,"z")
    lhcbStyle.SetTitleFont(lhcbFont)
    lhcbStyle.SetTitleFont(lhcbFont,"x")
    lhcbStyle.SetTitleFont(lhcbFont,"y")
    lhcbStyle.SetTitleFont(lhcbFont,"z")
    lhcbStyle.SetTitleSize(1.2*lhcbTSize,"x")
    lhcbStyle.SetTitleSize(1.2*lhcbTSize,"y")
    lhcbStyle.SetTitleSize(1.2*lhcbTSize,"z")

    lhcbStyle.SetLineWidth(lhcbWidth)
    lhcbStyle.SetFrameLineWidth(lhcbWidth)
    lhcbStyle.SetHistLineWidth(lhcbWidth)
    lhcbStyle.SetFuncWidth(lhcbWidth)
    lhcbStyle.SetGridWidth(lhcbWidth)
    lhcbStyle.SetLineStyleString(2,"[12 12]")
    lhcbStyle.SetMarkerStyle(20)
    lhcbStyle.SetMarkerSize(1.0)

    lhcbStyle.SetLabelOffset(0.010,"X")
    lhcbStyle.SetLabelOffset(0.010,"Y")

    lhcbStyle.SetOptStat(0)
    lhcbStyle.SetStatFormat("6.3g")
    lhcbStyle.SetOptTitle(0)
    lhcbStyle.SetOptFit(0)
    lhcbStyle.SetTitleOffset(0.95,"X")
    lhcbStyle.SetTitleOffset(0.95,"Y")
    lhcbStyle.SetTitleOffset(1.2,"Z")
    lhcbStyle.SetTitleFillColor(0)
    lhcbStyle.SetTitleStyle(0)
    lhcbStyle.SetTitleBorderSize(0)
    lhcbStyle.SetTitleFont(lhcbFont,"title")
    lhcbStyle.SetTitleX(0.0)
    lhcbStyle.SetTitleY(1.0)
    lhcbStyle.SetTitleW(1.0)
    lhcbStyle.SetTitleH(0.05)

    lhcbStyle.SetStatBorderSize(0)
    lhcbStyle.SetStatFont(lhcbFont)
    lhcbStyle.SetStatFontSize(0.05)
    lhcbStyle.SetStatX(0.9)
    lhcbStyle.SetStatY(0.9)
    lhcbStyle.SetStatW(0.25)
    lhcbStyle.SetStatH(0.15)

    lhcbStyle.SetPadTickX(1)
    lhcbStyle.SetPadTickY(1)

    lhcbStyle.SetNdivisions(505,"x")
    lhcbStyle.SetNdivisions(510,"y")

    ROOT.gROOT.SetStyle("lhcbStyle")
    ROOT.gROOT.ForceStyle()

    lhcbName = ROOT.TPaveText(ROOT.gStyle.GetPadLeftMargin() + 0.05,0.87 - ROOT.gStyle.GetPadTopMargin(),ROOT.gStyle.GetPadLeftMargin() + 0.20,0.95 - ROOT.gStyle.GetPadTopMargin(),'BRNDC')
    lhcbName.AddText("LHCb")
    lhcbName.SetFillColor(0)
    lhcbName.SetTextAlign(12)
    lhcbName.SetBorderSize(0)

    lhcbLabel = ROOT.TText()
    lhcbLabel.SetTextFont(lhcbFont)
    lhcbLabel.SetTextColor(1)
    lhcbLabel.SetTextSize(lhcbTSize)
    lhcbLabel.SetTextAlign(12)

    lhcbLatex = ROOT.TLatex()
    lhcbLatex.SetTextFont(lhcbFont)
    lhcbLatex.SetTextColor(1)
    lhcbLatex.SetTextSize(lhcbTSize)
    lhcbLatex.SetTextAlign(12)

    print '-------------------------'
    print 'Set LHCb Style - Feb 2012'
    print '-------------------------'

    return lhcbStyle

def InitCanvas(canvas,top=0.15,bottom=None,left=None,right=None) :
    if top is not None :
        canvas.SetTopMargin(top)
    if bottom is not None :
        canvas.SetBottomMargin(bottom)
    if left is not None :
        canvas.SetLeftMargin(left)
    # Set the right margin to 0.19 for COLZ plots.
    if right is not None :
        canvas.SetRightMargin(right)
    return

def InitHist(hist,title='',x='',y='') :
    """
    Initialise 1D histogram <hist> with title <title>, x axis title <x>, and y axis title <y>.
    If <y> is not given, a y axis title based on the number of entries per bin is set.
    """
    hist.SetTitle(title)
    hist.GetXaxis().SetTitle(x)
    if (y=='') :
        hist.GetYaxis().SetTitle('Entries / {0}'.format(hist.GetBinWidth(1)))
    else :
        hist.GetYaxis().SetTitle(y)
    return

def InitHist2(hist,title='',x='',y='',z='') :
    """
    Initialise 2D histogram <hist> with title <title>, x axis title <x>, y axis title <y>, and z axis title <z>.
    If <z> is not given, a z axis title based on the number of entries per bin is set.
    """
    hist.SetTitle(title)
    hist.GetXaxis().SetTitle(x)
    hist.GetYaxis().SetTitle(y)
    if (z=='') :
        hist.GetZaxis().SetTitle('Entries / ({0}*{1})'.format(hist.GetXaxis().GetBinWidth(1),hist.GetYaxis().GetBinWidth(1)))
    else :
        hist.GetZaxis().SetTitle(z)
    return

def InitGraph(graph,title='',x='',y='') :
    """
    Initialise graph <graph> with title <title>, x axis title <x>, and y axis title <y>.
    """
    graph.SetTitle(title)
    graph.GetXaxis().SetTitle(x)
    graph.GetYaxis().SetTitle(y)
    return

def SetStyleObj(obj,markerColor=ROOT.kBlack,lineColor=ROOT.kMagenta,fillColor=None,markerSize=1,markerStyle=20,lineWidth=2,lineStyle=1) :
    """
    Set the style of object <obj>:
    - marker color
    - line color
    - fill color
    - marker size
    - marker style
    - line width
    - line style
    A default value exists for the last four.
    """
    obj.SetMarkerColor(markerColor)
    obj.SetLineColor(lineColor)
    obj.SetFillStyle(1001)
    if fillColor is not None :
        obj.SetFillColor(fillColor)
    obj.SetMarkerSize(markerSize)
    obj.SetMarkerStyle(markerStyle)
    obj.SetLineWidth(lineWidth)
    obj.SetLineStyle(lineStyle)
    return

def AdjustOffsets(hist,axis,labelOffset=0.005,titleOffset=0.85) :
    """
    Adjust the offsets of histogram <hist> along the axis <axis>, by setting the label offset <labelOffset> and the title offset <titleOffset>.
    The argument <axis> can be 'X', 'Y', or 'Z'.
    The default values are those used for a TH2 object one wants to draw with the COLZ option.
    """
    hist.SetLabelOffset(labelOffset,axis)
    hist.SetTitleOffset(titleOffset,axis)
    return

def DrawObj(canvas,obj,leg=None,option='',folder='.') :
    """
    Draw object <obj> in canvas <canvas>, together with legend <leg> (if provided), with option <option> and save the result in a pdf file having the same name as the canvas, located in the folder <folder>.
    """
    name=canvas.GetName()
    canvas.cd()
    obj.Draw(option)
    if ('A' in option) :
        canvas.SetTickx(1)
        canvas.SetTicky(1)
    if (leg is not None) :
        leg.Draw()
    canvas.Update()
    canvas.Write()
    canvas.SaveAs(folder+'/'+name+'.pdf')
    canvas.Close()
    return

def DrawObjCompare(canvas,objList,colorDict,optionDict,leg,folder='.') :
    """
    Draw list of objects <objList> in canvas <canvas>, together with legend <leg>, with color <colorDict> and option <optionDict> and save the result in a pdf file having the same name as the canvas, located in the folder <folder>.
    This allows to draw histograms, graphs, and functions in the same canvas.
    However, since graphs do not have a 'GetMaximum()' method, the object to be drawn first is searched for among histograms and functions only (whose names start with 'h' and 'f', respectively). The remaining objects are drawn after.
    """
    name=canvas.GetName()
    canvas.cd()
    maxList = []
    for obj in objList :
        obj.SetMarkerColor(colorDict[obj.GetName()])
        obj.SetLineColor(colorDict[obj.GetName()])
        obj.SetFillColor(colorDict[obj.GetName()])
        if (obj.GetName().startswith('h') or obj.GetName().startswith('f')) :
            maxList.append(obj.GetMaximum())

    print maxList

    firstDraw = maxList.index(max(maxList))
    print firstDraw

    objList[firstDraw].Draw(optionDict[objList[firstDraw].GetName()])
    for obj in objList :
        if (obj != objList[firstDraw]) :
            obj.Draw(optionDict[obj.GetName()]+'SAME')
    leg.Draw()
    canvas.Update()
    canvas.Write()
    canvas.SaveAs(folder+'/'+name+'.pdf')
    canvas.Close()
    return

def CreateLegend(objList,labelDict,optionDict,x1=0.64,y1=0.59,x2=0.94,y2=0.89) :
    """
    Create legend <leg> starting from list of objects <objList>, dictionary of labels <labelDict>, and dictionary of options <optionDict>.
    The legend is displayed in the position identified by <x1>, <y1>, <x2>, and <y2>.
    If the option 'w' is specified, the fill color is set to white.
    The text size is set to '0.05'.
    """
    leg = ROOT.TLegend(x1,y1,x2,y2)
    leg.SetFillColor(ROOT.kWhite)
    for obj in objList :
        label = labelDict[obj.GetName()]
        option = optionDict[obj.GetName()]
        leg.AddEntry(obj,label,option)
    leg.SetTextSize(0.05)
    return leg

def CreateStats(hist,option,x1=0.64,y1=0.59,x2=0.94,y2=0.89) :
    """
    Create stats <stats> of histogram <hist> depending on option <option>.
    The stats are displayed in the position identified by <x1>, <y1>, <x2>, and <y2>.
    """
    stats = ROOT.TPaveStats(x1,y1,x2,y2,'brNDC')
    stats.SetName('stats')
    stats.SetBorderSize(1)
    stats.SetFillColor(0)
    stats.SetTextAlign(12)
    text = stats.AddText('Entries = {0}'.format(hist.GetEntries()))
    if ('I' in option) :
        text = stats.AddText('Integral = {0}'.format(hist.Integral()))
    N = hist.GetDimension()
    if (N == 1) :
        text = stats.AddText('Mean  = {0}'.format(hist.GetMean()));
        text = stats.AddText('RMS   = {0}'.format(hist.GetRMS()));
        if ('S' in option) :
            text = stats.AddText('Skewness = {0}'.format(hist.GetSkewness()))
        if ('K' in option) :
            text = stats.AddText('Kurtosis = {0}'.format(hist.GetKurtosis()))
        if ('U' in option) :
            text = stats.AddText('Underflow = {0}'.format(hist.GetBinContent(0)))
        if ('O' in option) :
            over = hist.GetNbinsX()+1
            text = stats.AddText('Overflow  = {0}'.format(hist.GetBinContent(over)))
    else :
        text = stats.AddText('Mean x = {0}'.format(hist.GetMean(1)))
        text = stats.AddText('Mean y = {0}'.format(hist.GetMean(2)))
        text = stats.AddText('RMS x  = {0}'.format(hist.GetRMS(1)))
        text = stats.AddText('RMS y  = {0}'.format(hist.GetRMS(2)))
        if ('S' in option) :
            text = stats.AddText('Skewness x = {0}'.format(hist.GetSkewness(1)))
            text = stats.AddText('Skewness y = {0}'.format(hist.GetSkewness(2)))
        if ('K' in option) :
            text = stats.AddText('Kurtosis x = {0}'.format(hist.GetKurtosis(1)))
            text = stats.AddText('Kurtosis y = {0}'.format(hist.GetKurtosis(2)))
        if ('U' in option or 'O' in option) :
            nbinsX = hist.GetNbinsX()
            nbinsY = hist.GetNbinsY()
            bl = hist.GetBinContent(0,0)
            bc = 0
            for i in range(1,nbinsX+1) :
                bc += hist.GetBinContent(i,0)
            br = hist.GetBinContent(nbinsX+1,0)
            cl = 0
            for i in range(1,nbinsY+1) :
                cl += hist.GetBinContent(0,i)
            cc = hist.GetEffectiveEntries()
            cr = 0
            for i in range(1,nbinsX+1) :
                cr += hist.GetBinContent(nbinsX+1,i)
            tl = hist.GetBinContent(0,nbinsY+1)
            tc = 0
            for i in range(1,nbinsX+1) :
                tc += hist.GetBinContent(i,nbinsY+1)
            tr = hist.GetBinContent(nbinsX+1,nbinsY+1)
            text = stats.AddText('{0}| {1}| {2}'.format(tl,tc,tr))
            text = stats.AddText('{0}| {1}| {2}'.format(cl,cc,cr))
            text = stats.AddText('{0}| {1}| {2}'.format(bl,bc,br))
    return stats

def TestStyle() :
    """
    Test the functions defined above in a quick and automatic way.
    """
    lhcbStyle = LHCbStyle()

    canvasDrawObj = ROOT.TCanvas('canvasDrawObj','canvasDrawObj',400,300)
    canvasDrawObjCompare = ROOT.TCanvas('canvasDrawObjCompare','canvasDrawObjCompare',400,300)
    canvasDrawObjCompare2 = ROOT.TCanvas('canvasDrawObjCompare2','canvasDrawObjCompare2',400,300)

    hist1 = ROOT.TH1F('hist1','hist1',100,0,10)
    hist2 = ROOT.TH1F('hist2','hist2',100,0,10)
    hist3 = ROOT.TH1F('hist3','hist3',100,0,10)

    hist2D = ROOT.TH2F('hist2D','hist2D',100,0,10,10,0,2)

    hist1.Fill(2,10)
    hist2.Fill(3,2)
    hist3.Fill(5,1)

    hist2D.Fill(2,1)
    hist2D.Fill(4,0)
    hist2D.Fill(5,2,10)

    ax = np.arange(0.,10.,2.)
    ay = ax*5.
    asize = len(ax)
    print ax
    print ay
    print asize
    graph = ROOT.TGraph(asize,ax,ay)
    graph.SetName('graph')

    # Test InitHist.
    InitHist(hist1,title='title for hist 1',x='x axis',y='y axis')
    InitHist(hist2,title='title for hist 2',x='x axis')
    InitHist(hist3,title='')

    # Test InitHist2.
    InitHist2(hist2D)

    # Test InitGraph.
    InitGraph(graph)

    markerColor=ROOT.kMagenta
    lineColor=ROOT.kMagenta
    fillColor=ROOT.kMagenta
    # Test SetStyleObj.
    SetStyleObj(hist1,markerColor,lineColor,fillColor)
    SetStyleObj(graph,markerColor,lineColor,fillColor)

    # Test AdjustOffsets.
    AdjustOffsets(hist2D,'Z')

    objList=[hist1,hist2,hist3]
    colorDict={hist1.GetName():ROOT.kRed,hist2.GetName():ROOT.kBlue,hist3.GetName():ROOT.kGreen}
    optionDrawDict={hist1.GetName():'E',hist2.GetName():'',hist3.GetName():'E'}
    labelDict={hist1.GetName():'first',hist2.GetName():'second',hist3.GetName():'third'}
    optionLegDict={hist1.GetName():'l',hist2.GetName():'p',hist3.GetName():'fw'}

    objList2=[hist1,graph]
    colorDict2={hist1.GetName():ROOT.kRed,graph.GetName():ROOT.kBlue}
    optionDrawDict2={hist1.GetName():'E',graph.GetName():'APC'}
    labelDict2={hist1.GetName():'first',graph.GetName():'second'}
    optionLegDict2={hist1.GetName():'l',graph.GetName():'pw'}

    print objList2
    print colorDict2
    print optionDrawDict2
    print labelDict2
    print optionLegDict2

    # Test DrawObj.
    DrawObj(canvasDrawObj,hist1)

    # Test CreateLegend.
    leg = CreateLegend(objList,labelDict,optionLegDict,0.5,0.5,0.8,0.7)
    leg2 = CreateLegend(objList2,labelDict2,optionLegDict2,0.6,0.2,0.9,0.5)

    # Test DrawObjCompare.
    DrawObjCompare(canvasDrawObjCompare,objList,colorDict,optionDrawDict,leg)
    DrawObjCompare(canvasDrawObjCompare2,objList2,colorDict2,optionDrawDict2,leg2)

    option = 'IUO'
    # Test CreateStats.
    CreateStats(hist1,option)

    return

LHCbStyle()