SNDLHCplotter.py

import os
import sys
from argparse import ArgumentParser
from datetime import date

import ROOT
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from SNDstyle import init_style, writeSND

today = date.today().strftime('%d%m%y')
"""
    SNDLHCplotter.py    A multi-purpose histogram plotter for SND@LHC (D. Centanni 2023)

    Command-Line Usage
    ------------------
    To run the script, open a terminal and execute the following command:
    python -i SNDLHCplotter.py [options] [arguments]

    | Option            | Description                                                                                           |
    | ----------------- | ----------------------------------------------------------------------------------------------------- |
    | --help            | Show the help message and exit.                                                                       |
    | -f, --inputFile   | Specify the input file path.                                                                          |
    | -c, --inputCanvas | Specify the input canvas file path.                                                                   |
    | -e, --extratext   | Add extratext below the SND@LHC writing.                                                              |
    | -hname            | Specify the name(s) of the histogram(s) to be plotted.                                                |
    | --scale           | Specify the scale factor to be applied to each histogram.                                             |
    | --auto            | Enables the auto mode: automatically plots histograms according to the arguments provided.            |
    | --dataMC          | Enables the Data-MonteCarlo comparison mode: data histogram must contain DATA in its name! (for now). |

    Examples
    --------
    1. Auto-mode, single histogram:
        python -i SNDLHCplotter.py -f histofile.root -hname Nscifi_hits -e Preliminary --auto
    2. Auto-mode, multi histograms:
        python -i SNDLHCplotter.py -f histofile.root -hname Nscifi_hits1 Nscifi_hits3 Nscifi_hits3 -e Preliminary --auto
    3. Auto-mode, Data-Montecarlo comparison:
        python -i SNDLHCplotter.py -f histofile.root -hname DATA_Nscifi_hits MC_Nscifi_hits --auto --dataMC
    4. ...

    Still WIP
"""


def load_hists(histfile, query=None):
    f = ROOT.TFile.Open(histfile)
    histlist = {}
    keylist = f.GetListOfKeys()
    for key in keylist:
        if query is not None and key.GetName() not in query:
            continue
        hist = f.Get(key.GetName())
        # check if histogram is readable
        try:
            hist.SetDirectory(ROOT.gROOT)
        except:
            print('### WARNING ###: key "'+str(key.GetName())+'" does not correspond to valid hist.')
            continue
        hist.SetName(key.GetName())
        histlist[key.GetName()] = hist
    if len(histlist) == 0: raise Exception('ERROR: histlist is empty!')
    f.Close()
    return histlist

def getHistFromfiles(filelist, hname, labellist):
    histlist = {}
    if len(file_list) != len(labellist): raise Exception('N. of files and labels mismatches!')
    for i_file, f in enumerate(file_list):
        fin = ROOT.TFile.Open(f)
        hist = fin.Get(hname)
        try:
            hist.SetDirectory(ROOT.gROOT)
        except:
            print('### WARNING ###: Name "'+str(hname)+'" does not correspond to valid hist.')
            continue
        hist.SetName(labellist[i_file]+'_'+hname)
        histlist[labellist[i_file]+'_'+hname] = hist
        fin.Close()
    if len(histlist) == 0: raise Exception('ERROR: histlist is empty!')
    return histlist
    

def createCanvas(noPlots):
    if noPlots == 1:
        xPad = 1; yPad = 1; width = 550; height = 0.90*width
    elif noPlots == 2:
        xPad = 2; yPad = 1; width = 600; height = 0.50*width
    elif noPlots == 3:
        xPad = 3; yPad = 1; width = 900; height = 0.4*width
    elif noPlots == 4:
        xPad = 2; yPad = 2; width = 600; height = width
    else:
        xPad = 3; yPad = 2; width = 800; height = 0.55*width
    noPadPerCanv = xPad * yPad
    nCanvases = int(noPlots/6)+1
    canvs = {}
    for i in range(nCanvases):
        #canvs['canv'+str(i)] = ROOT.TCanvas("canv"+str(i), "Variables", int(width), int(height))
        canvs['canv'+str(i)] = ROOT.TCanvas("canv"+str(i), "Variables", 1000, int(0.55*1000/0.55))
        canvs['canv'+str(i)].Divide(xPad, yPad)
    #canv = ROOT.TCanvas("canv", "Variables", int(width), int(height))
    #canv = ROOT.TCanvas("canv", "Variables", 1200, int(0.55*1200))
    return canvs

def drawSingleHisto(hist, canvas=None, xaxtitle=None, yaxtitle=None, 
            label='auto', color=None, logy=False, drawoptions='',
	        extratext=None,topmargin=None, bottommargin=None,
	        leftmargin=None, rightmargin=None,
	        xaxlabelfont=None, xaxlabelsize=None, outpath='', rebin=None, sigma=list(), scale=1., xaxrange=None, yaxrange=None):

    if not canvas:
        canvas = ROOT.TCanvas("c", "c", 800, 600)

    if color is None: color = ROOT.kAzure-4

    if xaxlabelfont is None: xaxlabelfont = 4 
    if xaxlabelsize is None: xaxlabelsize = 15
    yaxlabelfont = 4; yaxlabelsize = 15
    axtitlefont = 6; axtitlesize = 18
    legendfont = 5
    if leftmargin is None: leftmargin = 0.15
    if rightmargin is None: rightmargin = 0.05
    if topmargin is None: topmargin = 0.05
    if bottommargin is None: bottommargin = 0.15
    canvas.SetBottomMargin(bottommargin)
    canvas.SetLeftMargin(leftmargin)
    canvas.SetRightMargin(rightmargin)
    canvas.SetTopMargin(topmargin)

    pentryheight = 0.12
    plegendbox = ([leftmargin+0.30,1-topmargin-pentryheight-0.01, 1-rightmargin-0.03,1-topmargin-0.03])
    hist.SetLineColor(color)
    hist.SetLineWidth(2)

    if rebin is not None:
        hist.Rebin(int(hist.GetNbinsX()/rebin))

    if scale !=1.:
        hist.Scale(scale)
    
    if label is not None:
        leg = ROOT.TLegend(plegendbox[0],plegendbox[1],plegendbox[2],plegendbox[3])
        leg.SetTextFont(10*legendfont+2)
        leg.SetFillColor(ROOT.kWhite)
        leg.SetBorderSize(0)
        if label=='auto': 
            label = hist.GetTitle()
            if label=='':
                label = hist.GetName()
        leg.AddEntry(hist,label,"l")
        leg.AddEntry(0, 'Entries: {:.2e}'.format(hist.GetEntries()), '')
        leg.AddEntry(0, 'Integral: {:.2e}'.format(hist.Integral()), '')
        leg.AddEntry(0, 'Mean: {:.2e}'.format(hist.GetMean()), '')
        leg.AddEntry(0, 'Std dev: {:.2e}'.format(hist.GetStdDev()), '')
        #leg.AddEntry(0,'#sigma: {:.2e}'.format(sigma[0]), '')
        #leg.AddEntry(0,'#sigma error: {:.2e}'.format(sigma[1]), '')

    xax = hist.GetXaxis()
    #xax.SetNdivisions(5,4,0,ROOT.kTRUE)
    xax.SetLabelFont(10*xaxlabelfont+3)
    xax.SetLabelSize(xaxlabelsize)
    if xaxtitle is not None:
        xax.SetTitle(xaxtitle)
    xax.SetTitleFont(10*axtitlefont+3)
    xax.SetTitleSize(axtitlesize)
    xax.SetTitleOffset(1.2)
    xax.CenterTitle(True)
    if xaxrange:
        xax.SetRangeUser(float(xaxrange[0]), float(xaxrange[1]))
    if not logy:
        hist.SetMaximum(hist.GetMaximum()*1.35)
        hist.SetMinimum(0.)
    else:
        hist.SetMaximum(hist.GetMaximum()*100)
        #hist.SetMinimum(hist.GetMaximum()/1e7)
        canvas.SetLogy()
    yax = hist.GetYaxis()
    yax.SetMaxDigits(3)
    yax.SetNdivisions(8,4,0,ROOT.kTRUE)
    yax.SetLabelFont(10*yaxlabelfont+3)
    yax.SetLabelSize(yaxlabelsize)
    if yaxtitle is not None: 
        yax.SetTitle(yaxtitle)
    yax.SetTitleFont(10*axtitlefont+3)
    yax.SetTitleSize(axtitlesize)
    yax.SetTitleOffset(1.5)
    yax.CenterTitle(True)
    if yaxrange:yax.SetRangeUser(float(yaxrange[0]), float(yaxrange[1]))
    
    hist.Draw(drawoptions)
    ROOT.gPad.RedrawAxis()
    writeSND(canvas, extratext=extratext)
    if label is not None: leg.DrawClone("same")
    ROOT.gPad.Update()
    canvas.Draw()
    canvas.SaveAs(outpath+hist.GetName()+'.pdf', 'pdf')


def drawDATAMC(histlist, c1=None, figname='DATA-MC', xaxtitle=None, yaxtitle=None,
	    normalize=False, lumi=None, dolegend=True, labellist=None, logy=False, extra_text = '', rebin=None, outpath='.', xaxrange=[]):

    if not c1:
        c1 = ROOT.TCanvas()
    
    xaxlabelfont = 4 
    xaxlabelsize = 22
    yaxlabelfont = 4; yaxlabelsize = 22
    axtitlefont = 6; axtitlesize = 26
    legendfont = 5
    leftmargin = 0.15
    rightmargin = 0.05
    topmargin = 0.05
    bottommargin = 0.15
    c1.SetBottomMargin(bottommargin)
    c1.SetLeftMargin(leftmargin)
    c1.SetRightMargin(rightmargin)
    c1.SetTopMargin(topmargin)


    pentryheight = 0.06
    nentries = 1 + len(histlist)
    if nentries>3: pentryheight = pentryheight*0.8
    plegendbox = ([leftmargin+0.45,1-topmargin-pentryheight*nentries, 1-rightmargin-0.03,1-topmargin-0.03])
    
    if rebin is not None:
        for hist in histlist:
            hist.Rebin(int(hist.GetNbinsX()/rebin))

    pairs = list()
    for hist in histlist:
        hist.SetStats(0)
        pairs.append([hist.GetMaximum(), hist])
    maxpair = max(pairs,key=lambda item:item[0])


    if normalize:
        for hist in histlist:
            if hist.Integral() == 0: 
                print(hist.GetName(), 'has null integral, skipping')
                return            
            hist.Scale(1./hist.Integral())
    elif lumi:
        for hist in histlist:
            if hist.Integral() == 0: 
                print(hist.GetName(), 'has null integral, skipping')
                return  
            if not 'DATA' in hist.GetName():
                hist.Scale(lumi)
            else:
                hist.Scale(1.)
    
    if not logy:
        maxpair[1].SetMaximum(maxpair[1].GetMaximum()*1.2)
        maxpair[1].SetMinimum(0.)
        for h in histlist:
            h.SetMaximum(maxpair[1].GetMaximum()*1.2)
    else:
        maxpair[1].SetMaximum(maxpair[1].GetMaximum()*10)
        for h in histlist:
            h.SetMaximum(maxpair[1].GetMaximum()*10)
        c1.SetLogy()
        
    data_index = -1
    for i, h in enumerate(histlist):
        if 'DATA' in h.GetName():
            h.SetMarkerColor(ROOT.kBlack)
            h.SetLineColor(ROOT.kBlack)
            h.SetMarkerStyle(8)
            h.SetMarkerSize(0.5)
            data_index = i
        else:
            h.SetLineWidth(1)
            h.SetFillColor(ROOT.TColor.GetColor("#7d99d1"))
            h.SetLineColor(ROOT.TColor.GetColor('#0000ee'))
            h.SetFillStyle(1001)

    legend = ROOT.TLegend(plegendbox[0],plegendbox[1],plegendbox[2],plegendbox[3])
    legend.SetNColumns(1)
    legend.SetName('legend')
    legend.SetFillColor(ROOT.kWhite)
    legend.SetTextFont(10*legendfont+3)
    legend.SetBorderSize(0)
    for i,hist in enumerate(histlist):
        label = hist.GetTitle()
        if labellist is not None: label = labellist[i]
        else:
            if 'DATA' in hist.GetName():
                if lumi:
                    legend.AddEntry(hist, 'Data L_{int} = '+str(lumi)+' fb^{-1}', "PEL")
                else:
                    legend.AddEntry(hist, 'Data', "PEL")
            else:
                legend.AddEntry(hist, 'Background', "FEL")
        legend.AddEntry(hist, 'Integral: {:.2e}'.format(hist.Integral()), '')
        legend.AddEntry(hist, 'Mean: {:.2e}'.format(hist.GetMean()), '')
    
    for h in histlist:
        xax = h.GetXaxis()
        xax.SetLabelSize(xaxlabelsize)
        xax.SetLabelFont(10*xaxlabelfont+3)
        if xaxtitle is not None:
            xax.SetTitle(xaxtitle)
        xax.SetTitleFont(10*axtitlefont+3)
        xax.SetTitleSize(axtitlesize)
        xax.SetTitleOffset(1.2)
        if xaxrange:
            xax.SetRangeUser(xaxrange[0], xaxrange[1])
        # Y-axis layout
        yax = h.GetYaxis()
        yax.SetMaxDigits(3)
        yax.SetLabelFont(10*yaxlabelfont+3)
        yax.SetLabelSize(yaxlabelsize)
        if yaxtitle is not None:
            yax.SetTitle(yaxtitle)
        yax.SetTitleFont(10*axtitlefont+3)
        yax.SetTitleSize(axtitlesize)
        yax.SetTitleOffset(1.2)
        yax.CenterTitle(True)
        hist.SetMaximum(maxpair[1].GetMaximum()*1.2)

    for i in range(len(histlist)):
        drawopt = ''
        if i!= data_index:
            if histlist[i].GetEntries() < 20: drawopt = 'E'
            histlist[i].Draw("HIST SAME "+drawopt)
    histlist[data_index].Draw("* SAME")
    
    
    ROOT.gPad.RedrawAxis()
    writeSND(c1, extratext=extra_text)
    if dolegend: legend.DrawClone("same")
    ROOT.gPad.Update()
    c1.Draw()
    c1.SaveAs(outpath+figname+'.pdf', 'pdf')

def draw2dHisto(hist, canvas=None,xaxtitle=None, yaxtitle=None, 
            label=None, drawoptions='COLZ',
	        extratext=None,topmargin=None, bottommargin=None,
	        leftmargin=None, rightmargin=None,
	        xaxlabelfont=None, xaxlabelsize=None, outpath=''):

    if not canvas:
        canvas = ROOT.TCanvas()

    if xaxlabelfont is None: xaxlabelfont = 4 
    if xaxlabelsize is None: xaxlabelsize = 22
    yaxlabelfont = 4; yaxlabelsize = 22
    axtitlefont = 6; axtitlesize = 26
    legendfont = 5
    if leftmargin is None: leftmargin = 0.15
    if rightmargin is None: rightmargin = 0.05
    if topmargin is None: topmargin = 0.05
    if bottommargin is None: bottommargin = 0.15
    canvas.SetBottomMargin(bottommargin)
    canvas.SetLeftMargin(leftmargin)
    canvas.SetRightMargin(rightmargin)
    canvas.SetTopMargin(topmargin)

    pentryheight = 0.15
    plegendbox = ([leftmargin+0.45,1-topmargin-pentryheight, 1-rightmargin-0.03,1-topmargin-0.03])
    
    if label is not None:
        hist.SetTitle(label)
    
    hist.SetStats(0)

    xax = hist.GetXaxis()
    #xax.SetNdivisions(5,4,0,ROOT.kTRUE)
    xax.SetLabelFont(10*xaxlabelfont+3)
    xax.SetLabelSize(xaxlabelsize)
    if xaxtitle is not None:
        xax.SetTitle(xaxtitle)
    xax.SetTitleFont(10*axtitlefont+3)
    xax.SetTitleSize(axtitlesize)
    xax.SetTitleOffset(1.2)

    yax = hist.GetYaxis()
    yax.SetMaxDigits(3)
    yax.SetNdivisions(8,4,0,ROOT.kTRUE)
    yax.SetLabelFont(10*yaxlabelfont+3)
    yax.SetLabelSize(yaxlabelsize)

    if yaxtitle is not None: 
        yax.SetTitle(yaxtitle)
    yax.SetTitleFont(10*axtitlefont+3)
    yax.SetTitleSize(axtitlesize)
    yax.SetTitleOffset(1.2)
    
    hist.Draw(drawoptions)
    ROOT.gPad.RedrawAxis()
    writeSND(canvas, extratext=extratext, text_in=False)
    ROOT.gPad.Update()
    canvas.Draw()
    canvas.SaveAs(outpath+hist.GetName()+'.pdf', 'pdf')

def drawMultiHisto(histlist, c1=None, figname='multihisto', xaxtitle=None, yaxtitle=None,
	    normalize=False, dolegend=True, labellist=None, 
	    colorlist=None, logy=False, drawoptions='', extra_text = '', rebin=None, outpath='.', scale=1., xaxrange=None):

    if not c1: c1 = ROOT.TCanvas("c", "c", 800, 600)

    if colorlist is None:
        colorlist = ([ROOT.kAzure-4, ROOT.kRed, ROOT.kGreen+1, ROOT.kViolet, ROOT.kMagenta-9,
                        ROOT.kTeal-6, ROOT.kAzure+6, ROOT.kOrange, ROOT.kGreen-1])
    if( len(histlist)>len(colorlist) ):
        raise Exception('ERROR: histogram list is longer than color list')
    if(labellist is not None and len(labellist)!=len(histlist)):
        raise Exception('ERROR: length of label list does not agree with histogram list')
    
    xaxlabelfont = 4 
    xaxlabelsize = 22
    yaxlabelfont = 4; yaxlabelsize = 22
    axtitlefont = 6; axtitlesize = 26
    legendfont = 5
    leftmargin = 0.15
    rightmargin = 0.05
    topmargin = 0.05
    bottommargin = 0.15
    c1.SetBottomMargin(bottommargin)
    c1.SetLeftMargin(leftmargin)
    c1.SetRightMargin(rightmargin)
    c1.SetTopMargin(topmargin)

    pentryheight = 0.08
    nentries = 1 + len(histlist)
    if nentries>3: pentryheight = pentryheight*0.8
    plegendbox = ([leftmargin+0.30,1-topmargin-pentryheight*nentries, 1-rightmargin-0.03,1-topmargin-0.03])
    
    if rebin is not None:
        for hist in histlist:
            hist.Rebin(int(hist.GetNbinsX()/rebin))
    
    if scale !=1.:
        for hist in histlist: hist.Scale(scale)

    if normalize:
        for hist in histlist:
            hist.Scale(1./hist.Integral())
    
    pairs = list()
    for hist in histlist:
        hist.SetStats(0)
        pairs.append([hist.GetMaximum(), hist])
    print(pairs)
    maxpair = max(pairs,key=lambda item:item[0])
    print('Max is', maxpair)
    
    if not logy:
        maxpair[1].SetMaximum(maxpair[1].GetMaximum()*1.2)
        maxpair[1].SetMinimum(0.)
        for h in histlist:
            h.SetMaximum(maxpair[1].GetMaximum()*1.2)
    else:
        maxpair[1].SetMaximum(maxpair[1].GetMaximum()*10)
        for h in histlist:
            h.SetMaximum(maxpair[1].GetMaximum()*10)
        c1.SetLogy()
    
    """if len(histlist) == 3:
        histlist[2].SetMarkerColor(ROOT.kBlack)
        histlist[2].SetLineColor(ROOT.kBlack)
        histlist[2].SetMarkerStyle(8)
        histlist[2].SetMarkerSize(0.5)
        histlist[1].SetLineWidth(1)
        histlist[1].SetFillColor(ROOT.TColor.GetColor("#7d99d1"))
        histlist[1].SetLineColor(ROOT.TColor.GetColor('#0000ee'))
        histlist[1].SetFillStyle(1001)
        histlist[0].SetLineWidth(1)
        histlist[0].SetFillColor(ROOT.TColor.GetColor("#ff0000"))
        histlist[0].SetLineColor(ROOT.TColor.GetColor('#ff0000'))
        histlist[0].SetFillStyle(3554)
    else:"""
    for i,hist in enumerate(histlist):
        hist.SetLineWidth(2)
        hist.SetLineColor(colorlist[i])
        hist.SetMarkerSize(0)

    legend = ROOT.TLegend(plegendbox[0],plegendbox[1],plegendbox[2],plegendbox[3])
    legend.SetNColumns(1)
    legend.SetName('legend')
    legend.SetFillColor(ROOT.kWhite)
    legend.SetTextFont(10*legendfont+3)
    legend.SetBorderSize(0)
    for i,hist in enumerate(histlist):
        label = hist.GetName()
        if labellist is not None: label = labellist[i]
        legend.AddEntry(hist,label,"FL")
        #legend.AddEntry(hist, 'Integral: {:.2e}'.format(hist.Integral()), '')
        legend.AddEntry(hist, 'Mean: {:.2e}'.format(hist.GetMean()), '')
        legend.AddEntry(hist, 'Std dev: {:.2e}'.format(hist.GetStdDev()), '')
    

    for h in histlist: 
        xax = h.GetXaxis()
        xax.SetLabelSize(xaxlabelsize)
        xax.SetLabelFont(10*xaxlabelfont+3)
        if xaxtitle is not None:
            xax.SetTitle(xaxtitle)
        elif 'Momentum' in h.GetTitle():
            xax.SetTitle('Momentum [GeV/c]')
        elif 'Energy' in h.GetTitle():
            xax.SetTitle('Energy [GeV]')
        xax.SetTitleFont(10*axtitlefont+3)
        xax.SetTitleSize(axtitlesize)
        xax.SetTitleOffset(1.2)
        xax.CenterTitle(True)
        if xaxrange:
            xax.SetRangeUser(xaxrange[0], xaxrange[1])
        # Y-axis layout
        yax = h.GetYaxis()
        yax.SetMaxDigits(3)
        yax.SetLabelFont(10*yaxlabelfont+3)
        yax.SetLabelSize(yaxlabelsize)
        if yaxtitle:
            yax.SetTitle(yaxtitle)
        elif 'Momentum' in h.GetTitle() and not yaxtitle:
            yax.SetTitle('dN/dpdt [(GeV/c)^{-1}s^{-1}]')
        elif 'Energy' in h.GetTitle() and not yaxtitle:
            yax.SetTitle('dN/dEdt [GeV^{-1}s^{-1}]')
        yax.SetTitleFont(10*axtitlefont+3)
        yax.SetTitleSize(axtitlesize)
        yax.SetTitleOffset(1.2)
        yax.CenterTitle(True)

    """if len(histlist) == 3:
        for hist in histlist[:2]:
            hist.Draw('same '+drawoptions)
        histlist[2].Draw("* SAME")
    else:"""
    histlist[0].Draw(drawoptions)
    for hist in histlist[1:]:
        hist.Draw('same '+drawoptions)
    
    ROOT.gPad.RedrawAxis()
    writeSND(c1, extratext=extra_text)
    if dolegend: legend.DrawClone("same")
    ROOT.gPad.Update()
    c1.Draw()
    c1.SaveAs(outpath+figname+'.pdf', 'pdf')

def MultiCanvas(query=None):
    if len(options.inputFile) == 1: f = options.inputFile[0]
    histlist = load_hists(f)
    if query == None:
        Varlist = [histlist.keys()]
    else:
        Varlist = query
    noPlots = len(Varlist)
    global canvs
    canvs = createCanvas(noPlots)
    for icanv, canv in enumerate(canvs.values()):
        ipad = 0
        min = 6*icanv
        max = 6*(icanv+1)
        if max > len(Varlist): max= len(Varlist)
        for var in Varlist[min:max]:
            sel_histo = histlist[var] 
            #print('Var is', var, 'Selected histos', hlist, 'name', hlist[0].GetName(), hlist[1].GetName(), hlist[2].GetName())
            pad = canv.cd(ipad+1)
            logy = False
            norm = True
            axtitle = 'a.u.'
            if options.lumi:
                norm = False
                axtitle='N'
                logy=True
                if options.norm: 
                    norm = True
                    axtitle= 'a.u.'
            #drawDATAMC(hlist, c1=pad, xaxtitle=var, yaxtitle=axtitle, normalize=norm, extra_text='Comparison', logy=logy, lumi=options.lumi) There can be an option to plot TMVA-like canvas
            drawSingleHisto(sel_hist, c1=pad, xaxtitle=sel_hist.GetXaxis().GetTitle(), yaxtitle=axtitle, extratext=options.extratext, logy=logy, drawoptions='HIST', outpath=outpath, scale=options.scalefactor, label=sel_hist.GetTitle())
            ipad+=1
        canv.SaveAs("canvas_"+str(icanv)+".pdf", "pdf") 

parser = ArgumentParser()
parser.add_argument("-f", nargs='+', dest="inputFile", help="input files", required=False)
parser.add_argument("-labels", nargs='+', dest="labels", help="list of labels", required=False, default=None)
parser.add_argument("-c", "--inputCanvas", dest="inputCanvas", help="single input canvas", required=False)
parser.add_argument("-e", "--extratext", dest="extratext", help="extratext written below SND@LHC", default=None, required=False)
parser.add_argument('-hname', nargs='+', dest="hname", help='List of histos to be drawn', required=False)
parser.add_argument("--scale", dest="scalefactor", help="scale factor", required=False, type=float, default=1.)
parser.add_argument("--lumi", dest="lumi", help="luminosity factor", required=False, type=float, default=None)
parser.add_argument("--auto", dest="auto", action='store_true', help='Enables automatic mode',required=False, default=False)
parser.add_argument("--sep", dest="sep", action='store_true', help='Enables multiple histos in same canvas',required=False, default=False)
parser.add_argument("--norm", dest="norm", action='store_true', help='Normalizes multi-hist plotting',required=False, default=False)
parser.add_argument("--dataMC", dest="dataMC", help='Enables dataMC comparison mode: data histogram must contain DATA in its name', action='store_true', required=False, default=False)
parser.add_argument("-xrange", nargs='+', dest="xaxrange", help="X axis range", required=False, default=None)
parser.add_argument("-yrange", nargs='+', dest="yaxrange", help="Y axis range", required=False, default=None)

options = parser.parse_args()

if options.inputFile and len(options.inputFile) > 1 and len(options.hname)>1: raise Exception('Multi-file & Multi-histos not yet implemented!')
if options.inputFile and len(options.inputFile) > 1 and len(options.hname)==1 and options.labels == None: raise Exception('Please provide labellist for different input files!')
singlefile= False
if options.inputFile:
    if len(options.inputFile) < 2:
        options.inputFile = options.inputFile[0]
        tmp = options.inputFile.split('.')
        singlefile = True
        if options.hname and len(options.hname)> 1:
            Hlist = load_hists(options.inputFile, query=options.hname)
        else:
            Hlist = load_hists(options.inputFile)
    else:
        tmp = [str(today)]
        if options.hname and len(options.hname) < 2:
            file_list = options.inputFile
            Hlist = getHistFromfiles(file_list, options.hname[0], options.labels)
            print(Hlist)
    outpath = 'plots_'+tmp[0]+'/'
    if not os.path.exists(outpath):
            os.makedirs(outpath)
    
        

canvases    = {}
extratext=''
if options.extratext:
    extratext=options.extratext

xaxrange = None
yaxrange = None
if options.xaxrange:
    xaxrange = list(options.xaxrange)
if options.yaxrange:
    yaxrange = list(options.yaxrange)

if options.inputFile:
    init_style()
    if not options.hname and options.auto:
        for i_h,h in enumerate(Hlist.values()):
            if i_h not in canvases.keys():
                canvases[i_h] = ROOT.TCanvas("c"+str(i_h), "c"+str(i_h), 800, 600)
            htype = h.IsA().GetName()
            if 'TH1' in htype:
                drawSingleHisto(h, canvases[i_h], drawoptions='HIST', extratext=extratext, logy=True, outpath=outpath)
            elif 'TH2' in htype:
                draw2dHisto(h, canvases[i_h], extratext=extratext, outpath=outpath)
    elif options.auto and len(options.hname) < 2:
        if singlefile:
            i_h = 0
            options.hname = options.hname[0]
            canvases[i_h] = ROOT.TCanvas("c"+str(i_h), "c"+str(i_h), 800, 800)
            htype = Hlist[options.hname].IsA().GetName()
            sel_hist = Hlist[options.hname]
            if 'TH1' in htype:
                drawSingleHisto(Hlist[options.hname], canvases[i_h], drawoptions='HIST', extratext=extratext, logy=False, outpath=outpath, scale=options.scalefactor, label=sel_hist.GetTitle(), xaxtitle=sel_hist.GetXaxis().GetTitle(), yaxtitle=sel_hist.GetYaxis().GetTitle(), xaxrange=xaxrange, yaxrange=yaxrange) #,xaxrange[])
            elif 'TH2' in htype:
                draw2dHisto(Hlist[options.hname], canvases[i_h], extratext=extratext, outpath=outpath)
        else:
            i_h = 0
            options.hname = options.hname[0]
            canvases[i_h] = ROOT.TCanvas("c"+str(i_h), "c"+str(i_h), 800, 800)
            for hist in Hlist.values():
                if not 'TH1' in hist.IsA().GetName(): 
                    print('Not supported!')
                    continue
            drawMultiHisto(list(Hlist.values()), canvases[i_h], drawoptions='HIST', extra_text=extratext, outpath=outpath, scale=options.scalefactor, yaxtitle=list(Hlist.values())[0].GetYaxis().GetTitle(), xaxtitle=list(Hlist.values())[0].GetXaxis().GetTitle(), normalize=options.norm)
    elif options.auto and len(options.hname)>1:
        if options.sep:
            _hlist = list(Hlist.values())
            canvases = createCanvas(len(_hlist))
            for icanv, c in enumerate(canvases.values()):
                ipad = 0
                min = 6*icanv
                max = 6*(icanv+1)
                if max > len(_hlist): max= len(_hlist)
                for _h in _hlist[min:max]:
                    pad = c.cd(ipad+1)
                    htype = _h.IsA().GetName()
                    if 'TH1' in htype:
                        logy=True
                        drawSingleHisto(_h, pad, drawoptions='HIST', extratext=extratext, logy=logy, outpath=outpath, scale=options.scalefactor, xaxtitle=_h.GetXaxis().GetTitle(), yaxtitle=_h.GetYaxis().GetTitle())
                        ipad+=1
                    elif 'TH2' in htype:
                        draw2dHisto(_h, pad, extratext=extratext, outpath=outpath)
                        ipad+=1
                    else:
                        raise Exception('Object type not recognized.')    
        else:
            i_h = 0
            canvases[i_h] = ROOT.TCanvas("c"+str(i_h), "c"+str(i_h), 800, 800)
            if options.dataMC:
                drawDATAMC(list(Hlist.values()), canvases[i_h], xaxtitle=list(Hlist.values())[0].GetXaxis().GetTitle(), yaxtitle=list(Hlist.values())[0].GetYaxis().GetTitle(), normalize=True, extra_text='DATA-MC Comparison')
            else:
                drawMultiHisto(list(Hlist.values()), canvases[i_h], logy=False, drawoptions='HIST', extra_text=extratext, outpath=outpath, scale=options.scalefactor, yaxtitle='Counts', xaxtitle='diff (#mu_{out}-#mu_{in})', normalize=options.norm, labellist=options.labels)


#################################################################################

elif options.inputCanvas:
    f = ROOT.TFile.Open(options.inputCanvas)
    keylist = f.GetListOfKeys()
    clist = list()
    for key in keylist:
        canvas = f.Get(key.GetName())
        canvas.SetName(key.GetName())
        clist.append(canvas)
    #f.Close()
    histlist = list()
    legendlist = list()
    for c in clist:
        _list = c.GetListOfPrimitives()
        for l in _list:
            #print(l.IsA().GetName())
            if 'TH1' in l.IsA().GetName():
                histlist.append(l)
            elif 'TLegend' in l.IsA().GetName():
                legendlist.append(l)
    init_style()
    if len(histlist)==1:
        canvases[0] = ROOT.TCanvas("c"+str(0), "c"+str(0), 800, 600)
        h = histlist[0]
        drawSingleHisto(h, canvases[0], drawoptions=h.GetDrawOption(), extratext=extratext, logy=True, outpath='')
    elif len(histlist)== 2:
        #for p in legendlist[0].GetListOfPrimitives():
        #    print(p.GetLabel())
        canvases[0] = ROOT.TCanvas("c"+str(0), "c"+str(0), 800, 600)
        #drawDATAMC(histlist=histlist, c1=canvases[0], xaxtitle='US QDC', yaxtitle='a.u.', 
        #           dolegend=True, labellist=['Data: 39 fb^{-1}', 'MonteCarlo shifted'], figname='USQDCcomparison', rebin=110)
        drawDATAMC(histlist=histlist, c1=canvases[0], xaxtitle='N SciFi hits', yaxtitle='a.u.', 
                   dolegend=True, labellist=['Data: 39 fb^{-1}', 'MonteCarlo'], figname='Nsfhits', logy=True, extra_text=extratext)
            


# edit the last part in order to take and sort all of the items present in the .root file (frame, th1, tlegend, tpavetext...)