#include "GlobalFunctions.hh"
#include "BackgroundPdf.hpp"

using namespace std;
using namespace RooFit;
using namespace RooStats;


//pdfs background
RooExponential *  BackgroundPdf::bkg_exp1(RooRealVar *B_plus_M){
    RooRealVar *exp = getRooRealVar("exp_par1",false);
    BkgRooVars.push_back(exp);
    RooExponential *bkg_exp = new RooExponential("bkg_exp1","ExponentialBckGnd1",*B_plus_M, *exp);
    if (verboseLevel <2) bkg_exp->Print();
    return bkg_exp;
}
RooExponential *  BackgroundPdf::bkg_exp2(RooRealVar *B_plus_M){
    RooRealVar *exp = getRooRealVar("exp_par2",false);
    BkgRooVars.push_back(exp);
    RooExponential *bkg_exp = new RooExponential("bkg_exp2","ExponentialBckGnd2",*B_plus_M, *exp);
    if (verboseLevel <2) bkg_exp->Print();
    return bkg_exp;
}
RooExpAndGauss *  BackgroundPdf::ExpG(RooRealVar *B_plus_M){
    RooRealVar *mean = getRooRealVar("bkg_mean",false);
    RooRealVar *sigma = getRooRealVar("bkg_sigma",false);
    RooRealVar *decay = getRooRealVar("bkg_decay",false);
    BkgRooVars.push_back(mean);
    BkgRooVars.push_back(sigma);
    BkgRooVars.push_back(decay);
    RooExpAndGauss *ExpG = new RooExpAndGauss("ExpG", "ExpG", *B_plus_M, *mean, *sigma, *decay);
    if (verboseLevel <2) ExpG->Print();
    return ExpG;
}
RooDoubleCB    *  BackgroundPdf::bkg_CB(RooRealVar *B_plus_M){
    RooRealVar *mean = getRooRealVar("bkg_mean",false);
    RooRealVar *sigma = getRooRealVar("bkg_sigma",false);
    RooRealVar *alpha1 = getRooRealVar("bkg_alpha1",false);
    RooRealVar *alpha2 = getRooRealVar("bkg_alpha2",false);
    RooRealVar *n1 = getRooRealVar("bkg_n1",false);
    RooRealVar *n2 = getRooRealVar("bkg_n2",false);
    BkgRooVars.push_back(mean);
    BkgRooVars.push_back(sigma);
    BkgRooVars.push_back(alpha1);
    BkgRooVars.push_back(alpha2);
    BkgRooVars.push_back(n1);
    BkgRooVars.push_back(n2);
    RooDoubleCB *CBBplus = new RooDoubleCB("CBbkg", "CBbkg", *B_plus_M, *mean, *sigma, *alpha1, *n1, *alpha2, *n2);
    if (verboseLevel <2) CBBplus->Print();
    return CBBplus;
}

RooAddPdf*  BackgroundPdf::getBplusBkgModel(RooRealVar *B_plus_M){
    RooAddPdf *BplusBckGndModel = NULL;
    RooExponential *exp1 = NULL; //could be defined right away, I keep it as it is in case fo future addons
    if (NoBackground){
        coutDebug("No background PDF!");
        return NULL;
    }
    RooRealVar  *zeroSlope   = new RooRealVar("zeroSlope","slope", 0.05, 0.0, 0.1); //Doesn't matter anyhow
    zeroSlope->setConstant("kTrue");
    RooRealVar *zeroF = new RooRealVar("zeroF","zeroF", 1.0, 1.0, 1.0);
    RooPolynomial *zeroLine = new RooPolynomial("zeroLine", "zeroLine", *B_plus_M, RooArgList(*zeroSlope));
    coutDebug("Created zero slope for bkgPdf.");

    if (SingleExponential){//No else-if, because then I can add OneCB
        coutDebug("Creating single exponential");
        exp1 = bkg_exp1(B_plus_M);
        if (!bkgOneCB && !ExpGaus) BplusBckGndModel = new RooAddPdf("BplusBckGndModel", "SingleExponentialBckGnd", RooArgList(*exp1, *zeroLine),RooArgList(*zeroF));
    }

    if (bkgOneCB) { //No else-ifs, because then I can add the previously defined functions
        coutDebug("Creating OneCB");
        RooDoubleCB *OneCB = NULL;
        OneCB = bkg_CB(B_plus_M);
        if (!SingleExponential) BplusBckGndModel = new RooAddPdf("BplusBckGndModel", "OneCBBckGnd", RooArgList(*OneCB, *zeroLine),RooArgList(*zeroF));
        else{
            RooRealVar *f = getRooRealVar("bkg_f",false);
            BkgRooVars.push_back(f);
            BplusBckGndModel = new RooAddPdf("BplusBckGndModel", "SingleExponentialWithCBBckGnd", RooArgList(*OneCB, *exp1),RooArgList(*f));
            coutDebug("Setting OneCB+SingleExp");
            if (verboseLevel<3) BplusBckGndModel->Print();
        }
    }
    else if (DoubleExponential){
        coutDebug("Creating double exponential");
        RooExponential *exp2 = bkg_exp2(B_plus_M);
        RooRealVar *f = getRooRealVar("bkg_f",false);
        BkgRooVars.push_back(f);
        BplusBckGndModel = new RooAddPdf("BplusBckGndModel", "DoubleExponentialBckGnd", RooArgList(*exp1, *exp2),RooArgList(*f));
    }
    else if (ExpGaus){
        coutDebug("Creating ExpGaus");
        RooExpAndGauss *expGauss = ExpG(B_plus_M);
        RooRealVar *f = getRooRealVar("bkg_f",false);
        BkgRooVars.push_back(f);
        BplusBckGndModel = new RooAddPdf("BplusBckGndModel", "ExpGausBckGnd", RooArgList(*exp1, *expGauss),RooArgList(*f));
    }
    if (NoBackground) coutDebug("No background PDF!");
    else if (verboseLevel < 2) BplusBckGndModel->Print();
    return BplusBckGndModel;
}

void BackgroundPdf::setAllRooVarsConstant(){
    for (auto var : BkgRooVars){
        var->setConstant(kTRUE);
    }
}