#include "ParamValues.hpp"

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

void ParamValues::printProperties(){
    coutDebug("Name\t\t\t" + name);
    coutDebug("LaTeX name\t\t" + latexName);
    coutDebug("Fixed\t\t\t" + string(fix ? "yes" : "no"));
    coutDebug("Constrained\t\t" + string(constrain ? "yes" : "no"));
    coutDebug("val\t\t\t" + to_string(value));
    coutDebug("valError\t\t" + to_string(error));
    coutDebug("lowRange\t\t" + to_string(ranges[0]));
    coutDebug("upRange\t\t\t" + to_string(ranges[1]));
}

void ParamValues::SetParamValues(string s_string, bool fitRef){
    if (s_string == "mean"){
        name = "sig_mean";
        latexName = "B^{+} #mu";
        value = PDGMASS.B_PLUS;
        error = 0.5;
        constrain = KshortChannel ? false : false;
        fix       = KshortChannel ? false : false;
        ranges[0] = 5250.0;
        ranges[1] = 5320.0;
    }
    else if (s_string == "sigma"){
        name = "sig_sigma";
        latexName = "B^{+} #sigma_{1}";
        value = OneCB ? 30.0 : 15.0;
        error = 1.0;
        constrain = KshortChannel ? false : false;
        fix       = KshortChannel ? true : !fitRef;
        ranges[0] = 5.0;
        ranges[1] = 150.;
    }
    else if (s_string == "sigma2"){
        name = "sig_sigma2";
        latexName = "B^{+} #sigma_{2}";
        value = anyGaussain ? 18.0 : 12.0;
        error = 1.0;
        constrain = KshortChannel ? false : false;
        fix       = KshortChannel ? anyGaussain : anyGaussain;
        ranges[0] = 5.0;
        ranges[1] = DoubleGaussian ? 150. : 100.;
    }
    else if (s_string == "alpha"){
        name = "sig_alpha";
        latexName = "B^{+} #alphaCB_{1}";
        value =  (RightCB || OneCB) ? 1.5 : -1.5;
        error = 0.1;
        constrain = KshortChannel ? false : false;
        fix       = KshortChannel ? true : true;
        ranges[0] = (RightCB || OneCB) ? 0.01 : -3.5;
        ranges[1] = (RightCB || OneCB) ? 3.5  : -0.01;
    }
    else if (s_string == "alpha2"){
        name = "sig_alpha2";
        latexName = "B^{+} #alphaCB_{2}";
        value = 1.5;
        error = 1.0;
        constrain = KshortChannel ? false : false;
        fix = KshortChannel ? true : true;
        ranges[0] = 0.01;
        ranges[1] = 3.5;
    }
    else if (s_string == "n"){
        name = "sig_n";
        latexName = "B^{+} #nCB_{1}";
        value = 3.1;
        error = 0.1;
        constrain = KshortChannel ? false : false;
        fix = KshortChannel ? true : true;
        ranges[0] = 2.1;
        ranges[1] = 4.0;
    }
    else if (s_string == "n2"){
        name = "sig_n2";
        latexName =" B^{+} #nCB_{2}";
        value = 4.0;
        error = 0.1;
        constrain = KshortChannel ? false : false;
        fix = KshortChannel ? true : true;
        ranges[0] = 1.0;
        ranges[1] = 6.0;
    }

    else if (s_string == "f"){
        name = "sig_f";
        latexName = "f__{B^{+}}";
        if (DoubleGaussian) value = 0.5;
        else if (DoubleCB) value = 0.55; //0.5; CB_left >= CB_right
        else value = 0.0;
        error = 0.02;
        constrain = KshortChannel ? DoubleCB : DoubleCB;
        fix       = KshortChannel ? false : false;
        ranges[0] = KshortChannel ?  0.05 : 0.15;
        ranges[1] = KshortChannel ?  0.95 : (DoubleGaussian ? 0.95 : 0.55);
    }
    else if (s_string == "exp_par1"){
        name = "exp_par1";
        latexName = "#lambda_{1}";
        value = -0.0015;
        error = 0.0001;
        constrain = false;
        fix       = false;
        ranges[0] = -0.02;
        ranges[1] = 0.0;
    }
    else if (s_string == "exp_par2"){
        name = "exp_par2";
        latexName = "#lambda_{0}";
        value = -0.0016;
        error = 0.0001;
        constrain = false;
        fix       = false;
        ranges[0] = -0.016;
        ranges[1] = 0.0;
    }
    else if (s_string == "bkg_mean"){ //ExpGaus mean
        name = "bkg_mean";
        latexName = "B^{+} #mu";
        value = ExpGaus? 5144.0 : PDGMASS.B_PLUS;
        error = 30.0;
        constrain = true;
        fix       = false;
        ranges[0] = ExpGaus ? 5115.0 : PDGMASS.B_PLUS-10;
        ranges[1] = ExpGaus ? 5175.0 : PDGMASS.B_PLUS+10;
    }
    else if (s_string == "bkg_sigma"){
        name = "bkg_sigma";
        latexName = "B^{+} #sigma_{bkg}";
        value = 11.0;
        error = 1.0;
        constrain = ExpGaus;
        fix       = !ExpGaus;
        ranges[0] = 2.0;
        ranges[1] = 50.;
    }
    else if (s_string == "bkg_decay"){
        name = "bkg_decay";
        latexName = "B^{+} decay coeff";
        value =  5140;
        error = 0.1;
        constrain = false;
        fix       = false;
        ranges[0] = 5000.0;
        ranges[1] = 5150.0;
    }
    else if (s_string == "bkg_alpha1"){
        name = "bkg_alpha1";
        latexName = "B^{+} #alphaCB_{1}^{bkg}";
        value = 1.5;
        error = 1.0;
        constrain = false;
        fix       = true;
        ranges[0] = 0.01;
        ranges[1] = 3.5;
    }
    else if (s_string == "bkg_alpha2"){
        name = "bkg_alpha2";
        latexName = "B^{+} #alphaCB_{2}^{bkg}";
        value = 1.5;
        error = 1.0;
        constrain = false;
        fix       = true;
        ranges[0] = 0.01;
        ranges[1] = 3.5;
    }
    else if (s_string == "bkg_n1"){
        name = "bkg_n1";
        latexName = "B^{+} #nCB_{1}^{bkg}";
        value = 5.1;
        error = 0.1;
        constrain = false;
        fix       = true;
        ranges[0] = 0.1;
        ranges[1] = 20.;
    }
    else if (s_string == "bkg_n2"){
        name = "bkg_n2";
        latexName =" B^{+} #nCB_{2}^{bkg}";
        value = 1.4;
        error = 0.1;
        constrain = false;
        fix       = true;
        ranges[0] = 0.1;
        ranges[1] = 30.;
    }
    else if (s_string == "bkg_f"){
        name = "bkg_f";
        latexName = "f^{bckgnd}_{B^{+}}";
        value = 0.5;
        error = 0.01;
        constrain = false;
        fix       = false;
        ranges[0] = 0.15;
        ranges[1] = 0.99;
    }
    if (verboseLevel<3) printProperties();
};

FitValuesWithName getSignalSavedFitValue(string s_name){ //This would have been so much nicer in Python
    if      (s_name == "sig_mean")    return FitValuesSignal.sig_mean  ;
    else if (s_name == "sig_sigma")   return FitValuesSignal.sig_sigma ;
    else if (s_name == "sig_sigma2")  return FitValuesSignal.sig_sigma2;
    else if (s_name == "sig_alpha")   return FitValuesSignal.sig_alpha ;
    else if (s_name == "sig_alpha2")  return FitValuesSignal.sig_alpha2;
    else if (s_name == "sig_n")       return FitValuesSignal.sig_n     ;
    else if (s_name == "sig_n2")      return FitValuesSignal.sig_n2    ;
    else if (s_name == "sig_f")       return FitValuesSignal.sig_f     ;
    else{
        coutERROR("Parameter " + s_name + "  was not found in the FitValuesSignal vector!");
        coutERROR("Returning 0.");
        return FitValuesWithName("",0.0,0.0);
    }
}

FitValuesWithName getBackgroundSavedFitValue(string s_name){ //This would have been so much nicer in Python
    if (s_name == "bkg_mean")        return FitValuesBackground.bkg_mean  ;
    else if (s_name == "bkg_sigma")  return FitValuesBackground.bkg_sigma ;
    else if (s_name == "bkg_decay")  return FitValuesBackground.bkg_decay ;
    else if (s_name == "bkg_alpha1")  return FitValuesBackground.bkg_alpha ;
    else if (s_name == "bkg_alpha2") return FitValuesBackground.bkg_alpha2;
    else if (s_name == "bkg_n1")      return FitValuesBackground.bkg_n     ;
    else if (s_name == "bkg_n2")     return FitValuesBackground.bkg_n2    ;
    else if (s_name == "bkg_f")      return FitValuesBackground.bkg_f     ;
    else{
        coutERROR("Parameter " + s_name + "  was not found in the FitValuesBackground vector!");
        coutERROR("Returning 0.");
        return FitValuesWithName("",0.0,0.0);
    }
}

RooRealVar *ParamValues::getRooRealVar(string s_name, bool sig, bool fitRef){
    SetParamValues(s_name, fitRef);
    RooRealVar *variable  = new RooRealVar(name.c_str(), latexName.c_str(), value, ranges[0], ranges[1]);
    if (fix && (sig ? FixParameters : FixBkgParameters)){
        coutDebug("Fixing parameter " + string(variable->GetName()));
        sig ? variable->setVal(getSignalSavedFitValue(variable->GetName()).val) : variable->setVal(getBackgroundSavedFitValue(variable->GetName()).val);
        variable->setConstant("kTrue");
    }
    if (constrain && (sig ? ConstrParameters: ConstrBkgParameters)){
        coutDebug("Adding constrain to " + string(variable->GetName()));
        if (fix) coutWarning("Cannot constrain parameter " + string(variable->GetName()) + " which is already fixed. Ignore constraining!");
        else{
            string ConstrName = name + "Constr";            
            RooGaussian *gauss_contrain_sig = sig ? new RooGaussian(ConstrName.c_str(), ConstrName.c_str(), *variable, RooConst(getSignalSavedFitValue(variable->GetName()).val), RooConst(getSignalSavedFitValue(variable->GetName()).err)) : new RooGaussian();
            if (verboseLevel<2 && sig) gauss_contrain_sig->Print();

            RooGaussian *gauss_contrain_bkg = sig ? new RooGaussian() : new RooGaussian(ConstrName.c_str(), ConstrName.c_str(), *variable, RooConst(getBackgroundSavedFitValue(variable->GetName()).val), RooConst(getBackgroundSavedFitValue(variable->GetName()).err));
            if (verboseLevel<2 && !sig) gauss_contrain_bkg->Print();

            sig ? ConsPDF->add(*gauss_contrain_sig) : BkgConsPDF->add(*gauss_contrain_bkg);
            sig ? ConsParameter->add(*variable) : BkgConsParameter->add(*variable);
            sig ? AtLeastOneConstrainFound = true : BkgAtLeastOneConstrainFound = true;
            if (verboseLevel<2) sig ? ConsPDF->Print() : BkgConsPDF->Print() ;
        }
    }
    coutDebug("Created variable: ");
    if (verboseLevel<2) variable->Print();
    return variable;
}

void setSignalSavedFitValue(string s_name, double d_val, double d_err){ //This would have been so much nicer in Python

    coutDebug("Setting " + s_name +" to " + to_string(d_val) + " +- " + to_string(d_err));
    if      (s_name == "sig_mean")    {
        FitValuesSignal.sig_mean.val = d_val;
        FitValuesSignal.sig_mean.err = d_err;}
    else if (s_name == "sig_sigma")   {
        FitValuesSignal.sig_sigma.val = d_val;
        FitValuesSignal.sig_sigma.err = d_err;}
    else if (s_name == "sig_sigma2")  {
        FitValuesSignal.sig_sigma2.val = d_val;
        FitValuesSignal.sig_sigma2.err = d_err;}
    else if (s_name == "sig_alpha")   {
        FitValuesSignal.sig_alpha.val = d_val;
        FitValuesSignal.sig_alpha.err = d_err;}
    else if (s_name == "sig_alpha2")  {
        FitValuesSignal.sig_alpha2.val = d_val;
        FitValuesSignal.sig_alpha2.err = d_err;}
    else if (s_name == "sig_n")       {
        FitValuesSignal.sig_n.val = d_val;
        FitValuesSignal.sig_n.err = d_err;}
    else if (s_name == "sig_n2")      {
        FitValuesSignal.sig_n2.val = d_val;
        FitValuesSignal.sig_n2.err = d_err;}
    else if (s_name == "sig_f")       {
        FitValuesSignal.sig_f.val = d_val;
        FitValuesSignal.sig_f.err = d_err;}
    else if (s_name == "sig_effSigma"){
        FitValuesSignal.sig_effSigma.val = d_val;
        FitValuesSignal.sig_effSigma.err = d_err;}
    else{
        coutWarning("Parameter " + s_name + " was not found in the FitValuesSignal vector!");
        coutWarning("Cannot set anything!");
    }
    return;
}

void setBackgroundSavedFitValue(string s_name, double d_val, double d_err){ //This would have been so much nicer in Python

    coutDebug("Setting " + s_name +" to " + to_string(d_val) + " +- " + to_string(d_err));
    if      (s_name == "bkg_mean")    {
        FitValuesBackground.bkg_mean.val = d_val;
        FitValuesBackground.bkg_mean.err = d_err;}
    else if (s_name == "bkg_sigma")   {
        FitValuesBackground.bkg_sigma.val = d_val;
        FitValuesBackground.bkg_sigma.err = d_err;}
    else if (s_name == "bkg_alpha1")  {
        FitValuesBackground.bkg_alpha.val = d_val;
        FitValuesBackground.bkg_alpha.err = d_err;}
    else if (s_name == "bkg_alpha2")   {
        FitValuesBackground.bkg_alpha2.val = d_val;
        FitValuesBackground.bkg_alpha2.err = d_err;}
    else if (s_name == "bkg_n1")  {
        FitValuesBackground.bkg_n.val = d_val;
        FitValuesBackground.bkg_n.err = d_err;}
    else if (s_name == "bkg_n2")       {
        FitValuesBackground.bkg_n2.val = d_val;
        FitValuesBackground.bkg_n2.err = d_err;}
    else if (s_name == "bkg_f")      {
        FitValuesBackground.bkg_f.val = d_val;
        FitValuesBackground.bkg_f.err = d_err;}
    else{
        coutWarning("Parameter " + s_name + " was not found in the FitValuesBackground vector!");
        coutWarning("Cannot set anything!");
    }
    return;
}

//Set parameters to get the values needed by the constrained fit
void ParamValues::setValuesAndErrors(bool sig){
    for (auto var : (sig ? SigRooVars : BkgRooVars)){
        coutDebug("Setting " + string(var->GetName()));
        if (verboseLevel < 3) var->Print();
        sig ? setSignalSavedFitValue(var->GetName(),var->getVal(),var->getError()) :  setBackgroundSavedFitValue(var->GetName(),var->getVal(),var->getError());
    }
    if (sig) setSignalSavedFitValue("sig_effSigma",getEffSigmaFromResult(),0.04);

    return;
}