import argparse

parser = argparse.ArgumentParser(description="")
# parser.add_argument("inputfile" , help = "Path to the input ROOT file")
parser.add_argument("elesel"    , help = "Define if pfpf, pflp, lplp")
parser.add_argument("-d", "--doubleg"  , dest = "doubleg", help = "Define if use 1 or 2 gaussian for the signal model", default = '0')
args = parser.parse_args()

import os, sys, math
sys.path.insert(0, os.environ['HOME'] + '/.local/lib/python2.7/site-packages')

import ROOT
from ROOT import gSystem
gSystem.Load('libRooFit')
from ROOT import RooFit, RooRealVar, RooDataSet, RooArgList, RooTreeData, RooArgSet, RooAddPdf, RooFormulaVar
from ROOT import RooGaussian, RooExponential, RooChebychev
from uncertainties import ufloat

cut_ele      = ''
if args.elesel =='pfpf':
  cut_ele = '(e1pf==1 && e2pf==1) && '
if args.elesel =='pflp':
  cut_ele = '((e1pf==1 && e2pf==0) || (e2pf==1 && e1pf==0)) && '
if args.elesel =='lplp':
  cut_ele = '(e1pf==0 && e2pf==0) &&'

tree = ROOT.TChain('my_ttree')
tree.Add('group2_multicands_data_ALL.root')


BMass_    = 5.277
JPsiMass_ = 3.096916
PsiPMass_ = 3.686109

BMass      = RooRealVar("BMass"     , "BMass"   , BMass_  )
JPsiMass   = RooRealVar("JPsiMass"  , "JPsiMass", JPsiMass_)
PsiPMass   = RooRealVar("PsiPMass"  , "PsiPMass", PsiPMass_)

Bmass         = RooRealVar("Bmass"  , "e^{+}e^{-}K mass", 4.8, 5.8, "GeV")
Bmll          = RooRealVar("Bmll"   , "Bmll"  , 0, 6);
e1pf          = RooRealVar("e1pf"   , "e1pf"  , 0, 2);
e2pf          = RooRealVar("e2pf"   , "e2pf"  , 0, 2);


thevars = RooArgSet()
thevars.add(Bmass)
thevars.add(Bmll)
thevars.add(e1pf)
thevars.add(e2pf)

fulldata   = RooDataSet('fulldata', 'fulldataset', tree,  RooArgSet(thevars))

q2binning = [
                1,     ## 0
                8.41,  ## 1 = Jpsi (2.9-3.3)
                10.89, ## 2
                12.86, ## 3 = psiPrime
                14.18, ## 4
                19,
]


c1 = ROOT.TCanvas()

yields = {}
sigmas = {}

for ibin in range(len(q2binning)-1):
    cut = cut_ele + '(Bmll*Bmll > %s && Bmll*Bmll < %s)'%(q2binning[ibin], q2binning[ibin+1])
    print cut
    data       = fulldata.reduce(RooArgSet(Bmass,Bmll), cut)

    mean        = RooRealVar ("mass"         , "mean"          ,  BMass_,    5,   5.3, "GeV")
    sigma       = RooRealVar ("#sigma_{1}"   , "sigma"         ,  0.028,     0,   10, "GeV")
    signalGauss = RooGaussian("signalGauss"  , "signal gauss"  ,  Bmass,  mean,sigma)

    sigma2       = RooRealVar ("#sigma_{2}"    , "sigma2"         ,  0.048,     0,   0.07, "GeV")
    signalGauss2 = RooGaussian("signalGauss2"  , "signal gauss2"  ,  Bmass,  mean,sigma2)
    f1           = RooRealVar ("f1"            , "f1"             ,  0.8  ,     0.,   1.)
    gaus         = RooAddPdf  ("gaus"          , "gaus1+gaus2"    , RooArgList(signalGauss,signalGauss2), RooArgList(f1))

    pol_c1      = RooRealVar ("p1"           , "coeff x^0 term",    0.5,   -10, 10);
    pol_c2      = RooRealVar ("p2"           , "coeff x^1 term",    0.5,   -10, 10);
    pol_c3      = RooRealVar ("p3"           , "coeff x^2 term",    0.5,   -10, 10);
    slope       = RooRealVar ("slope"        , "slope"         ,    0.5,   -10, 10);
    bkg_exp     = RooExponential("bkg_exp"   , "exponential"   ,  slope,   Bmass  );
    bkg_pol     = RooChebychev("bkg_pol"     , "2nd order pol" ,  Bmass, RooArgList(pol_c1,pol_c2));

    nsig        = RooRealVar("Yield"         , "signal frac"   ,   4000,     0,   1000000);
    nbkg        = RooRealVar("nbkg"          , "bkg fraction"  ,   1000,     0,   550000);

    fitFunction = RooAddPdf ("fitfunction" , "fit function"  ,  RooArgList(signalGauss, bkg_exp), RooArgList(nsig, nbkg))
    if int(args.doubleg) == 2:
        fitFunction = RooAddPdf ("fitfunction" , "fit function"  ,  RooArgList(gaus, bkg_exp), RooArgList(nsig, nbkg))
    r = fitFunction.fitTo(data, RooFit.Extended(True), RooFit.Save(), RooFit.Range( 4.8, 5.8,))

    frame = Bmass.frame()
    data.plotOn(frame, RooFit.Binning(25), RooFit.MarkerSize(.7))
    fitFunction.plotOn(frame, );
    fitFunction.plotOn(frame, RooFit.Components("bkg_exp")     , RooFit.LineStyle(ROOT.kDotted), RooFit.LineColor(ROOT.kRed));
    if int(args.doubleg) != 2:
        fitFunction.plotOn(frame, RooFit.Components("signalGauss") , RooFit.LineStyle(ROOT.kDashed), RooFit.LineColor(ROOT.kBlue), RooFit.FillColor(ROOT.kBlue), RooFit.DrawOption("F"), RooFit.FillStyle(3004))
        fitFunction.plotOn(frame, RooFit.Components("signalGauss") , RooFit.LineStyle(ROOT.kDashed), RooFit.LineColor(ROOT.kBlue))

    elif int(args.doubleg) == 2:
        fitFunction.plotOn(frame, RooFit.Components("signalGauss") , RooFit.LineStyle(ROOT.kDashed), RooFit.LineColor(ROOT.kGreen+1))
        fitFunction.plotOn(frame, RooFit.Components("signalGauss2"), RooFit.LineStyle(ROOT.kDashed), RooFit.LineColor(ROOT.kOrange+1));
        fitFunction.plotOn(frame, RooFit.Components("gaus") , RooFit.LineStyle(ROOT.kDashed), RooFit.LineColor(ROOT.kBlue), RooFit.FillColor(ROOT.kBlue), RooFit.DrawOption("F"), RooFit.FillStyle(3004))
        fitFunction.plotOn(frame, RooFit.Components("gaus") , RooFit.LineStyle(ROOT.kDashed), RooFit.LineColor(ROOT.kBlue));

    parList = RooArgSet (nsig,sigma,sigma2, mean)
# fitFunction.plotOn(frame, RooFit.Components("signalGauss2"), RooFit.LineStyle(ROOT.kDashed), RooFit.LineColor(ROOT.kGreen+2));

    fitFunction.paramOn(frame, RooFit.Parameters(parList), RooFit.Layout(0.62,0.86,0.88))

    frame.GetYaxis().SetTitleOffset(1.35)
    frame.getAttText().SetTextSize(0.022)
    frame.getAttText().SetTextFont(42)
    frame.getAttLine().SetLineColor(0)
    frame.SetTitle('')

    frame.Draw()
    c1.SaveAs('save_fit_data_%s_%s_Mauro_NoCleaning.pdf'%(ibin,args.elesel))

    resSigma1 = ufloat (sigma.getVal(),  sigma.getError())
    resSigma2 = ufloat (0,  0)
    resF      = ufloat (1,  0)
    if int(args.doubleg) == 2:
        resSigma2 = ufloat (sigma2.getVal(), sigma2.getError())
        resF      = ufloat (f1.getVal(),     f1.getError())

    totSigma = resF*(resSigma1**2) + (1-resF)*(resSigma2**2)

    totSigma_v = math.sqrt(totSigma.n)
    totSigma_e = totSigma.s/2/math.sqrt(totSigma.n)


    yields[ibin] =  [nsig.getVal(), nsig.getError()]
    sigmas[ibin] =  [totSigma_v, totSigma_e]


for k,v in yields.items():
    print 'bin ', k, '\t B0 yield: ', v[0], '+/-', v[1]

for k,v in sigmas.items():
    print 'bin ', k, '\t B0 sigma: ', v[0], '+/-', v[1]




#pragma link C++ class RooTreeData ;
#pragma link C++ class RooTreeData::PlotOpt ;
#pragma link C++ class RooTruthModel ;