//create MC weights for B+->Kst+mumu
//signal B+ mass is fitted to generate sPlots
//david gerick
//Renata Kopecna

#include "GlobalFunctions.hh"
#include "MassFit.cpp"
#include "Paths.hpp"
#include "Utils.hpp"
#include "BmassShape/ParamValues.hpp"

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

//////////////////////////////////////////////////////
///  quickFit()
//////////////////////////////////////////////////////
/// Function is imported from MassFit.hpp
///
/// the function merges up and down events into a new file.
/// from the fit to the B+ mass distribution, a sPlot weight is saved to a new Branch in this merged file.
/// 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 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
///
/// OPTIONAL:
/// To determine the signal yield of the reference/resonance channel, the boolean UseOnlyJpsiEvents can be set.
/// Then only B+ -> J/psi K*+ events are fitted for S and B. The Q^2 bin #9 is taken, with 8.0 GeV^2 < Q^2 < 11.0 GeV^2
/// Also the seperation into DD and LL tracks for the Kshort Channel is possible!
///
///
///
//////////////////////////////////////////////////////
///  quickFitAll()
//////////////////////////////////////////////////////
///
/// Function to fit all years, sWeights all data years
/// Possible for data or MC. Most options of quickFit() are identically choosable
///
///
///
//////////////////////////////////////////////////////
///  NtrackWeight()
//////////////////////////////////////////////////////
///
/// the sPlot results from the quickFit() function are used to weight the nTrack signal distribution.
/// the ratio of MC / data of this nTrack distribution is used to create 1D weights for the MC sample,
/// which is later used as input for the BDT selection///
///
/// MC can be reweighted also without sWeights!
///
/// NEWLY ADDED: the sPlot results from the quickFit() function are used to 2 * 1D-weight the MC sample.
/// It will be weighted in nTracks first and then in the transverse momentum of the B+ (B_plus_PT)
///
/// Weight branches can be re-defined in GlobalFunctions.hh (see firstMCweight and seconMCweight)
///
///
///
//////////////////////////////////////////////////////
///  WeightAll()
//////////////////////////////////////////////////////
///
/// run the quickFit for both years, then apply NtrackWeights to both years.
///
///
///
//////////////////////////////////////////////////////
///  ReweightMCOnly()
//////////////////////////////////////////////////////
///
/// Do not fit the data but use the sWeighted data to re-weight the MC.
/// Including options to re-weight Reference channel MC and PHSP MC.
///
///
//////////////////////////////////////////////////////
///  GetSignalAndBckgndEstimation()
///
/// Uses the quickFit() function with the optional boolean to only fit the Jpsi events to get the estimation of
/// signal candidates for the reference channel before pre-selection. Via the PDG Branching Ratios, the signal
/// yield of the signal channel is calculated. From the total number of candidates for the signal channel within
/// the B-plus mass window estimated as 2*sigma window
/// the background is determined via B = N - S, with S the calculated signal estimation from
/// first part
///
/// (not very suitable for pi0)
///

//Initilize MC tracks histograms
TH1D *get_hist_MC_w(TH1D *hist, string name, bool is2D){
    TH1D* hist_nTracks_MC_w = (TH1D*) hist->Clone((is2D ? seconMCweight : firstMCweight+name).c_str());
    if (is2D) hist_nTracks_MC_w->SetTitle((seconMCweight+" MC 2D weights [norm.]").c_str());
        else hist_nTracks_MC_w->SetTitle((firstMCweight+" MC weights [norm.]").c_str());
    return hist_nTracks_MC_w;
}

//Initialize
TH1D *get_hist(string name, bool is2D, bool isMC){
    TH1D* hist = NULL;
    if (!isMC){
        if (is2D) hist = new TH1D(seconMCweight.c_str(), (seconMCweight+" Yield [norm.]").c_str(), secondnBins, seconMCrange[0], seconMCrange[1]);
        else      hist = new TH1D(firstMCweight.c_str() ,(firstMCweight+" Yield [norm.]").c_str(), firstnBins, firstMCrange[0], firstMCrange[1]);
    }
    else{
        if (is2D) hist = new TH1D((seconMCweight+name).c_str(),(seconMCweight+" MC weighted by "+firstMCweight+" [norm.]").c_str(), secondnBins, seconMCrange[0], seconMCrange[1]);
        else      hist = new TH1D((firstMCweight+name).c_str(),(firstMCweight+" MC [norm.]").c_str(), firstnBins, firstMCrange[0], firstMCrange[1]);
    }
    coutDebug("Hist name: " + string(hist->GetName()));
    return hist;
}


//Get MC weights
void getWeightHist(TChain *tree, string weightName, TH1D *hist, TH1D *hist_w, string TMmethod, bool KshortDecayInVelo, bool gammaTM){

    tree->SetBranchStatus("*",0);
    //activate needed branches
    tree->SetBranchStatus(TMmethod.c_str(),1);
    tree->SetBranchStatus(weightName.c_str(),1);
    if(Kst2Kspiplus)tree->SetBranchStatus("KshortDecayInVeLo",1);

    //assign variables
    Double_t weight = 0;
    Int_t weight_I = 0;
    bool isInt = false;
    if (weightName.find("Track") != std::string::npos) isInt = true; //check if the value is integer or not
    Int_t KshortDecayInVeLo = 0;
    Int_t TMed = 0;
    Int_t TMed_gammas = 0;

    tree->SetBranchAddress(TMmethod.c_str(), &TMed);
    tree->SetBranchAddress("TM_gammas",&TMed_gammas);
    if (isInt) tree->SetBranchAddress(weightName.c_str() , &weight_I);
    else tree->SetBranchAddress(weightName.c_str() , &weight);

    if(Kst2Kspiplus)tree->SetBranchAddress("KshortDecayInVeLo" , &KshortDecayInVeLo);

    int nEvts = tree->GetEntries();
    for(int i=0; i < nEvts ; i++){
        if (0ul == (i % 10000ul) || nEvts == i + 1) coutDebug("Read MC event " + to_string(i) + "/" + to_string(nEvts));
        tree->GetEntry(i);
        if (!isTM(TMmethod,TMed,gammaTM,TMed_gammas)) continue;

        //only write the correct DD or LL tracks into the histogram!
        if(Kst2Kspiplus && SplitDDandLL){
            if (KshortDecayInVelo ^ KshortDecayInVeLo) continue; //^ is a xor operator
        }
        isInt ? hist->Fill(weight_I) : hist->Fill(TMath::Log(weight));
    }

    //normalize the histogram by its integral
    hist->Sumw2();
    hist->Scale(1./hist->Integral());

    //divide histograms (hist_w = clone of hist_data)
    hist_w->Divide(hist);

    //reset branch links because root is stupid
    tree->ResetBranchAddresses();

    return;

}

//Find the weight for each event
void fillWeightHist(double &w, double &delta_w, int bin, bool TM, TH1D *h_w){
    if (TM){
        w       = h_w->GetBinContent(bin);
        delta_w = h_w->GetBinError(bin);
    }
    else{
        w       =  1.0;
        delta_w = -1.0;
    }
    return;
}

double getDelta_w2D(double w, double delta_w, double w2D, double delta_w2D){
    if(delta_w < 0.0 || delta_w2D < 0.0)//keep delta_w2D negative, if both deltas are negative
        return (-TMath::Sqrt(TMath::Power(delta_w*w2D,2)+TMath::Power(delta_w2D*w,2)));
    else
        return (TMath::Sqrt(TMath::Power(delta_w*w2D,2)+TMath::Power(delta_w2D*w,2)));
}

int quickFitAll(string SignalShape = "OneCB", string BckShape = "SingleExponential", bool sWeight=true, bool GetShapeFromMC = true,
                bool ConstrainParameters = true, int Run = 1){
    //Creates sWeights in data

    bool UseOnlyJpsiEvents = true; //Set based on if reweighted by Jpsi or not
    bool UseOnlyMuMuEvents = false;
    std::vector<string> years = yearsMC(false,UseOnlyJpsiEvents,Run);

    for(unsigned int y = 0; y < years.size(); y++){
        if(Kst2Kspiplus && SplitDDandLL){
           if(quickFit(years.at(y), false, sWeight, UseOnlyJpsiEvents, UseOnlyMuMuEvents, true, GetShapeFromMC, SignalShape, BckShape, ConstrainParameters) == 0){
               coutERROR("Failed quickFit() for " + years.at(y) + " (LL tracks). Exit!");
               return 0;
           }
           if(quickFit(years.at(y), false, sWeight, UseOnlyJpsiEvents, UseOnlyMuMuEvents, false, GetShapeFromMC, SignalShape, BckShape, ConstrainParameters) == 0){
               coutERROR("Failed quickFit() for " + years.at(y) + " (DD tracks). Exit!");
               return 0;
           }
        }
        else{
            if(quickFit(years.at(y), false, sWeight, UseOnlyJpsiEvents, UseOnlyMuMuEvents, false, GetShapeFromMC, SignalShape, BckShape, ConstrainParameters) == 0){
               coutERROR("Failed quickFit() for " + years.at(y) + ". Exit!");
               return 0;
           }
        }
    }
    return 1;
}

int NtrackWeight(string year = "2011", const bool ReferenceChannel = false, bool PHSP = false, bool ReweightInBplusPT = true, bool KshortDecayInVelo = true, bool sWeightUse = true) {
    //the bollean is called sWeightUse in order not to be confused with sWeight branches

    coutInfo("Start MonteCarlo re-weighting: " + year + string(Kst2Kspiplus && SplitDDandLL ? (KshortDecayInVelo ? " (LL tracks)" : " (DD tracks)") : ""));

    if (!checkMC(ReferenceChannel,PHSP)) return false;
    //Make ROOT shutup
    gStyle -> SetOptStat(0);
    LHCbStyle();
    gROOT->SetBatch(kTRUE);

    //get data and MC file(s)
    TChain * tree     = new TChain("DecayTree");
    TChain * treeMCTM = new TChain("DecayTreeTruthMatched");
    TFile * output;

    bool GetWeightsFromJpsiChannel = ReweightByRefChannel;
    string MCyear = year;
    if(KshortChannel){
      if(ReweightByRefChannel && !checkRefYear(year))
        GetWeightsFromJpsiChannel = false;
    }
    else{ //K+ pi0 channel
      if (GetWeightsFromJpsiChannel && !checkRefYear(year)) MCyear = "2016"; //Use 2016 MC for reweighting 2017 and 2018 signal MC
    }
    //Decide if reweighting by Reference channel or not
    bool loadRefChannel = (AlwaysUseRefChannelData || GetWeightsFromJpsiChannel || ReferenceChannel);

    ///DATA
    //load sWeighted OnlyJpsi data OR load the full data-set of Jpsi and non-resonant mumu data events (sWeighted)
    tree->Add(GetBDTinputFile(year,false,loadRefChannel,false,KshortDecayInVelo).c_str());
    coutDebug("Adding " + GetBDTinputFile(year,false,loadRefChannel,false,KshortDecayInVelo));

    ///MC
    //load MonteCarlo sample of up and down type. truthmatched and non-truthmatched
    treeMCTM->Add(GetInputFile(MCyear,"down",true, true, loadRefChannel, false, false).c_str());
    treeMCTM->Add(GetInputFile(MCyear,"up",  true, true, loadRefChannel, false, false).c_str());

    coutDebug("Adding " + GetInputFile(MCyear,"down",true, true, loadRefChannel, false, false));
    coutDebug("Adding " + GetInputFile(MCyear,"up",  true, true, loadRefChannel, false, false));
    checkEntries(tree);
    checkEntries(treeMCTM);

    //deactivate all other branches
    tree->SetBranchStatus("*",0);
    tree->SetBranchStatus(firstMCweight.c_str(),1);
    tree->SetBranchStatus(seconMCweight.c_str(),1);
    if (sWeightUse) tree->SetBranchStatus("N_Bplus_sw",1);
    //assign variables
    Int_t nTracks = 0;
    Double_t sWeights = 1.0;
    Double_t B_plus_PT = 0.0;

    //link variables to branches
    tree->SetBranchAddress(firstMCweight.c_str() , &nTracks);
    tree->SetBranchAddress(seconMCweight.c_str() , &B_plus_PT);

    if (sWeightUse) tree->SetBranchAddress("N_Bplus_sw",&sWeights);

    //Make histograms
    TH1::SetDefaultSumw2();
    TH2::SetDefaultSumw2();

    TH1D* hist_nTracks = get_hist("",false,false);
    TH1D* hist_nTracks_MC             = get_hist("_MC",false,true);
    TH1D* hist_nTracks_MC_TM          = get_hist("_MC_TM",false,true);
    TH1D* hist_nTracks_MC_TM_rndGamma = get_hist("_MC_TM_rndGamma",false,true);
    TH1D* hist_nTracks_MC_noPi0TM     = get_hist("_MC_noPi0TM",false,true);

    TH1D* hist_BplusPT = get_hist("",true,false);
    TH1D* hist_BplusPT_MC             = get_hist("_weighted",true,true);
    TH1D* hist_BplusPT_MC_TM          = get_hist("_weighted_TM",true,true);
    TH1D* hist_BplusPT_MC_TM_rndGamma = get_hist("_weighted_TM_rndGamma",true,true);
    TH1D* hist_BplusPT_MC_noPi0TM     = get_hist("_weighted_noPi0TM",true,true);

    //histogram to check correlations between nTracks and B_plus_PT
    TH2D* hCorrelationCheck =  new TH2D((firstMCweight+"_"+seconMCweight+"_correlation").c_str(),("correlation between "+firstMCweight+" and "+seconMCweight).c_str(),
                                         firstnBins,firstMCrange[0], firstMCrange[1], secondnBins, seconMCrange[0], seconMCrange[1]);

    ///-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/
    /// load data (nTracks and B_plus_PT) into histograms
    ///-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/-/

    Int_t nEvents = tree->GetEntries();
    Int_t nEventsMCTM = treeMCTM->GetEntries();

    //loop over data events
    coutInfo("Loop over data sample from " + year + TheDecay + " to fill histogram!");
    if (sWeightUse){
        for(int i=0; i < nEvents; i++){
            if (0ul == (i % 10000ul) || nEvents == i + 1) coutDebug("Read data event " + to_string(i) + "/" + to_string(nEvents));
            tree->GetEntry(i);
            hist_nTracks->Fill(nTracks, sWeights);
            hist_BplusPT->Fill(TMath::Log(B_plus_PT), sWeights);
            hCorrelationCheck->Fill(nTracks, TMath::Log(B_plus_PT), sWeights);
        }
    }
    else {
        for(int i=0; i < nEvents; i++){
            if (0ul == (i % 10000ul) || nEvents == i + 1) coutDebug("Read data event " + to_string(i) + "/" + to_string(nEvents));
            tree->GetEntry(i);
            hist_nTracks->Fill(nTracks);
            hist_BplusPT->Fill(TMath::Log(B_plus_PT));
            hCorrelationCheck->Fill(nTracks, TMath::Log(B_plus_PT));
        }
    }

    Double_t CorrelationCoefficent = hCorrelationCheck->GetCorrelationFactor();
    coutDebug("The correlation coefficient between " + firstMCweight + " and " + seconMCweight + " is: " + to_string(CorrelationCoefficent));

    //normalize histograms by its integral
    hist_nTracks->Scale(1./hist_nTracks->Integral());
    hist_BplusPT->Scale(1./hist_BplusPT->Integral());

    //clone normalized histogram for determining the ratio between data/MC of nTracks string
    //Make three histograms, each for given TM method
    TH1D* hist_nTracks_MC_w             = get_hist_MC_w(hist_nTracks,"_MC_weights", false);
    TH1D* hist_nTracks_MC_w_TM          = get_hist_MC_w(hist_nTracks,"_MC_weights_TM", false);
    TH1D* hist_nTracks_MC_w_TM_rndGamma = get_hist_MC_w(hist_nTracks,"_MC_weights_TM_rndGamma", false);
    TH1D* hist_nTracks_MC_w_noPi0TM     = get_hist_MC_w(hist_nTracks,"_MC_weights_noPi0TM", false);
    //clone normalized histogram for determining the ratio between data/MC
    TH1D* hist_BplusPT_MC_w2D             = get_hist_MC_w(hist_BplusPT,"_MC_weights", true);
    TH1D* hist_BplusPT_MC_w2D_TM          = get_hist_MC_w(hist_BplusPT,"_MC_weights_TM", true);
    TH1D* hist_BplusPT_MC_w2D_TM_rndGamma = get_hist_MC_w(hist_BplusPT,"_MC_weights_TM_rndGamma", true);
    TH1D* hist_BplusPT_MC_w2D_noPi0TM     = get_hist_MC_w(hist_BplusPT,"_MC_weights_noPi0TM", true);

    ///////////////////////////////
    /// reweight in nTracks
    ///////////////////////////////

    treeMCTM->GetEntry(0);
    coutInfo("Loop over MC sample from " + year + TheDecay + " to fill histogram for 1D reweighting in " + firstMCweight + "!");
    getWeightHist(treeMCTM,firstMCweight,hist_nTracks_MC,hist_nTracks_MC_w,"TMedBKGCAT",KshortDecayInVelo, false);
    getWeightHist(treeMCTM,firstMCweight,hist_nTracks_MC_TM,hist_nTracks_MC_w_TM,"TMed",KshortDecayInVelo, true);
    if (!Kst2Kspiplus){
        getWeightHist(treeMCTM,firstMCweight,hist_nTracks_MC_TM_rndGamma,hist_nTracks_MC_w_TM_rndGamma,"TMed",KshortDecayInVelo, false);
        getWeightHist(treeMCTM,firstMCweight,hist_nTracks_MC_noPi0TM,hist_nTracks_MC_w_noPi0TM,"TMed_noPi0",KshortDecayInVelo, true);
    }

    ///////////////////////////////
    /// reweight in B_plus_PT
    ///////////////////////////////

    if(ReweightInBplusPT){
        coutInfo("Loop over MC sample from " + year + TheDecay + " to fill histogram for 1D reweighting in " + seconMCweight + "!");
        getWeightHist(treeMCTM,seconMCweight,hist_BplusPT_MC,hist_BplusPT_MC_w2D,"TMedBKGCAT",KshortDecayInVelo, false);
        getWeightHist(treeMCTM,seconMCweight,hist_BplusPT_MC_TM,hist_BplusPT_MC_w2D_TM,"TMed",KshortDecayInVelo, true);
        if (!Kst2Kspiplus){
            getWeightHist(treeMCTM,seconMCweight,hist_BplusPT_MC_TM_rndGamma,hist_BplusPT_MC_w2D_TM_rndGamma,"TMed",KshortDecayInVelo, false);
            getWeightHist(treeMCTM,seconMCweight,hist_BplusPT_MC_noPi0TM,hist_BplusPT_MC_w2D_noPi0TM,"TMed_noPi0",KshortDecayInVelo, true);
        }
    }

    ////////////////////////////////////////////////////////////////////
    /// Plot the histograms to canvases
    ////////////////////////////////////////////////////////////////////

    string canvasPath = "";
    if(!PHSP && !ReferenceChannel){
        canvasPath = GetControlPlots(year,ReferenceChannel,false,KshortDecayInVelo,sWeightUse,1).c_str();
        drawKolmogorovTest(hist_nTracks, hist_nTracks_MC, canvasPath,"_BKGCAT");
        drawKolmogorovTest(hist_nTracks, hist_nTracks_MC_TM, canvasPath,"_TM");
        drawKolmogorovTest(hist_nTracks, hist_nTracks_MC_TM_rndGamma, canvasPath,"_TM_rndGamma");
        //Plot correlation between B_plus and nTracks
        canvasPath = GetControlPlots(year,ReferenceChannel,PHSP,KshortDecayInVelo,sWeightUse,4).c_str();
        drawWeightCorrelation(hCorrelationCheck, CorrelationCoefficent, canvasPath);
    }

    //Plot the ratio
    canvasPath = GetControlPlots(year,ReferenceChannel,PHSP,KshortDecayInVelo,sWeightUse,2).c_str();
    drawWeightRatio(hist_nTracks_MC_w, canvasPath, false, "_BKGCAT");
    drawWeightRatio(hist_nTracks_MC_w_TM, canvasPath, false, "_TM");
    drawWeightRatio(hist_nTracks_MC_w_TM_rndGamma, canvasPath, false,"_TM_rndGamma");

    //Plot the ratio for B_plus_PT
    canvasPath = GetControlPlots(year,ReferenceChannel,PHSP,KshortDecayInVelo,sWeightUse,3).c_str();
    drawWeightRatio(hist_BplusPT_MC_w2D, canvasPath, true, "_BKGCAT");
    drawWeightRatio(hist_BplusPT_MC_w2D_TM, canvasPath, true, "_TM");
    drawWeightRatio(hist_BplusPT_MC_w2D_TM_rndGamma, canvasPath, true,"_TM_rndGamma");

    //////////////////////////////////
    /// SAVE WEIGHTS TO MC
    //////////////////////////////////

    treeMCTM->ResetBranchAddresses();
    treeMCTM->SetBranchStatus("*",1);
    coutInfo("Copy MC Trees... ");
    if (!ReferenceChannel && !PHSP && year=="2015") return 1; //Skip 2015 signal MC
    TTree * newtreeMCTM;
    TChain * reweightedTreeMCTM;
    //if ReweightByRefChannel == true, the ratios above have been created with Jpsi data and MC
    //but if we don't run for the ReferenceChannel, we want to load the correct Tree (so delete previous treeMCTM and treeMC and load new (Signal MC or PHSP)

    if(GetWeightsFromJpsiChannel && !ReferenceChannel){

        //to be safe, close trees of reference channel MC:
        delete treeMCTM;
        reweightedTreeMCTM = new TChain("DecayTreeTruthMatched");
        if (reweightedTreeMCTM == NULL) coutERROR("Failed to create your new tree!");

        //load trees from Signal MC files:
        reweightedTreeMCTM  ->Add(GetInputFile(year,"down",true,true,ReferenceChannel,PHSP,false).c_str());
        reweightedTreeMCTM  ->Add(GetInputFile(year,"up",  true,true,ReferenceChannel,PHSP,false).c_str());

        coutDebug("Loading " + GetInputFile(year,"down",true,true,ReferenceChannel,PHSP,false) );
        coutDebug("Loading " + GetInputFile(year,"up",  true,true,ReferenceChannel,PHSP,false) );

        //remove an old BDTresponse branch, that might still be active in the tree after pre-selection:
        TBranch * bTM = (TBranch *)reweightedTreeMCTM->GetBranch(TMVAmethod+"response");
        if(bTM != NULL) reweightedTreeMCTM->SetBranchStatus("*response", 0);

        //Get new event numbers for the Signal or PHSP trees:
        nEventsMCTM = reweightedTreeMCTM->GetEntries();

        //copy the loaded tree according to the cuts
        if(Kst2Kspiplus && SplitDDandLL){
            coutDebug("Apply track cut on MC Tree: " + string(KshortDecayInVelo ? "LL tracks!" : "DD tracks!"));
            newtreeMCTM = reweightedTreeMCTM->CopyTree(Form("KshortDecayInVeLo == %i", KshortDecayInVelo ? 1 : 0));
        }
        else{
            newtreeMCTM = reweightedTreeMCTM->CopyTree("");
        }
        coutDebug("Truthmatched Tree completed!");

    }

    else{ //if not forced to be reweighted from Jpsi channel, just copy the tree used above for the ratio:

      //remove an old BDTresponse branch, that might still be active in the tree after pre-selection:
      TBranch * bTM = (TBranch *)treeMCTM->GetBranch(TMVAmethod+"response");
      if(bTM != NULL){
        treeMCTM->SetBranchStatus("*response", 0);
      }

      if(Kst2Kspiplus && SplitDDandLL){
          coutDebug("Apply track cut on MC Tree: " +string(KshortDecayInVelo ? "LL tracks!" : "DD tracks!"));
          newtreeMCTM = treeMCTM->CopyTree(Form("KshortDecayInVeLo == %i", KshortDecayInVelo ? 1 : 0));      
      }
      else{
          coutDebug("Copying tree...");
          newtreeMCTM = treeMCTM->CloneTree();
      }
      coutDebug("Truthmatched Tree completed!");
    }

    coutDebug("Old TM Tree entries:\t" + to_string(nEventsMCTM));
    coutDebug("Copied TM Tree entries:\t" + to_string(newtreeMCTM->GetEntries()));
    coutInfo("Adding branches to MC tree.");

    //create new TBranches: each for different TM method and coresponding weight error branch
    double w = 1., w2D = 1., delta_w = 0., delta_w2D = 0.;
    double w_TM = 1., w2D_TM  = 1., delta_w_TM  = 0., delta_w2D_TM  = 0.;
    double w_noPi0TM = 1., w2D_noPi0TM = 1., delta_w_noPi0TM = 0., delta_w2D_noPi0TM = 0.;
    double w_TM_rndGamma = 1., w2D_TM_rndGamma  = 1., delta_w_TM_rndGamma  = 0., delta_w2D_TM_rndGamma  = 0.;

    TBranch* Bra_w        = newtreeMCTM->Branch(Form("weight_%s",        firstMCweight.c_str()),&w,  Form("weight_%s/D",  firstMCweight.c_str()));
    TBranch* Bra_w2D      = newtreeMCTM->Branch(Form("weight2D_%s",      firstMCweight.c_str()),&w2D,Form("weight2D_%s/D",firstMCweight.c_str()));
    TBranch* Bra_delta_w  = newtreeMCTM->Branch(Form("delta_weight_%s",  firstMCweight.c_str()),&delta_w,  Form("delta_weight_%s/D",  firstMCweight.c_str()));
    TBranch* Bra_delta_w2D= newtreeMCTM->Branch(Form("delta_weight2D_%s",firstMCweight.c_str()),&delta_w2D,Form("delta_weight2D_%s/D", firstMCweight.c_str()));

    TBranch* Bra_w_TM        = newtreeMCTM->Branch(Form("weight_%s_TM",        firstMCweight.c_str()),&w_TM ,  Form("weight_%s_TM/D",  firstMCweight.c_str()));
    TBranch* Bra_w2D_TM      = newtreeMCTM->Branch(Form("weight2D_%s_TM",      firstMCweight.c_str()),&w2D_TM ,Form("weight2D_%s_TM/D",firstMCweight.c_str()));
    TBranch* Bra_delta_w_TM  = newtreeMCTM->Branch(Form("delta_weight_%s_TM",  firstMCweight.c_str()),&delta_w_TM ,  Form("delta_weight_%s_TM/D",  firstMCweight.c_str()));
    TBranch* Bra_delta_w2D_TM= newtreeMCTM->Branch(Form("delta_weight2D_%s_TM",firstMCweight.c_str()),&delta_w2D_TM ,Form("delta_weight2D_%s_TM/D",firstMCweight.c_str()));

    TBranch* Bra_w_noPi0TM        = newtreeMCTM->Branch(Form("weight_%s_noPi0TM",         firstMCweight.c_str()),&w_noPi0TM,  Form("weight_%s_noPi0TM/D",  firstMCweight.c_str()));
    TBranch* Bra_w2D_noPi0TM      = newtreeMCTM->Branch(Form("weight2D_%s_noPi0TM",       firstMCweight.c_str()),&w2D_noPi0TM,Form("weight2D_%s_noPi0TM/D",firstMCweight.c_str()));
    TBranch* Bra_delta_w_noPi0TM  = newtreeMCTM->Branch(Form("delta_weight_%s_noPi0TM",   firstMCweight.c_str()),&delta_w_noPi0TM,  Form("delta_weight_%s_noPi0TM/D",  firstMCweight.c_str()));
    TBranch* Bra_delta_w2D_noPi0TM= newtreeMCTM->Branch(Form("delta_weight2D_%s_noPi0TM", firstMCweight.c_str()),&delta_w2D_noPi0TM,Form("delta_weight2D_%s_noPi0TM/D",firstMCweight.c_str()));

    TBranch* Bra_w_TM_rndGamma        = newtreeMCTM->Branch(Form("weight_%s_TM_rndGamma",        firstMCweight.c_str()),&w_TM_rndGamma ,  Form("weight_%s_TM_rndGamma/D",  firstMCweight.c_str()));
    TBranch* Bra_w2D_TM_rndGamma      = newtreeMCTM->Branch(Form("weight2D_%s_TM_rndGamma",      firstMCweight.c_str()),&w2D_TM_rndGamma ,Form("weight2D_%s_TM_rndGamma/D",firstMCweight.c_str()));
    TBranch* Bra_delta_w_TM_rndGamma  = newtreeMCTM->Branch(Form("delta_weight_%s_TM_rndGamma",  firstMCweight.c_str()),&delta_w_TM_rndGamma ,  Form("delta_weight_%s_TM_rndGamma/D",  firstMCweight.c_str()));
    TBranch* Bra_delta_w2D_TM_rndGamma= newtreeMCTM->Branch(Form("delta_weight2D_%s_TM_rndGamma",firstMCweight.c_str()),&delta_w2D_TM_rndGamma ,Form("delta_weight2D_%s_TM_rndGamma/D",firstMCweight.c_str()));


    nEventsMCTM = newtreeMCTM->GetEntries();
    coutInfo("Looping over MC tree with " + to_string(nEventsMCTM) + " entries.");

    //renew links to branches in new tree:
    //status is already set to 1 for all
    Int_t TMedBKGCAT, TMed, TMed_noPi0;
    Int_t gammaTMed;
    Int_t nTracksMC;
    Double_t B_plus_PTMC;

    newtreeMCTM->SetBranchAddress("TMedBKGCAT" , &TMedBKGCAT);
    newtreeMCTM->SetBranchAddress("TMed" , &TMed);
    if (!Kst2Kspiplus){
        newtreeMCTM->SetBranchAddress("TMed_noPi0", &TMed_noPi0);
        newtreeMCTM->SetBranchAddress(gammaTMbranch.c_str() , &gammaTMed);
    }
    newtreeMCTM->SetBranchAddress((firstMCweight).c_str() , &nTracksMC);
    newtreeMCTM->SetBranchAddress((seconMCweight).c_str() , &B_plus_PTMC);

    Int_t nTrackBin = 1, nBplusBin = 1;

    //loop over MC events (Only the TruthMatched events get weights!)
    coutInfo("Loop over MC data from " + year + TheDecay + " to get weights for MC");
    for(int i=0; i < nEventsMCTM; i++)
    {
        if (0ul == (i % 10000ul) || nEventsMCTM == i + 1) coutDebug("Read MC event " + to_string(i) + "/" + to_string(nEventsMCTM));
        newtreeMCTM->GetEntry(i);

        //1D weights
        if(nTracksMC < firstMCrange[0] || nTracksMC > firstMCrange[1]){
            coutInfo(firstMCweight + " out of range: " + to_string(nTracksMC) + " in event: " + to_string(i));
            w       =  1.0;
            delta_w = -1.0; //set negative for when no value could be obtained
            w_TM        =  1.0;
            delta_w_TM  = -1.0;
            w_TM_rndGamma        =  1.0;
            delta_w_TM_rndGamma  = -1.0;
            w_noPi0TM       =  1.0;
            delta_w_noPi0TM = -1.0;
        }
        else{
            nTrackBin = hist_nTracks->FindBin(nTracksMC);
            fillWeightHist(w,             delta_w,             nTrackBin, TMedBKGCAT,                         hist_nTracks_MC_w);
            fillWeightHist(w_TM,          delta_w_TM,          nTrackBin, isTM("TMed",TMed,true, gammaTMed),  hist_nTracks_MC_w_TM);
            fillWeightHist(w_TM_rndGamma, delta_w_TM_rndGamma, nTrackBin, isTM("TMed",TMed,false,gammaTMed),  hist_nTracks_MC_w_TM_rndGamma);
            fillWeightHist(w_noPi0TM,     delta_w_noPi0TM,     nTrackBin, isTM("TMed",TMed,true, gammaTMed),  hist_nTracks_MC_w_noPi0TM);
        }

        //2 x 1D weights
        B_plus_PTMC = TMath::Log(B_plus_PTMC);
        if(B_plus_PTMC < seconMCrange[0] || B_plus_PTMC > seconMCrange[1]){
            coutInfo(seconMCweight + " out of range: " + to_string(B_plus_PTMC) + " in event: " + to_string(i));
            w2D       =  1.0;
            delta_w2D = -1.0; //set negative for when no value could be obtained
            w2D_TM    =  1.0;
            delta_w2D_TM = -1.0;
            w2D_noPi0TM       =  1.0;
            delta_w2D_noPi0TM = -1.0;
        }
        else{
            nBplusBin = hist_BplusPT->FindBin(B_plus_PTMC);

            fillWeightHist(w2D,             delta_w2D,             nBplusBin, TMedBKGCAT,                        hist_BplusPT_MC_w2D);
            fillWeightHist(w2D_TM,          delta_w2D_TM,          nBplusBin, isTM("TMed",TMed,true,gammaTMed),  hist_BplusPT_MC_w2D_TM);
            fillWeightHist(w2D_TM_rndGamma, delta_w2D_TM_rndGamma, nBplusBin, isTM("TMed",TMed,false,gammaTMed), hist_BplusPT_MC_w2D_TM_rndGamma);
            fillWeightHist(w2D_noPi0TM,     delta_w2D_noPi0TM,     nBplusBin, isTM("TMed",TMed,true,gammaTMed),  hist_BplusPT_MC_w2D_noPi0TM);
        }

        //save uncertainties on weights BEFORE weights, as w2D gets overwritten later on:
        delta_w2D             = getDelta_w2D(w,             delta_w,             w2D,             delta_w2D);
        delta_w2D_TM          = getDelta_w2D(w_TM,          delta_w_TM,          w2D_TM,          delta_w2D_TM);
        delta_w2D_TM_rndGamma = getDelta_w2D(w_TM_rndGamma, delta_w_TM_rndGamma, w2D_TM_rndGamma, delta_w2D_TM_rndGamma);
        delta_w2D_noPi0TM     = getDelta_w2D(w_noPi0TM,     delta_w_noPi0TM,     w2D_noPi0TM,  delta_w2D_noPi0TM);

        //Fill the branches with uncertainties on weights
        Bra_delta_w->Fill();
        Bra_delta_w2D->Fill();
        Bra_delta_w_TM->Fill();
        Bra_delta_w2D_TM->Fill();
        Bra_delta_w_TM_rndGamma->Fill();
        Bra_delta_w2D_TM_rndGamma->Fill();
        Bra_delta_w_noPi0TM->Fill();
        Bra_delta_w2D_noPi0TM->Fill();

        //prevent weight <= 0!
        if(w<=0.)w=1.;
        if(w2D<=0.)w2D=1.;
        if(w_TM<=0.)w_TM=1.;
        if(w2D_TM<=0.)w2D_TM=1.;
        if(w_TM_rndGamma<=0.)w_TM_rndGamma=1.;
        if(w2D_TM_rndGamma<=0.)w2D_TM_rndGamma=1.;
        if(w_noPi0TM<=0.)w_noPi0TM=1.;
        if(w2D_noPi0TM<=0.)w2D_noPi0TM=1.;

        //Recalculate 2D weights
        w2D*=w;
        w2D_TM*=w_TM;
        w2D_TM_rndGamma*=w_TM_rndGamma;
        w2D_noPi0TM*=w_noPi0TM;

        //Fill weights
        Bra_w->Fill();
        Bra_w2D->Fill();
        Bra_w_TM->Fill();
        Bra_w2D_TM->Fill();
        Bra_w_TM_rndGamma->Fill();
        Bra_w2D_TM_rndGamma->Fill();
        Bra_w_noPi0TM->Fill();
        Bra_w2D_noPi0TM->Fill();

    }

    //create output file for re-weighted and combined trees (truthmatched and non-truthmatched)
    coutInfo("Save re-weighted tree to file: " + GetBDTinputFile(year,true,ReferenceChannel,PHSP,KshortDecayInVelo));
    output = new TFile(GetBDTinputFile(year,true,ReferenceChannel,PHSP,KshortDecayInVelo).c_str(),"RECREATE");

    if(!output->IsOpen()){
        coutERROR("Could not create output file. Abort!");
        return 0;
    }

    output->cd();
    coutInfo("Saving new tree of MC data " + year + TheDecay + " to file!");
    newtreeMCTM->Write("",TObject::kWriteDelete);

    hist_nTracks_MC_w->Clear();
    hist_nTracks_MC_w_TM->Clear();
    hist_nTracks_MC_w_TM_rndGamma->Clear();
    hist_nTracks_MC_w_noPi0TM->Clear();

    hist_nTracks->Clear();
    hist_nTracks_MC->Clear();
    hist_nTracks_MC_TM->Clear();
    hist_nTracks_MC_TM_rndGamma->Clear();
    hist_nTracks_MC_noPi0TM->Clear();

    hist_BplusPT->Clear();
    hist_BplusPT_MC->Clear();
    hist_BplusPT_MC_TM->Clear();
    hist_BplusPT_MC_TM_rndGamma->Clear();
    hist_BplusPT_MC_noPi0TM->Clear();

    hist_BplusPT_MC_w2D->Clear();
    hist_BplusPT_MC_w2D_TM->Clear();
    hist_BplusPT_MC_w2D_TM_rndGamma->Clear();
    hist_BplusPT_MC_w2D_noPi0TM->Clear();

    output->Close();
    gROOT->Reset();


    return 1;

}

int NtrackWeightAllKplus(const bool ReferenceChannel = false, bool PHSP = false, bool ReweightInBplusPT = true){
    vector <string> years = yearsVector(true,ReferenceChannel,PHSP,12);
    for (auto year: years)
        if (NtrackWeight(year, ReferenceChannel, PHSP, ReweightInBplusPT, false, true)==0) return 0;

    return 1;
}


int WeightYear(string year = "2011", bool WeightBplusPT = true, bool GetShapeFromMC = true, bool sWeight = true){

    bool ConstrainParameters = KshortChannel ? false : true;

    string SignalFunction = Kst2Kspiplus ? "OneCB" : "OneCB";
    string BackGroundFunction = Kst2Kspiplus ? "SingleExponential" : "SingleExponential";

    if((ReweightByRefChannel && (checkRefYear(year) || !KshortChannel)) || AlwaysUseRefChannelData){ //reweight data by RefChannel MC, for Kshort only Run1
        if(Kst2Kspiplus && SplitDDandLL){

            //fit Jpsi channel only (create sWeights)
            if(quickFit(year, false, sWeight, true, false, true, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " Jpsi data (LL tracks) failed!");
                return 0;
            }
            if(quickFit(year, false, sWeight, true, false, false, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " Jpsi data (DD tracks) failed!");
                return 0;
            }

            //fit all events (Jpsi + non-resonant, but no sWeights!)
            if(quickFit(year, false, sWeight, false, false, true, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " data (LL tracks) failed!");
                return 0;
            }
            if(quickFit(year, false, sWeight, false, false, false, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " data (DD tracks) failed!");
                return 0;
            }

            //Apply sWeights from Jpsi channel to MC
            if(NtrackWeight(year, false, false, WeightBplusPT, true, sWeight) == 0){
                coutERROR("Creating weights for nTrack distribution for " + year + " MC (LL tracks) failed!");
                return 0;
            }
            if(NtrackWeight(year, false, false, WeightBplusPT, false, sWeight) == 0){
                coutERROR("Creating weights for nTrack distribution for " + year + " MC (DD tracks) failed!");
                return 0;
            }
        }
        else{
            //fit Jpsi channel only (create sWeights) for pi0 or don't split LL/DD
            if(quickFit(year, false, true, true, false, false, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " Jpsi data failed!");
                return 0;
            }
            //fit all events (Jpsi + non-resonant, but no sWeights!)
            if(quickFit(year, false, true, false, false, false, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " data failed!");
                return 0;
            }
            //Apply sWeights from Jpsi channel to MC (and no need to add RefChannel, since the weights are aded automatically there
            if(NtrackWeight(year, false, false, WeightBplusPT, false, sWeight) == 0){
                coutERROR("Creating weights for nTrack distribution for " + year + " MC failed!");
                return 0;
            }
        }
    }
    else{ //reweight data by signal MC
        if(Kst2Kspiplus && SplitDDandLL){
            if(quickFit(year, false, sWeight, false, false, true, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " data (LL tracks) failed!");
                return 0;
            }
            if(quickFit(year, false, sWeight, false, false, false, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " data (DD tracks) failed!");
                return 0;
            }
            if(NtrackWeight(year, false, false, WeightBplusPT, true, sWeight) == 0){
                coutERROR("Creating weights for nTrack distribution for " + year + " MC (LL tracks) failed!");
                return 0;
            }
            if(NtrackWeight(year, false, false, WeightBplusPT, false, sWeight) == 0){
                coutERROR("Creating weights for nTrack distribution for " + year + " MC (DD tracks) failed!");
                return 0;
            }
        }
        else{
            //signal channel
            if(quickFit(year, false, sWeight, false, false, false, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " data failed!");
                return 0;
            }
            if(NtrackWeight(year, false, false, WeightBplusPT, false, sWeight) == 0){
                coutERROR("Creating weights for nTrack distribution for " + year + " MC failed!");
                return 0;
            }
            //reference channel
            if(quickFit(year,false, sWeight, true, false, false, GetShapeFromMC, SignalFunction, BackGroundFunction, ConstrainParameters)  == 0){
                coutERROR("Fitting the B+ mass for " + year + " data failed!");
                return 0;
            }
            if(NtrackWeight(year, true, false, WeightBplusPT, false, sWeight) == 0){
                coutERROR("Creating weights for nTrack distribution for " + year + " MC failed!");
                return 0;
            }
        }
    }

    coutInfo("COMPLETED WEIGHTS FOR YEAR " + year);

    return 1;
}

int WeightAll(bool WeightBplusPT = true, Int_t Run = 1, bool sWeight = true){

    bool GetShapeFromMC = true;

    std::vector<string> years = yearsVector(false,false,false, Run);
    for(unsigned int y = 0; y < years.size(); y++){
        if(WeightYear(years.at(y), WeightBplusPT, GetShapeFromMC, sWeight) == 0){
            coutERROR("sWeighting of the year " + years.at(y) + "did not succeed!");
            return 0;
        }
    }
    coutInfo("Weighting of all data and re-weighting of all signalMC was done!");

    return 1;

}

int ReweightMCOnly(string year = "2011", bool ReferenceChannel = false, bool PHSP = false, bool WeightBplusPT = true, bool sWeight = true){
    if (checkMC(ReferenceChannel,PHSP)==false) return 0;

    if(NtrackWeight(year, ReferenceChannel, PHSP, WeightBplusPT, false, sWeight) == 0){
        coutERROR("Creating weights for nTrack distribution for " + year + string(Kst2Kspiplus && SplitDDandLL ? " (DD tracks) " : "")  + " MC failed!");
        return 0;
    }

    if(Kst2Kspiplus && SplitDDandLL){
        if(NtrackWeight(year, ReferenceChannel, PHSP, WeightBplusPT, true, sWeight) == 0){
            coutERROR("Creating weights for nTrack distribution for " + year + " (LL tracks) MC failed!");
            return 0;
        }
    }

    return 1;
}

int ReweightSignalMC(bool WeightBplusPT = true, Int_t Run = 1, bool sWeight = true){

    std::vector<string> years = yearsMC(false,false,Run);

    for(unsigned int y = 0; y < years.size(); y++){
        if(ReweightMCOnly(years.at(y), false, false, WeightBplusPT, sWeight) == 0){
            coutERROR("Could not reweight signal MC for " + years.at(y) + ".Exit program");
            return 0;
        }
    }

    return 1;

}
int ReweightPHSPMC(bool WeightBplusPT = true, Int_t Run = 1, bool sWeight = true){

    std::vector<string> years = yearsMC(false,true,Run);

    for(unsigned int y = 0; y < years.size(); y++){
        if(ReweightMCOnly(years.at(y), false, true, WeightBplusPT, sWeight) == 0){
            coutERROR("Could not reweight PHSP MC for " + years.at(y) + ".Exit program");
            return 0;
        }
    }

    return 1;

}
int ReweightReferenceMC(bool WeightBplusPT = true, Int_t Run = 1, bool sWeight = true){

    std::vector<string> years = yearsMC(true,false,Run);

    for(unsigned int y = 0; y < years.size(); y++){
        if(ReweightMCOnly(years.at(y), true, false, WeightBplusPT, sWeight) == 0){
            coutERROR("Could not reweight Reference MC for " + years.at(y) + ".Exit program");
            return 0;
        }
    }

    return 1;
}

int ReweightAllMC(bool WeightBplusPT = true, Int_t Run = 1, bool sWeight = true){
    if (ReweightSignalMC(WeightBplusPT,Run,sWeight) == 0) return 0;
    if (ReweightReferenceMC(WeightBplusPT,Run,sWeight) == 0) return 0;
    if (ReweightPHSPMC(WeightBplusPT,Run,sWeight) == 0) return 0;
    return 1;
}
int ReweightAll(bool WeightBplusPT = true, Int_t Run = 1, bool sWeight = true){
    if (WeightAll(WeightBplusPT,Run,sWeight) == 0) return 0;
    if (ReweightAllMC(WeightBplusPT,Run,sWeight) == 0) return 0;
    return 1;
}

int GetSignalAndBckgndEstimation(bool KshortDecaysInVelo = true, Int_t Run = 1){ //assumes S = N - B

    std::vector<string> years = yearsData(Run);

    gStyle -> SetOptStat(0);
    LHCbStyle();
    gROOT->SetBatch(kTRUE);
    //lhcbStyle->SetOptTitle(1);

    bool ConstrainParameters = false; //maybe not needed for signal estimation

    if(Kst2Kspiplus){
        ConstrainParameters = true;
    }

    string SignalFunction = "OneCB";
    string BackGroundFunction;
    if(Kst2Kspiplus){
      if(KshortDecaysInVelo)    BackGroundFunction = "SingleExponential";
      else                      BackGroundFunction = "DoubleExponential";
    }
    else  BackGroundFunction = "SingleExponential";


    //Get the yield of the reference channel from the quickfit, then scale it via the Branching Ratios of both decays
    Int_t SignalYield = 0;
    for(unsigned int y = 0; y < years.size(); y++)
        SignalYield += quickFit(years.at(y).c_str(), true, false, true, false, KshortDecaysInVelo, false, SignalFunction, BackGroundFunction, ConstrainParameters) * BR_sig / BR_ref;

    //get number of candidates in the signal channel within the B_plus mass window
    TChain * tree   = new TChain("DecayTree");

    for(unsigned int y = 0; y < years.size(); y++) tree->Add(GetBDTinputFile(years.at(y),false,false,false,KshortDecaysInVelo).c_str());

    string q2Cuts = getMuMucut();
    string sVariable = (UseDTF ? "B_plus_M_DTF" : "B_plus_M");
    string massCuts = sVariable + "> " + to_string(FitValuesSignal.sig_mean.val - SignalRegionNsigma * FitValuesSignal.sig_effSigma.val)
                        + " && " +
                      sVariable + " < " +to_string(FitValuesSignal.sig_mean.val + SignalRegionNsigma * FitValuesSignal.sig_effSigma.val);
    string cut = "(" + q2Cuts + ") && (" + massCuts + ")";


    //data: create histograms from tree
    TCanvas * c1 = new TCanvas("c1", "c1");
    tree->Draw(Form("%s>>B_plus_M_plot", sVariable.c_str()), cut.c_str());
    TH1 * histo = ((TH1 *) gPad->GetPrimitive("B_plus_M_plot"));

    //... get number of entries from histogram as N = B + S (!!!)
    Int_t EventsAfterPreSelection = histo->GetEntries();

    //correct title and axis labels and save histogram to PDF
    string title = "B^{+} mass after preselection";
    if(Kst2Kspiplus && SplitDDandLL)title.append(KshortDecaysInVelo ? " [LL tracks]" : " [DD tracks]");
    histo->GetXaxis()->SetTitle("m(B^{+}) ( MeV/c^{2} )");
    histo->GetYaxis()->SetTitle(Form("Events / ( %.1f MeV/c^{2} )", histo->GetBinWidth(1)));
    histo->SetTitle(title.c_str());
    histo->Draw();
    if(Kst2Kpluspi0Resolved || Kst2Kpluspi0Merged){
        c1->Print(Form("%s/SignalYieldHistos/%s_Run%i_SignalYield.eps", path_to_output_KplusPizero.c_str(), TheDecay.c_str(), Run));
        c1->SaveAs(Form("%s/SignalYieldHistos/%s_Run%i_SignalYield.root", path_to_output_KplusPizero.c_str(), TheDecay.c_str(), Run));
    }
    else{
        c1->Print(Form("%s/SignalYieldHistos/%s%s_Run%i_SignalYield.eps", path_to_output_KshortPiplus.c_str(), TheDecay.c_str(),(SplitDDandLL ? (KshortDecaysInVelo ? "_LL" : "_DD") : ""), Run));
        c1->SaveAs(Form("%s/SignalYieldHistos/%s%s_Run%i_SignalYield.root", path_to_output_KshortPiplus.c_str(), TheDecay.c_str(),(SplitDDandLL ? (KshortDecaysInVelo ? "_LL" : "_DD") : ""), Run));
    }

    delete c1;
    delete histo;

    if (Kst2Kspiplus){
        coutInfo("[SIGNAL]: " + string(Kst2Kspiplus && SplitDDandLL ? (KshortDecaysInVelo ? "LL Tracks: " : "DD Tracks: ") : "") + to_string(SignalYield));
        coutInfo("[BCKGND]: " + string(Kst2Kspiplus && SplitDDandLL ? (KshortDecaysInVelo ? "LL Tracks: " : "DD Tracks: ") : "") + to_string(EventsAfterPreSelection - SignalYield));
    }
    else{
        coutInfo(TheDecay + "[SIGNAL]: " + to_string(SignalYield));
        coutInfo(TheDecay + "[BCKGND]: " + to_string(EventsAfterPreSelection - SignalYield));
    }
    return 1;

}


