#include "GlobalFunctions.hh"
#include "SignalPdf.hpp"

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

//pdfs signal
RooGaussian* SignalPdf::GaussBplus1(RooRealVar *B_plus_M, RooRealVar *mean, bool fitRef){
    RooRealVar *sigma = getRooRealVar("sigma",true,fitRef);
    SigRooVars.push_back(sigma);
    RooGaussian* gaus = new RooGaussian("GaussBplus1","GaussBplus1",*B_plus_M, *mean, *sigma);
    if (verboseLevel <2) gaus->Print();
    return gaus;
}
RooGaussian* SignalPdf::GaussBplus2(RooRealVar *B_plus_M, RooRealVar *mean, bool fitRef){
    RooRealVar *sigma = getRooRealVar("sigma2",true,fitRef);
    SigRooVars.push_back(sigma);
    RooGaussian *gaus = new RooGaussian("GaussBplus2","GaussBplus2",*B_plus_M, *mean, *sigma);
    if (verboseLevel <2) gaus->Print();
    return gaus;
}
RooCBShape* SignalPdf::CBBplus1(RooRealVar *B_plus_M, RooRealVar *mean, bool fitRef){ //could be shortened using the vectors
    RooRealVar *sigma = getRooRealVar("sigma",true,fitRef);
    RooRealVar *alpha = getRooRealVar("alpha",true);
    RooRealVar *n     = getRooRealVar("n",true);
    SigRooVars.push_back(sigma);
    SigRooVars.push_back(alpha);
    SigRooVars.push_back(n);

    RooCBShape *CB = new RooCBShape("CBBplus1", "CBBplus1", *B_plus_M, *mean, *sigma, *alpha, *n);
    if (verboseLevel <2) CB->Print();
    return CB;
}
RooCBShape* SignalPdf::CBBplus2(RooRealVar *B_plus_M, RooRealVar *mean, bool fitRef){
    RooRealVar *sigma = getRooRealVar("sigma2",true,fitRef);
    RooRealVar *alpha = getRooRealVar("alpha2",true);
    RooRealVar *n     = getRooRealVar("n2",true);
    SigRooVars.push_back(sigma);
    SigRooVars.push_back(alpha);
    SigRooVars.push_back(n);

    RooCBShape *CB = new RooCBShape("CBBplus2", "CBBplus2", *B_plus_M, *mean, *sigma, *alpha, *n);
    if (verboseLevel <2) CB->Print();
    return CB;
}
RooDoubleCB* SignalPdf::CBBplus(RooRealVar *B_plus_M, RooRealVar *mean, bool fitRef){
    RooRealVar *sigma  = getRooRealVar("sigma",true,fitRef);
    RooRealVar *alpha  = getRooRealVar("alpha",true);
    RooRealVar *alpha2 = getRooRealVar("alpha2",true);
    RooRealVar *n      = getRooRealVar("n",true);
    RooRealVar *n2     = getRooRealVar("n2",true);
    SigRooVars.push_back(sigma);
    SigRooVars.push_back(alpha);
    SigRooVars.push_back(alpha2);
    SigRooVars.push_back(n);
    SigRooVars.push_back(n2);

    RooDoubleCB *CB = new RooDoubleCB("OneCBBplus", "OneCBBplus", *B_plus_M, *mean, *sigma, *alpha, *n, *alpha2, *n2);

    if (verboseLevel <2) CB->Print();
    return CB;
}

//Final shape
RooAddPdf* SignalPdf::getBplusMassModel(RooRealVar *B_plus_M, RooRealVar *mean, bool fitRef){
    RooAddPdf * BplusMassModel = NULL;
    if (NoSig) return NULL;

    //Adding a line of 0 to pdf, to have a pdf and not RooCBShape or so
    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));    

    //Create gaussians
    if (anyGaussain){
        RooGaussian* Gauss1 = GaussBplus1(B_plus_M, mean, fitRef);
        if (DoubleGaussian){
            RooRealVar *f_Bplus  = getRooRealVar("sig_f",true);
            SigRooVars.push_back(f_Bplus);
            RooGaussian* Gauss2 = GaussBplus2(B_plus_M, mean, fitRef);
            BplusMassModel = new RooAddPdf("BplusMassModel", "DoubleGaussBplus", RooArgList(*Gauss1,*Gauss2),RooArgList(*f_Bplus));
        }
        else{
            BplusMassModel = new RooAddPdf("BplusMassModel", "SingleGaussBplus", RooArgList(*Gauss1, *zeroLine),RooArgList(*zeroF));
        }
    }
    else{ //Create CB
        if (LeftCB || RightCB){
            string name = LeftCB ? "LeftCBBplus" : "RightCBBplus";
            RooCBShape *OneSidedCB = CBBplus1(B_plus_M, mean, fitRef);
            BplusMassModel = new RooAddPdf("BplusMassModel", name.c_str(), RooArgList(*OneSidedCB, *zeroLine),RooArgList(*zeroF));
        }
        else if (DoubleCB){
            RooCBShape *RightCB = CBBplus1(B_plus_M, mean, fitRef);
            RooCBShape *LeftCB  = CBBplus2(B_plus_M, mean, fitRef);
            RooRealVar *f_Bplus = getRooRealVar("sig_f",true);
            SigRooVars.push_back(f_Bplus);
            BplusMassModel = new RooAddPdf("BplusMassModel", "DoubleCrystalBall", RooArgList(*LeftCB, *RightCB), RooArgList(*f_Bplus));
        }
        else{ //if (OneCB)
            RooDoubleCB *OneCB = CBBplus(B_plus_M, mean, fitRef);
            BplusMassModel = new RooAddPdf("BplusMassModel", "CBBplus", RooArgList(*OneCB, *zeroLine),RooArgList(*zeroF));
        }
    }
    if (verboseLevel <2) BplusMassModel->Print();
    return BplusMassModel;

}


void SignalPdf::setAllRooVarsConstant(){
    for (auto var : SigRooVars){
        var->setConstant(kTRUE);
    }

}