#!/usr/bin/env python
"""

Script to merge the histograms obtained by the different groups with Ex1*

"""
import os,sys
import datetime

from math import pow,exp,pi,sqrt
from ROOT import TCanvas, TFile, gROOT, gStyle,TPaveText,TGraph,TF1,TLine,TBox
from ROOT import *
#-------------------------------------------------------------
# CONTROL CARDS
#-------------------------------------------------------------
DMASS = 1864.86 #MeV
SigMassLow  = DMASS - 20. #1846.
SigMassHigh = DMASS + 20. #1886.


def waittocontinue(summary=None,next=None):
    if summary is not None: print summary
    if next is not None: print "Next :",next
    _wait = True
    _stop = False
    while _wait :
        userinput = raw_input('continue (y/[n]/x) ? ')
        if len(userinput)==0: continue
        if userinput[0]=='x':
            _wait = False
            _stop = True
        if userinput[0]=='y':
            _wait = False 
    return _stop        

def popup(_title,_text):
    try:
        import tkMessageBox
        import Tkinter
        window=Tkinter.Tk()
        window.wm_withdraw()
        tkMessageBox.showinfo(title=_title,message=_text)
    except :
        print "Unable to pop up window"
        print "*"*80
        print "%s : %s"%(_title,_text)
        print "*"*80
    
from optparse   import Option, OptionParser
def getOpt():
    parser = OptionParser()
    baseOptions(parser)
    opt = parser.parse_args()[0]
    _range = []
    for a_range in opt.range.split("-"):
        _min,_max = a_range.split(",")
        #print 'min,max=',_min,_max,int(_min),int(_max)
        _range += range(int(_min),int(_max)+1)
    opt.range = _range    
    return opt
                    
def baseOptions(parser):
    '''
    Option arguments for mass histogram merger script
    '''

    # ---

    parser.add_option("--test",
                      help="Test",
                      dest="test",
                      default=False,
                      action="store_true")
    
    parser.add_option("--verbose",
                      help="Verbose",
                      dest="verbose",
                      default=False,
                      action="store_true")
    
    parser.add_option("-r","--range",
                      help="Range of group number (min,max - comma separated)",
                      dest="range",
                      default='0,30', # 1,5-101,103-106,108 => [1,2,3,4,5,101,102,103,106,107,108]
                      action="store")

    parser.add_option("--popup",
                      help="Allow popup windows",
                      dest="popup",
                      default=False,
                      action="store_true")

    parser.add_option("-d","--result-directory",
                      help="Result directory",
                      dest="result_directory",
                      default=os.path.join("/afs/cern.ch/user/j/jmarchan/LHCbMasterclass/DEPOT/",'LHCbMasterclass-%s'%(datetime.date.isoformat(datetime.date.today()))),
                      action="store")


def massFunction(x,par):
    return par[0]*exp(-1.*pow(x[0]-par[1],2)/(2.*pow(par[2],2)))+par[3]
def massFunctionSig(x,par):
    return par[0]*exp(-1.*pow(x[0]-par[1],2)/(2.*pow(par[2],2)))
def massFunctionBdf(x,par):
    return par[0]


if __name__ == "__main__":

    opt = getOpt()

    ##print "directory=",opt.result_directory
    
    gStyle.SetOptStat(0)
    gStyle.SetOptTitle(0)
    
    #mergeHisto(opt)

    trend_evts = []
    canv = TCanvas("Results","Results")
    canv.Divide(1,2)

    first = True
    hsum = None
    paves = []
    rootfiles_in_result_directory = [f for f in  os.listdir(opt.result_directory) if os.path.splitext(f)[1]=='.root']
    Nsamples = 0
    for combination in opt.range:
        root_filenames = [f for f in rootfiles_in_result_directory if f.rfind("MassHistogram_%d_"%(combination))>-1]
        ##print " root_filenames", root_filenames
        if len(root_filenames)==0:
            print 'xxx Distribution en masse pour la combinaison %d introuvable'%(combination)
            if opt.popup:
                popup('?????',"Qui n'a pas valide l'exercice 1 ? (%d)"%(combination)) 
            continue
        root_filename = os.path.join(opt.result_directory,root_filenames[0])
        ##print "root file",root_filename
        if not os.path.exists(root_filename):
            print 'XXX Gros foutage de gueule !!!'
            continue
        ##print "On a trouve le fichier pour la combinaison %d "%(combination)
        root_file = TFile(root_filename)
        
        h = root_file.Get('hist')
        xaxis = h.GetXaxis()
        wbin = (xaxis.GetXmax()-xaxis.GetXmin())/xaxis.GetNbins()
        h.SetStats(0)
        h.GetXaxis().SetTitle("Mass(K-#pi) [MeV]")
        h.GetYaxis().SetTitle("# of evts / %.2f MeV"%(wbin))
        h.GetXaxis().SetTitleSize(0.08)
        h.GetYaxis().SetTitleOffset(0.5)
        h.GetYaxis().SetTitleSize(0.065)
        print '--> Drawing combination %d'%(combination)
        p1 = canv.cd(1)
        p1.SetBottomMargin(0.23)
        h.Draw()

        n_comb = h.GetEntries()
        pave1a  =  TPaveText(.15,.85,.35,.95,"NDC")
        pave1a.SetName("TP1a%d"%(combination))
        pave1a.AddText("Combinaison %d"%(combination))
        pave1a.SetBorderSize(0)
        pave1a.SetFillColor(0)
        pave1a.SetTextAlign(12)
        canv.cd(1)
        pave1a.Draw()

        pave1b  =  TPaveText(.75,.85,.95,.95,"NDC")
        pave1b.SetName("TP1b%d"%(combination))
        pave1b.AddText("%d evts"%(n_comb))
        pave1b.SetBorderSize(1)
        pave1b.SetFillColor(0)
        pave1b.SetTextAlign(32)
        canv.cd(1)
        pave1b.Draw()

        if first :
            gROOT.cd()
            hsum = h.Clone("hsum")
            hsum.SetStats(0)
            hsum.GetXaxis().SetTitle("Mass(K-#pi) [MeV]")
            hsum.GetYaxis().SetTitle("# of evts / %.2f MeV"%(wbin))
            hsum.GetXaxis().SetTitleSize(0.08)
            hsum.GetYaxis().SetTitleOffset(0.5)
            hsum.GetYaxis().SetTitleSize(0.065)
            first=False
        else:
            hsum.Add(h)    

        Nsamples+=1     

        p2 = canv.cd(2)
        p2.SetBottomMargin(0.23)
        hsum.Draw()
        n_tot = hsum.GetEntries()
        pave2  =  TPaveText(.75,.85,.95,.95,"NDC")
        pave2.SetName("TP2%d"%(combination))
        pave2.AddText("%d evts"%(n_tot))
        pave2.SetBorderSize(1)
        pave2.SetLineColor(2)
        pave2.SetFillColor(0)
        pave2.SetTextColor(2)
        pave2.SetTextAlign(32)
        pave2.SetTextSize(0.1)
        canv.cd(2)
        pave2.Draw()

        paves += [pave1a,pave1b,pave2]

        canv.Update()

        trend_evts.append( (combination,n_comb)) 

        #=============================================================
        if waittocontinue(): break        
        #=============================================================

    if Nsamples==0:
        print "No histograms found, please check range and directory"
        print "opt.range: (-r) ",opt.range
        print "opt.result_directory: (-d)",opt.result_directory
        print "rootfiles_in_result_directory",rootfiles_in_result_directory
        sys.exit()
        
    print "-"*60    
    print "We have all results : %d events from %d samples"%(n_tot,Nsamples)
    print "Next : summary plots"
    print "-"*60    

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================

    canv.cd(1)
    gtrend_evts = TGraph()
    for ipoint,(icombination,nevts) in enumerate(trend_evts):
        #gtrend_evts.SetPoint(ipoint,icombination,nevts)
        gtrend_evts.SetPoint(ipoint,ipoint,nevts)
    gtrend_evts.SetMarkerStyle(20)
    gtrend_evts.SetMarkerSize(1.)
    #gtrend_evts.GetXaxis().SetTitle("Combination number")
    gtrend_evts.GetXaxis().SetTitle("Sample")
    gtrend_evts.GetYaxis().SetTitle("# of evts analyzed")
    gtrend_evts.GetXaxis().SetTitleSize(0.08)
    gtrend_evts.GetYaxis().SetTitleSize(0.08)
    gtrend_evts.GetYaxis().SetTitleOffset(0.5)
    gtrend_evts.GetYaxis().SetTitleSize(0.065)
    gtrend_evts.Draw("AP")
    
    pave1c  =  TPaveText(.55,.90,.90,.99,"NDC")
    pave1c.SetName("TP1c")
    pave1c.AddText("%d samples ; %d evts"%(Nsamples,n_tot))
    pave1c.SetBorderSize(0)
    pave1c.SetLineColor(0)
    pave1c.SetFillColor(0)
    pave1c.SetTextColor(2)
    pave1c.SetTextAlign(32)
    pave1c.Draw("same")

    canv.cd(2)
    hsum.Draw()

    canv.Update()

    print "-"*60    
    print "We have the summary plots"
    print "Next : Final mass histogram"
    print "-"*60    

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================

    canv2 = TCanvas("Mass distribution","Mass distribution")
    canv2.SetBottomMargin(0.23)

    hsum.Draw()

    print "-"*60    
    print "We have the final mass histogram"
    print "Next : set signal region between %.2f and %.2f MeV"%(SigMassLow,SigMassHigh)
    print "-"*60    

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================

    _line_ymax = hsum.GetMaximum()/2.

    line1 = TLine(SigMassLow,0.,SigMassLow,_line_ymax)
    line1.SetLineWidth(2)
    line1.Draw()

    line2 = TLine(SigMassHigh,0.,SigMassHigh,_line_ymax)
    line2.SetLineWidth(2)
    line2.Draw()

    canv2.Update()

    print "-"*60    
    print "Signal region set"
    print "Next : Count events in signal and background region"
    print "-"*60    

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================

    hsig = hsum.Clone()
    nbins_signalregion = 0
    for ibin in range(hsig.GetNbinsX()):
        xbin = hsig.GetBinCenter(ibin+1)
        if (xbin<SigMassLow) or (xbin>SigMassHigh): hsig.SetBinContent(ibin+1,0.)
        else: nbins_signalregion+=1
    hbdf = hsum.Clone()
    nbins_sidebands=0
    for ibin in range(hbdf.GetNbinsX()):
        xbin = hbdf.GetBinCenter(ibin+1)
        if (xbin>=SigMassLow) and (xbin<=SigMassHigh): hbdf.SetBinContent(ibin+1,0.)
        else: nbins_sidebands+=1
    hsig.SetFillColor(2)    
    hsig.SetFillStyle(3001)    
    hsig.SetLineColor(2)    
    hbdf.SetFillColor(4)    
    hbdf.SetFillStyle(3001)    
    hbdf.SetLineColor(4)    

    hsig.Draw('SAME')
    hbdf.Draw('SAME')

    xbinning = (hsum.GetXaxis().GetXmax()-hsum.GetXaxis().GetXmin())/hsum.GetNbinsX() # MeV/bin
        
    n_sidebands    = hbdf.Integral()    
    n_signalregion = hsig.Integral()    
    x_bdf = 1.*nbins_signalregion*n_sidebands/nbins_sidebands
    e_bdf = sqrt(n_sidebands)*1.*nbins_signalregion/nbins_sidebands
    x_sig = n_signalregion - x_bdf
    e_sig = sqrt(n_signalregion + e_bdf*e_bdf)

    print "--> side-bands size : %f MeV"%(nbins_sidebands*xbinning)
    print "--> signal region size : %f MeV"%(nbins_signalregion*xbinning)
    print "--> %5d events in the side-bands"%(n_sidebands)
    print "--> %5d events in the signal region"%(n_signalregion)
    print "--> %.2f +- %.2f background events in the signal region"%(x_bdf,e_bdf)
    print "==> %.2f +- %.2f D candidates"%(x_sig,e_sig)

    pave3b  =  TPaveText(.5,.90,.5,.90,"NDC")
    pave3b.SetName("TP3b")
    pave3b.AddText("Signal region: %6d events "%(n_signalregion))
    pave3b.SetBorderSize(0)
    pave3b.SetLineColor(1)
    pave3b.SetFillColor(0)
    pave3b.SetFillStyle(4000)
    pave3b.SetTextColor(1)
    pave3b.SetTextAlign(23)
    pave3b.SetTextSize(0.04)
    pave3b.Draw()

    pave3c  =  TPaveText(.12,.80,.12,.80,"NDC")
    pave3c.SetName("TP3c")
    pave3c.AddText("Side bands: %6d events"%(n_sidebands))
    pave3c.SetBorderSize(0)
    pave3c.SetLineColor(1)
    pave3c.SetFillColor(0)
    pave3c.SetFillStyle(4000)
    pave3c.SetTextColor(1)
    pave3c.SetTextAlign(13)
    pave3c.SetTextSize(0.04)
    pave3c.Draw()

    canv2.Update()

    print "-"*60    
    print "Event counted"
    print "Next : Compute nb of background events"
    print "-"*60    

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================

    pave3d  =  TPaveText(.52,.26,.52,.26,"NDC")
    pave3d.SetName("TP3d")
    pave3d.AddText("%4.1f background events"%(x_bdf))
    pave3d.SetBorderSize(0)
    pave3d.SetLineColor(1)
    pave3d.SetFillColor(0)
    pave3d.SetFillStyle(4000)
    pave3d.SetTextColor(4)
    pave3d.SetTextAlign(21)
    pave3d.SetTextSize(0.04)
    pave3d.Draw()

    bdf_line = TLine(hbdf.GetXaxis().GetXmin(),1.*n_sidebands/nbins_sidebands,hbdf.GetXaxis().GetXmax(),1.*n_sidebands/nbins_sidebands)
    bdf_line.SetLineColor(4)
    bdf_line.SetLineWidth(2)
    bdf_line.SetLineStyle(1)
    bdf_line.Draw()
    
    bdf_box = TBox(SigMassLow,0.,SigMassHigh,1.*n_sidebands/nbins_sidebands)
    bdf_box.SetFillColor(4)
    bdf_box.SetFillStyle(3004)
    bdf_box.Draw()
    
    canv2.Update()

    print "-"*60    
    print "Background event computed"
    print "Next : Compute nb of signal events"
    print "-"*60    

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================

    pave3a  =  TPaveText(.52,.65,.52,.65,"NDC")
    pave3a.SetName("TP3a")
    pave3a.AddText("%.2f #pm %.2f D candidates"%(x_sig,e_sig))
    pave3a.SetBorderSize(0)
    pave3a.SetLineColor(1)
    pave3a.SetFillColor(0)
    pave3a.SetFillStyle(4000)
    pave3a.SetTextColor(2)
    pave3a.SetTextAlign(21)
    pave3a.SetTextSize(0.04)
    pave3a.Draw()

    canv2.Update()

    print "-"*60    
    print "Signal event computed"
    print "Next : Fit"
    print "-"*60    

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================

    canv3 = TCanvas("Fit","Fitted mass")
    canv3.SetBottomMargin(0.23)

    hsum.Draw()

    
    canv3.Update()
    
    print "-"*60    
    print "Raw distribution plotted"
    print "Next : Fit distribution w/ Gaussian + horizontal line "
    print "-"*60    

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================
    f_massfit = TF1('f_massfit',massFunction,hsum.GetXaxis().GetXmin(),hsum.GetXaxis().GetXmax(),4)
    f_massfit.SetParNames("C","Mean","Sigma","Bdf")
    f_massfit.SetParameter(1,1866.)
    f_massfit.SetParameter(2,15.)
    f_massfit.SetLineColor(2)
    
    res = hsum.Fit(f_massfit,"S","goff")
    f_massfit.Draw('SAME')
    
    Cst    = f_massfit.GetParameter(0)
    mean   = f_massfit.GetParameter(1)
    sigma  = f_massfit.GetParameter(2)
    bdf    = f_massfit.GetParameter(3)
    x_gaus = Cst*sigma*sqrt(2*pi)/xbinning

    print "--> Number of candidates under the gaussian function = %f"%(x_gaus)

    ## f_massSig = TF1('f_massSig',massFunctionSig,hsum.GetXaxis().GetXmin(),hsum.GetXaxis().GetXmax(),3)
    ## f_massSig.SetParNames("C","Mean","Sigma")
    ## f_massSig.SetParameter(0,Cst)
    ## f_massSig.SetParameter(1,mean)
    ## f_massSig.SetParameter(2,sigma)
    ## f_massSig.SetLineColor(2)
    ## f_massSig.Draw('SAME')
    
    f_massBdf = TF1('f_massBdf',massFunctionBdf,hsum.GetXaxis().GetXmin(),hsum.GetXaxis().GetXmax(),1)
    f_massBdf.SetParNames("C","Mean","Sigma")
    f_massBdf.SetParameter(0,bdf)
    f_massBdf.SetLineColor(4)

    f_massBdf.Draw('SAME')

    pave4  =  TPaveText(.12,.88,.22,.88,"NDC")
    pave4.SetName("TP4")
    pave4.AddText("N_{bdf}= %.2f per %.1f MeV"%(bdf,wbin))
    pave4.SetBorderSize(0)
    pave4.SetLineColor(1)
    pave4.SetFillColor(0)
    pave4.SetTextColor(4)
    pave4.SetTextAlign(13)
    pave4.SetTextSize(0.04)
    pave4.Draw()

    pave4a  =  TPaveText(.12,.64,.22,.82,"NDC")
    pave4a.SetName("TP4a")
    pave4a.AddText("Mean = %.2f MeV"%(mean))
    pave4a.AddText("#sigma = %.2f MeV"%(sigma))
    pave4a.AddText("N_{signal}= %.2f #pm %.2f"%(x_gaus,sqrt(x_gaus)))
    pave4a.SetBorderSize(0)
    pave4a.SetLineColor(1)
    pave4a.SetFillColor(0)
    pave4a.SetTextColor(2)
    pave4a.SetTextAlign(13)
    pave4a.SetTextSize(0.04)
    pave4a.Draw()

    _par=[Cst,mean,sigma,bdf]
    D_line_top = 1.03*massFunction([mean],_par)
    print "D_line_top=",D_line_top
    D_line = TLine(DMASS,0.,DMASS,D_line_top)
    D_line.SetLineColor(kGreen+2)
    D_line.SetLineWidth(2)
    D_line.SetLineStyle(2)
    D_line.Draw()

    pave4b  =  TPaveText(.56,.26,.56,.26,"NDC")
    pave4b.SetName("TP4a")
    pave4b.AddText("M_{D}= %.2f MeV"%(DMASS))
    pave4b.SetBorderSize(0)
    pave4b.SetLineColor(1)
    pave4b.SetFillColor(0)
    pave4b.SetFillStyle(4000)
    pave4b.SetTextColor(kGreen+2)
    pave4b.SetTextAlign(11)
    pave4b.SetTextSize(0.04)
    pave4b.Draw()
    
    canv3.Update()
    
    print "-"*60    
    print "Fit done"
    print "Next : END "

    #===================================================================================================================
    if waittocontinue(): sys.exit()
    #===================================================================================================================