//Libraries used for the B+ mass fit model
//Renata Kopecna
#ifndef MASSFIT_HPP
#define MASSFIT_HPP

#include "GlobalFunctions.hh"
#include "RooFit/RooDoubleCB/RooDoubleCB.h"
#include "BmassShape/SignalPdf.hpp"

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

//////////////////////////////////////////////////////
///  \brief GetsWeightPlots()
//////////////////////////////////////////////////////
///
/// Returns path where the sWeighting control plots are stored
///
//////////////////////////////////////////////////////
///  \brief massFit()
//////////////////////////////////////////////////////
/// Based on quickFit in nTrackWeights.cc, but made into a seperate file
/// in on order to be able to be run as a standalone code
///
/// Main purpose is to fit the Bmass distribution and save the yield of the fit.
/// the fitted B+ mass spectrum may be saved in pdf and root files
///
///
/// FEATURES:
/// - choose your background and signal model:
///         signal:     (double) Gaussian, left/right/double/two-tailed CrystalBall function
///         background: (double) Exponential, Exponential plus RooExpGaus or no background
/// - get signal shape from MC data if needed and restrict shape to these parameters
/// - choose if fit all data, Jpsi q^2 region or q^2 region without resonances
/// - choose if fit low/high q^2 region
/// - choose to fit LL/DD Kshort sample
/// - choose to fit up/down/both
/// - choose to fit per year or per run///
///
/// OPTIONAL:
/// - cut on TMVAresponse
/// - fit only randomly selected half of the data sample
/// - Calculate yield in a 2sigma/fixed window
/// - Calculate yield over the full region
/// - Turn off the fit of MC needed to get the shape parameters (usefull in loops)
/// - Save the mass plots in efficiency file
/// - Fit in bins of selected variable
///
/// - from the fit to the B+ mass distribution, a sPlot weight is calculated and saved to a new Branch in this merged file.
///     If only Jpsi is selected, Reference channel is used and if sWeight==true,
///     J/psi data only is saved
///
//////////////////////////////////////////////////////
///  \brief quickFit()
//////////////////////////////////////////////////////
///
/// massFit() with simplified options
/// suitable for reweighting MC and creating sPlots
///
/////////////////////////////////////////////////////////
///  efficiencyFit()
//////////////////////////////////////////////////////
///
/// \brief massFit() with simplified options
/// suitable for efficiency estimation fits
///
/////////////////////////////////////////////////////////
///  \brief basicYieldFit()
//////////////////////////////////////////////////////
///
/// massFit() with simplified options
/// suitable for getting yield information
///
/// Available also to be done for:
///     All years
///     All years and q2 regions
///     Both Runs
///
/////////////////////////////////////////////////////////
///  \brief massFitTestAll()
//////////////////////////////////////////////////////
///
/// Fits given data sample with different
/// signal and background shapes, useful
/// when deciding what shape describes
/// the data the best
///
//////////////////////////////////////////////////////
///  \brief PrintFitResults()
//////////////////////////////////////////////////////
/// Takes RooFitResult as an iput and prints its content
///
///
///


//////////////////////////////////////////////////////
/// TODO at some point
/// \param year
/// \param UseOnlyJpsiEvents
/// \param UseOnlyMuMuEvents
/// \param KshortDecaysInVelo
/// \param GetShapeFromMC
/// \param SigType
/// \param BckGndType
/// \param ConstrainParameters
///
///
/// 


class TMefficiencyClass{

public:
    string sVariable;
    std::vector<string> sBranchName;
    Int_t Bins;
    Float_t Range[2];
    std::vector<std::array<double, 2>> vVarRange;
    std::vector<double> binEdges;
    std::vector<double> binEdgesEquidistant;
    bool isEquidistant;
    Float_t Step;
    string cut;

    TMefficiencyClass(){ //default constructor
        sVariable   = "";
        sBranchName = {};
        Bins = 0;
        Range[0] = 0.;
        Range[1] = 0.;
        vVarRange = {};
        binEdges = {};
        binEdgesEquidistant = {};
        cut = "";

    }

    void fillBinEdges(){
        double step = (Range[1] - Range[0]) / Bins;
        for (int i = 0; i <= Bins; i++){
            binEdgesEquidistant.push_back(Range[0]+i*step);
        }
        return;
    }
    void isEq(string &varName){ //'_equal' has to appear at the end of the variable!
        //See if the '_equal' is present
        size_t found = varName.find("_equal");
        if (found != string::npos) isEquidistant = true;
        else isEquidistant = false;
        //Remove '_equal' from the string
        replace(varName,"_equal","");
    }
    bool check_vector_size(bool isEqui){
        if (!isEqui){
            if (int(binEdges.size()) != Bins+1){
                coutDebug("Bins edges should be " + to_string(Bins+1) + " and are " + to_string(binEdges.size()));
                coutERROR("Wrong number of bin edges! Check " + sVariable);
                return false;
            }
            else return true;
        }
        else{
            if (int(binEdgesEquidistant.size()) != Bins+1){
                coutDebug("Bins edges should be " + to_string(Bins+1) + " and are " + to_string(binEdgesEquidistant.size()));
                coutERROR("Wrong number of bin edges! Check " + sVariable);
                return false;
            }
            else return true;
        }
    }

    TMefficiencyClass(string varName){ // constructor

        //first, replace the variable name
        isEq(varName);
        coutDebug("Is equidistant? " + to_string(isEquidistant));

        //q^2:
        if (varName == "q2"){
            sVariable   = "q^{2} [MeV^{2}]";
            sBranchName = {"Q2"};
            Bins = 10;
            Range[0] = 0.; // in MeV^2
            Range[1] = 20000000.; // in MeV^2
            vVarRange = {{Range[0],Range[1]}};
            fillBinEdges();
            binEdges = {};
            cut = sBranchName.at(0);
        }
        else if (varName == "q2_binned"){ //binned as in the analysis
            sVariable   = "q^{2} [GeV^{2}]";
            sBranchName = {"Q2"};
            Bins = 11;
            Range[0] = 0.1e6; // in MeV^2
            Range[1] = 20e6;
            vVarRange = {{Range[0],Range[1]}};
            fillBinEdges();
            binEdges = {0.1e6, 0.98e6, 1.1e6, 2.5e6, 4.0e6, 6.0e6, 8.0e6, 11.0e6, 12.5e6, 15.0e6, 17.0e6, 20.0e6};
            cut = sBranchName.at(0);
        }
        else if (varName == "q2_binned_fit"){ //binned as in the analysis
            sVariable   = "q^{2} [GeV^{2}]";
            sBranchName = {"Q2"};
            Bins = 6;
            Range[0] = 1.1e6; // in MeV^2
            Range[1] = 20e6;
            vVarRange = {{Range[0],Range[1]}};
            fillBinEdges();
            binEdges = {0.1e6, 4.0e6, 8.0e6, 11.0e6, 12.5e6, 15.0e6, 20.0e6};
            cut = sBranchName.at(0);
        }


        //cos(Theta_k):
        else if (varName == "thetak"){
            sVariable   = "cos(#theta_{k})";
            sBranchName = {"B_plus_ThetaK"};
            Bins = 10;
            Range[0] = -1.0;
            Range[1] = 1.0;
            //Range[0] = 0.;
            //Range[1] = 3.142;
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {0.000, 1.053, 1.294, 1.488, 1.662, 1.827, 1.993, 2.163, 2.353, 2.585, 3.150};
            vVarRange = {{0.,3.142}};
            cut = "cos(B_plus_ThetaK)";
        }
        //cos(Theta_l):
        else if (varName == "thetal"){
            sVariable   = "cos(#theta_{l})";
            sBranchName = {"B_plus_ThetaL"};
            Bins = 10;
            Range[0] = -1.0;
            Range[1] = 1.0;
            //Range[0] = 0.;
            //Range[1] = 3.142;
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {0.000, 0.630, 0.896, 1.110, 1.302, 1.486, 1.670, 1.866, 2.086, 2.368, 3.150};
            vVarRange = {{0.,3.142}};
            cut = "cos(B_plus_ThetaL)";
        }
        //phi
        else if (varName == "phi"){
            sVariable   = "#phi";
            sBranchName = {"B_plus_Phi"};
            Bins = 10;
            Range[0] = -3.15;
            Range[1] = 3.15;
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {-3.150, -2.476, -1.832, -1.207, -0.593, 0.015, 0.624, 1.241, 1.870, 2.517, 3.150};
            vVarRange = {{Range[0],Range[1]}};
            cut = sBranchName.at(0);
        }
        else if (varName == "pi_zero_ETA"){
            sVariable   = "#eta(#pi^{0})";
            sBranchName = {"pi_zero_resolved_ETA_DTF"};
            Bins = 10;
            Range[0] = 1.5;
            Range[1] = 4.5;
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {};
            vVarRange = {{Range[0],Range[1]}};
            cut = sBranchName.at(0);
        }
        else if (varName == "pi_zero_ETA-K_plus_ETA"){
            sVariable   = "|#eta(#pi^{0})-#eta(K^{+})|";
            sBranchName = {"pi_zero_resolved_ETA_DTF","K_plus_ETA_DTF"};
            Bins = 15;
            Range[0] = 0.0;
            Range[1] = 1.0;
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {};
            vVarRange = {{0.0,5.0},{0.0,5.0}};
            cut = "TMath::Abs(pi_zero_resolved_ETA_DTF-K_plus_ETA_DTF)";
        }
        else if (varName == "pi_zero_P"){
            sVariable   = "log(P(#pi^{0}))";
            sBranchName = {"pi_zero_resolved_P"};
            Bins = 10;
            Range[0] = 2980;  //not log!
            Range[1] = 100000;//not log!
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {8.0, 8.75367, 9.01297, 9.21333, 9.38583, 9.55554, 9.72447, 9.90965, 10.1241, 10.405, 11.5};
            vVarRange = {{Range[0],Range[1]}};
            cut = "log(pi_zero_resolved_P)";//TODO fix equidistant binning, now based on 2016 PHSP
        }
        else if (varName == "pi_zero_P_DTF"){
            sVariable   = "log(P(#pi^{0}^{DTF}})";
            sBranchName = {"pi_zero_resolved_P_DTF"};
            Bins = 10;
            Range[0] = 2980;  //not log!
            Range[1] = 100000;//not log!
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {8.0, 8.7508, 9.01136, 9.20829, 9.38185, 9.55001, 9.723, 9.90559, 10.1195, 10.4024, 11.5};
            vVarRange = {{Range[0],Range[1]}};
            cut = "log(pi_zero_resolved_P_DTF)";//TODO fix equidistant binning, now based on 2016 PHSP
        }
        else if (varName == "pi_zero_PT"){
            sVariable   = "log(P_{T}(#pi^{0}))";
            sBranchName = {"pi_zero_resolved_PT"};
            Bins = 10;
            Range[0] = 800; //not log!
            Range[1] = 8000;//not log!
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {6.5, 6.81709, 6.91294, 7.01023, 7.11076, 7.21995, 7.34191, 7.48141, 7.64965, 7.88324, 9.0};
            vVarRange = {{Range[0],Range[1]}};
            cut = "log(pi_zero_resolved_PT)"; //TODO fix equidistant binning, now based on 2016 PHSP
        }
        else if (varName == "pi_zero_PT_DTF"){
            sVariable   = "log(P_{T}(#pi^{0})^{DTF})";
            sBranchName = {"pi_zero_resolved_PT_DTF"};
            Bins = 10;
            Range[0] = 800; //not log!
            Range[1] = 8000;//not log!
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {6.5, 6.8152, 6.91109, 7.00766, 7.10771, 7.21609, 7.33793, 7.47641, 7.64706, 7.88087, 9.0};
            vVarRange = {{Range[0],Range[1]}};
            cut = "log(pi_zero_resolved_PT_DTF)";//TODO fix equidistant binning, now based on 2016 PHSP
        }
        else if (varName == "K_plus_PT"){
            sVariable   = "log(P_{T} (K^{+})) [MeV^{2}]";
            sBranchName = {"K_plus_PT"};
            Bins = 10;
            Range[0] = 270; //not log!
            Range[1] = 16000;//not log!
            Step = (Range[1] - Range[0]) / Bins;
            fillBinEdges();
            binEdges = {5.6, 6.658, 6.91969, 7.12516, 7.30566, 7.47015, 7.64604, 7.8327, 8.04823, 8.35614, 9.7};
            vVarRange = {{Range[0],Range[1]}};
            cut =  "log(K_plus_PT)";
        }
        else{
            sVariable   = "";
            sBranchName = {};
            Bins = 1;
            Range[0] = 0.;
            Range[1] = 0.;
            vVarRange = {};
            fillBinEdges();
            binEdges={Range[0],Range[1]};
            Step = 0;
            cut = "";
            isEquidistant = false;
        }
        if (!check_vector_size(isEquidistant)){
            sVariable   = "";
            sBranchName = {};
            Bins = 0;
            Range[0] = 0.;
            Range[1] = 0.;
            vVarRange = {};
            fillBinEdges();
            Step = 0;
            cut = "";
        }

    }
    ~TMefficiencyClass();   //destuctor
};

TMefficiencyClass::~TMefficiencyClass(){}//destuctor

bool useExtraVarBool(string extraVar);

string GetsWeightPlots(string year, bool UseOnlyJpsiEvents, bool UseOnlyMuMuEvents, bool KshortDecaysInVelo, bool GetShapeFromMC, string SignalType, string BkgType, bool ConstrainParameters);

double massFit(string year, string magnet, int Run,
               bool MC, bool Preselected, bool TM, bool PHSP,       //input/output file selection
               bool UseOnlyJpsiEvents, bool UseOnlyMuMuEvents,      //signal/reference
               bool GetShapeFromMC, string SigType, string BkgType, bool ConstrainParameters, //shape
               bool KshortDecaysInVelo, bool UseLowQ2Range,         //LL/DD? q2range?
               Double_t TMVAcut,  int randomSubset,                 //TMVA options
               bool fixedMassRegion, bool yieldOverFullRange,       //yield calculation region
               bool sWeight,                                        //sWeight data?
               bool loopFit, bool IsEfficiency,                     //additional options
               string sExtraVar, int nExtraBin,                     //fit in bins of extra variable
               bool removeMultiple,                                 //Remove multiple candidates?
               bool weighted, bool weightedFromPi0, string whichWeight, //use weight in the fit?
               bool nonTM, string customTMbranch, bool gammaTM,     //TM options
               bool InclusiveSample);

int quickFit(string year, bool MC, bool sWeight, bool UseOnlyJpsiEvents, bool UseOnlyMuMuEvents, bool KshortDecaysInVelo, bool GetShapeFromMC, string SigType, string BkgType, bool ConstrainParameters);

int efficiencyFit(std::string year, string magnet, int Run ,
                  bool Preselected, bool TM, bool PHSP,         //input/output file selection
                  bool UseOnlyJpsiEvents, bool UseOnlyMuMuEvents,  //signal/reference
                  bool GetShapeFromMC, std::string SigType , std::string BckGndType , bool ConstrainParameters, //shape
                  bool KshortDecaysInVelo,           //LL/DD?
                  Double_t TMVAcut,                          //TMVA options
                  bool fixedMassRegion,             //yield calculation region
                  bool UseLowQ2Range,     //q2 ranges
                  string sExtraVar, int nExtraBin, bool removeMultiple,
                  bool weighted, bool weightedFromPi0, string whichWeight, string customTMbranch, bool gammaTM);


int basicYieldFit(std::string year, int Run ,
                  bool MC, bool PHSP,        //input/output file selection
                  bool UseOnlyJpsiEvents, bool UseOnlyMuMuEvents,  //signal/reference
                  bool GetShapeFromMC, std::string SigType , std::string BckGndType , bool ConstrainParameters, //shape
                  bool KshortDecaysInVelo, bool UseLowQ2Range ,  //LL/DD? q2range?
                  Double_t TMVAcut,                                       //TMVA options
                  bool fixedMassRegion, bool loopFit, bool removeMultiple);

int basicYieldFitAllYears(bool MC, bool PHSP,        //input/output file selection
                  bool UseOnlyJpsiEvents, bool UseOnlyMuMuEvents,  //signal/reference
                  bool GetShapeFromMC, std::string SigType , std::string BckGndType , bool ConstrainParameters, //shape
                  bool KshortDecaysInVelo, bool UseLowQ2Range ,  //LL/DD? q2range?
                  Double_t TMVAcut,                                       //TMVA options
                  bool fixedMassRegion, bool loopFit, bool removeMultiple);


int basicFitAllYearsAndRegions(bool MC, bool PHSP,        //input/output file selection
        bool GetShapeFromMC, std::string SigType , std::string BckGndType , bool ConstrainParameters, //shape
        bool KshortDecaysInVelo, bool UseLowQ2Range ,  //LL/DD? q2range?
        Double_t TMVAcut, bool removeMultiple  //TMVA options, remove multiple
        );


int basicYieldFitAllRuns(
                  bool MC,        //input/output file selection
                  bool UseOnlyJpsiEvents, bool UseOnlyMuMuEvents,  //signal/reference
                  bool GetShapeFromMC, std::string SigType , std::string BckGndType , bool ConstrainParameters, //shape
                  bool KshortDecaysInVelo, bool UseLowQ2Range ,  //LL/DD? q2range?
                  Double_t TMVAcut,                                       //TMVA options
                  bool fixedMassRegion, bool loopFit,bool removeMultiple//yield calculation region
                  );


//Print efficiencies and fit parameters
int PrintFitResults(RooFitResult* fitRes);

int fitJpsi(string year, bool MC, double TMVAcut, bool RemoveMultiple);

#endif // MASSFIT_HPP