//Truth matching of MC data for the decay B+->Kst+mumu
//david gerick
//Renata Kopecna


#include "GlobalFunctions.hh"
#include "Design.hpp"
#include "MassFit.hpp"
#include "Paths.hpp"
#include "TMsource.cpp"

bool FixParameters = false;


using namespace std;
using namespace RooFit ;

//////////////////////////////////////////////////////
///  MCtruth()
/// compares the particleID of all reconstructed particles in the chain with the trueID
/// ID-check of the (grand)mother particles will be performed!
///
///  TruthMatchAll()
/// run the MCtruth() for all four configurations of 2011/2012 and down/up
///
///  PlotTruthMatchingEfficienciesToPDF()
/// Take the four TGraphErrors from the TFile, plot them nicely into canvases and print it to pdf files in a dedicated folder!
///
///  PlotAllTruthMatchingEfficienciesToPDF()
/// The function name explains it all ;)
///

Double_t YieldBeforeTruthMatching = 0., YieldAfterTruthMatching = 0.;
Double_t YieldBeforeTruthMatchingErr = 0., YieldAfterTruthMatchingErr = 0.;

//1: passed gammas (both), 2: one passed, one converted, 3: both converted, 4: one passed, one random, 5: one converted, one random, 6: both failed
int passGammas(gamma_IDs gamma1, gamma_IDs gamma2, counters *counter,
               bool B_plus_TM, bool K_star_plus_TM, bool pi_zero_TM, bool gamma_TM, bool gamma_TM_full){

    bool pass1 = gamma1.isGammaTrue(counter,gamma_TM,gamma_TM_full,pi_zero_TM,K_star_plus_TM,B_plus_TM);
    bool pass2 = gamma2.isGammaTrue(counter,gamma_TM,gamma_TM_full,pi_zero_TM,K_star_plus_TM,B_plus_TM);

    bool conv1 = gamma1.isGammaConversion(counter,gamma_TM,pi_zero_TM,K_star_plus_TM);
    bool conv2 = gamma2.isGammaConversion(counter,gamma_TM,pi_zero_TM,K_star_plus_TM);

    if (pass1 && pass2){
        coutDebug("Both gammas passed.");
        return 1;
    }
    else{
        if (pass1 || pass2){
            if (conv1 || conv2){
                coutDebug("One gamma passed, one is converted.");
                return 2;
            }
            else{
                coutDebug("One gamma passed, one is random.");
                return 4;
            }
        }
        else{
            if (conv1 && conv2){
                coutDebug("Both gammas are converted.");
                return 3;
            }
            else{
                if (conv1 || conv2){
                    coutDebug("One gamma is converted, one is random.");
                    return 5;
                }
                else{
                    coutDebug("Both gammas failed.");
                    return 6;
                }
            }
        }
    }
}

bool passAllTM(B_plus_IDs B_plus,K_star_IDs K_star,J_psi_IDs J_psi,mu_IDs mu_plus, mu_IDs mu_minus,
               K_plus_IDs K_plus,pi_zero_IDs pi_zero,
               K_short_IDs K_short,pi_plus_IDs pi_plus,Ks_pi_IDs Ks_pi_plus, Ks_pi_IDs Ks_pi_minus,
               counters *counter,double pi_zero_resolved_M, double B_plus_M,
               TH1D *h_pi0_mismatched_mass,TH1D *h_Bplus_M_pi0mismatched, TH1D *h_Bplus_M_noPi0Constr,TH1D *h_pi0_resolved_mass,TH1D *h_Bplus_M_pi0random,
               bool B_plus_TM, bool K_star_plus_TM, bool K_plus_TM, bool pi_zero_TM,  bool mu_TM, bool J_psi_TM, bool pi_zero_TM_full ){

    if(!B_plus.isBplus(counter,B_plus_TM)) return false;
    coutDebug("B+:\t passed");

    if (!K_star.isKstar(counter,K_star_plus_TM, B_plus_TM))return false;
    coutDebug("K*:\t passed");

    if (!J_psi.isJpsi(counter,J_psi_TM, B_plus_TM))return false;
    coutDebug("J/psi:\t passed");

    if (!mu_plus.isMu(counter, mu_TM, J_psi_TM, B_plus_TM))return false;
    coutDebug("Mu+:\t passed");

    if (!mu_minus.isMu(counter, mu_TM, J_psi_TM, B_plus_TM))return false;
    coutDebug("Mu-:\t passed");

    if (Kst2Kspiplus){
        if (!K_short.isKshort(counter)) return false;
        if (!pi_plus.isPiPlus(counter)) return false;
        if (!Ks_pi_plus.isKsPi(counter)) return false;
        if (!Ks_pi_minus.isKsPi(counter)) return false;
    }
    else if (Kst2Kpluspi0Resolved){
        if (!K_plus.isKplus(counter, K_plus_TM, K_star_plus_TM, B_plus_TM))return false;
        coutDebug("K+:\t passed");
        h_Bplus_M_noPi0Constr->Fill(B_plus_M);
        if (pi_zero_TM_full){
            if (!pi_zero.isPi0(counter, pi_zero_TM, K_star_plus_TM, B_plus_TM, pi_zero_resolved_M,  B_plus_M, h_pi0_mismatched_mass, h_Bplus_M_pi0mismatched, h_pi0_resolved_mass)){
                h_Bplus_M_pi0random->Fill(B_plus_M);
                return false;
            }
            else coutDebug("pi0:\t passed");
        }
    }
    return true;
}

Int_t MCtruth(string year = "2011", string magnet = "down", bool ReferenceChannel = false, bool PHSP = false, bool B0 = false, bool K1 = false, bool Inc = false) {

    gSystem->Load("./Design_cpp.so");

    if(ReferenceChannel && PHSP){
        coutWarning("Cannot set boolean of reference channel and phase-space MC at the same time! Process Reference Channel!");
        PHSP = false;
    }

    //Load MC data from preselected TFile
    TChain * treeMC   = new TChain("DecayTree");
    string string_input = "";
    if (B0 || K1 || Inc)  string_input = GetInputFileBkg(year, magnet, true, ReferenceChannel, B0, K1, Inc);
    else  string_input = GetInputFile(year,magnet,true,true,ReferenceChannel,PHSP,false);
    treeMC->Add(string_input.c_str());

    //Save TMed data
    string string_output = "";
    if (B0 || K1 || Inc) string_output = GetInputFileBkg(year, magnet, true, ReferenceChannel, B0, K1, Inc);
    else string_output = GetInputFile(year,magnet,true,true,ReferenceChannel,PHSP,false);
    replace(string_output,".root","_TM.root");

    TFile * outputFile = outputFile = new TFile(string_output.c_str(), "RECREATE");
    Int_t nEvents = treeMC->GetEntries();
    if (nEvents==0){
        coutERROR("No events in the tree! Exit.");
        return 0;
    }

    //BKGCAT branch
    Int_t BKGCAT = -10;

    B_plus_IDs B_plus_ID;
    K_star_IDs K_star_ID;
    J_psi_IDs J_psi_ID(ReferenceChannel);
    mu_IDs mu_plus_ID(ReferenceChannel);
    mu_IDs mu_minus_ID(ReferenceChannel);
    K_plus_IDs K_plus_ID;
    pi_zero_IDs pi_zero_ID;
    gamma_IDs gamma1_ID;
    gamma_IDs gamma2_ID;
    K_short_IDs K_short_ID;
    pi_plus_IDs pi_plus_ID;
    Ks_pi_IDs Ks_pi_plus_ID;
    Ks_pi_IDs Ks_pi_minus_ID;


    //link variables to branches
    if(UseBKGCAT){
        treeMC -> SetBranchAddress( "B_plus_BKGCAT" , &BKGCAT );
    }
    if (Kst2Kpluspi0Resolved){//I need control plots to check pi0 TM and so on
        treeMC -> SetBranchAddress( "B_plus_TRUEID" ,       &B_plus_ID.B_plus_TRUEID );
        treeMC -> SetBranchAddress( "K_star_plus_TRUEID" ,  &K_star_ID.K_star_plus_TRUEID );
        treeMC -> SetBranchAddress( "mu_plus_TRUEID" ,      &mu_plus_ID.mu_TRUEID );
        treeMC -> SetBranchAddress( "mu_minus_TRUEID" ,     &mu_minus_ID.mu_TRUEID );

        treeMC -> SetBranchAddress( "K_star_plus_MC_MOTHER_ID" , &K_star_ID.K_star_plus_MOTHER_ID );
        treeMC -> SetBranchAddress( "mu_plus_MC_MOTHER_ID" ,     &mu_plus_ID.mu_MOTHER_ID );
        treeMC -> SetBranchAddress( "mu_minus_MC_MOTHER_ID" ,    &mu_minus_ID.mu_MOTHER_ID );
        treeMC -> SetBranchAddress( "mu_plus_MC_GD_MOTHER_ID" ,  &mu_plus_ID.mu_GD_MOTHER_ID );
        treeMC -> SetBranchAddress( "mu_minus_MC_GD_MOTHER_ID" , &mu_minus_ID.mu_GD_MOTHER_ID );

        if(Kst2Kpluspi0Resolved){
            treeMC -> SetBranchAddress( "B_plus_BKGCAT" ,           &BKGCAT );
            treeMC -> SetBranchAddress( "K_plus_TRUEID" ,           &K_plus_ID.K_plus_TRUEID );
            treeMC -> SetBranchAddress( "pi_zero_resolved_TRUEID" , &pi_zero_ID.pi_zero_TRUEID );
            treeMC -> SetBranchAddress( "gamma1_TRUEID" ,           &gamma1_ID.gamma_TRUEID );
            treeMC -> SetBranchAddress( "gamma2_TRUEID" ,           &gamma2_ID.gamma_TRUEID );

            treeMC -> SetBranchAddress( "K_plus_MC_MOTHER_ID" ,             &K_plus_ID.K_plus_MOTHER_ID );
            treeMC -> SetBranchAddress( "pi_zero_resolved_MC_MOTHER_ID",    &pi_zero_ID.pi_zero_MOTHER_ID );
            treeMC -> SetBranchAddress( "gamma1_MC_MOTHER_ID" ,             &gamma1_ID.gamma_MOTHER_ID );
            treeMC -> SetBranchAddress( "gamma2_MC_MOTHER_ID" ,             &gamma2_ID.gamma_MOTHER_ID );

            treeMC -> SetBranchAddress( "K_plus_MC_GD_MOTHER_ID" ,          &K_plus_ID.K_plus_GD_MOTHER_ID );
            treeMC -> SetBranchAddress( "pi_zero_resolved_MC_GD_MOTHER_ID", &pi_zero_ID.pi_zero_GD_MOTHER_ID );
            treeMC -> SetBranchAddress( "gamma1_MC_GD_MOTHER_ID" ,          &gamma1_ID.gamma_GD_MOTHER_ID );
            treeMC -> SetBranchAddress( "gamma2_MC_GD_MOTHER_ID" ,          &gamma2_ID.gamma_GD_MOTHER_ID );

            treeMC -> SetBranchAddress( "gamma1_MC_GD_GD_MOTHER_ID" , &gamma1_ID.gamma_GD_GD_MOTHER_ID );
            treeMC -> SetBranchAddress( "gamma2_MC_GD_GD_MOTHER_ID" , &gamma2_ID.gamma_GD_GD_MOTHER_ID );
        }

        if(Kst2Kspiplus){
            treeMC -> SetBranchAddress( "K_short_TRUEID" ,          &K_short_ID.K_short_TRUEID );
            treeMC -> SetBranchAddress( "pi_plus_TRUEID" ,          &pi_plus_ID.pi_plus_TRUEID );
            treeMC -> SetBranchAddress( "Ks_pi_plus_TRUEID" ,       &Ks_pi_plus_ID.Ks_pi_TRUEID );
            treeMC -> SetBranchAddress( "Ks_pi_minus_TRUEID" ,      &Ks_pi_minus_ID.Ks_pi_TRUEID );
            treeMC -> SetBranchAddress( "K_short_MC_MOTHER_ID" ,    &K_short_ID.K_short_MOTHER_ID );
            treeMC -> SetBranchAddress( "pi_plus_MC_MOTHER_ID" ,    &pi_plus_ID.pi_plus_MOTHER_ID );
            treeMC -> SetBranchAddress( "Ks_pi_plus_MC_MOTHER_ID" , &Ks_pi_plus_ID.Ks_pi_MOTHER_ID );
            treeMC -> SetBranchAddress( "Ks_pi_minus_MC_MOTHER_ID" ,&Ks_pi_minus_ID.Ks_pi_MOTHER_ID );

            treeMC -> SetBranchAddress( "K_short_MC_GD_MOTHER_ID" ,     &K_short_ID.K_short_GD_MOTHER_ID );
            treeMC -> SetBranchAddress( "pi_plus_MC_GD_MOTHER_ID" ,     &pi_plus_ID.pi_plus_GD_MOTHER_ID );
            treeMC -> SetBranchAddress( "Ks_pi_plus_MC_GD_MOTHER_ID" ,  &Ks_pi_plus_ID.Ks_pi_GD_MOTHER_ID );
            treeMC -> SetBranchAddress( "Ks_pi_minus_MC_GD_MOTHER_ID" , &Ks_pi_minus_ID.Ks_pi_GD_MOTHER_ID );

            treeMC -> SetBranchAddress( "Ks_pi_plus_MC_GD_GD_MOTHER_ID" ,  &Ks_pi_plus_ID.Ks_pi_GD_GD_MOTHER_ID );
            treeMC -> SetBranchAddress( "Ks_pi_minus_MC_GD_GD_MOTHER_ID" , &Ks_pi_minus_ID.Ks_pi_GD_GD_MOTHER_ID );
        }

        treeMC -> SetBranchAddress( "J_psi_TRUEID" ,        &J_psi_ID.J_psi_TRUEID );
        treeMC -> SetBranchAddress( "J_psi_MC_MOTHER_ID" ,  &J_psi_ID.J_psi_MOTHER_ID );

    }

    Int_t TRUE_events = 0;

    //counter for BKGCAT:
    Int_t BKGCAT_VETO[20] = {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0};
    Int_t BKGCAT_TRUE[20] = {0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0};

    //control plots for pi0
    Double_t pi_zero_resolved_M = 0.0;
    if (Kst2Kpluspi0Resolved) treeMC -> SetBranchAddress( "pi_zero_resolved_M" , &pi_zero_resolved_M );

    //Use either B mass from tuple or from DTF
    Double_t B_plus_M;
    UseDTF ? treeMC->SetBranchStatus( "B_plus_M_DTF", 1)        : treeMC->SetBranchStatus( "B_plus_M", 1);
    UseDTF ? treeMC->SetBranchAddress("B_plus_M_DTF", &B_plus_M): treeMC->SetBranchAddress("B_plus_M", &B_plus_M);

    //Control pi0 histograms
    TH1D * h_pi0_resolved_mass  = generalHistogram("h_pi0_resolved_mass" ,   "m_{#pi^{0}} resolved",      70, 100, 170, "m_{pi^{0}} resolved [MeV]", "Counts / 1MeV", kBlue);
    TH1D * h_pi0_mismatched_mass= generalHistogram("h_pi0_mismatched_mass" , "m_{#pi^{0}} IDed as gamma", 70, 100, 170, "m_{pi^{0}}[MeV] MisIDed",   "Counts / 1MeV", kRed);

    //Control B+ mass for pi0 correctly/wrong/no matched
    TH1D * h_Bplus_M_pi0mismatched= BmassHistogram("", "m_{B^{+}}, #pi^{0} not  #pi^{0}",  "m_{B^{+}}[MeV] mismatched #pi^{0}",kRed);
    TH1D * h_Bplus_M_pi0random    = BmassHistogram("h_Bplus_M_pi0random",     "m_{B^{+}}, #pi^{0} random",        "m_{B^{+}}[MeV] random #pi^{0}",kBlue);
    TH1D * h_Bplus_M_noPi0Constr  = BmassHistogram("h_Bplus_M_noPi0Constr",   "m_{B^{+}}, no #pi^{0} TM",         "m_{B^{+}}[MeV], no #pi^{0} TM",kBlack);

    TH1D * h_Bplus_M_TM_BKGCAT               = BmassHistogram("h_Bplus_M_TM_BKGCAT",                "m_{B^{+}} TMed",                           "m_{B^{+}}[MeV] TMed(BKG CAT)",kGreen+2);
    TH1D * h_Bplus_M_TMed                    = BmassHistogram("h_Bplus_M_TMed",                     "m_{B^{+}} TMed",                           "m_{B^{+}}[MeV] TMed",kGreen+2);
    TH1D * h_Bplus_M_TMed_noPi0              = BmassHistogram("h_Bplus_M_TMed_noPi0",               "m_{B^{+}} TMed, no #pi^{0}",               "m_{B^{+}}[MeV] TMed, no #pi^{0}",kGreen+2);
    TH1D * h_Bplus_M_TMed_noPi0_noBplus      = BmassHistogram("h_Bplus_M_TMed_noPi0_noBplus",       "m_{B^{+}} TMed, no #pi^{0}, B^{+}",        "m_{B^{+}}[MeV] TMed, no #pi^{0}, B^{+}",kGreen+2);
    TH1D * h_Bplus_M_TMed_noPi0_noKst_noBplus= BmassHistogram("h_Bplus_M_TMed_noPi0_noKst_noBplus", "m_{B^{+}} TMed, no #pi^{0}, K^{*}, B^{+}", "m_{B^{+}}[MeV] TMed, no #pi^{0}, K^{*}, B^{+}",kGreen+2);

    TH1D * h_Bplus_M_notTM_BKGCAT               = BmassHistogram("h_Bplus_M_notTM_BKGCAT",                "m_{B^{+}} not TMed",                           "m_{B^{+}}[MeV] not TMed(BKG CAT)",kRed);
    TH1D * h_Bplus_M_notTMed                    = BmassHistogram("h_Bplus_M_notTMed",                     "m_{B^{+}} not TMed",                           "m_{B^{+}}[MeV] not TMed",kRed);
    TH1D * h_Bplus_M_notTMed_noPi0              = BmassHistogram("h_Bplus_M_notTMed_noPi0",               "m_{B^{+}} not TMed, no #pi^{0}",               "m_{B^{+}}[MeV] not TMed, no #pi^{0}",kRed);
    TH1D * h_Bplus_M_notTMed_noPi0_noBplus      = BmassHistogram("h_Bplus_M_notTMed_noPi0_noBplus",     "m_{B^{+}} not TMed, no #pi^{0}, B^{+},",       "m_{B^{+}}[MeV] not TMed, no #pi^{0}, B^{+}",kRed);
    TH1D * h_Bplus_M_notTMed_noPi0_noKst_noBplus= BmassHistogram("h_Bplus_M_notTMed_noPi0_noKst_noBplus", "m_{B^{+}} not TMed, no #pi^{0}, K^{*}, B^{+}", "m_{B^{+}}[MeV] not TMed, no #pi^{0}, K^{*}, B^{+}",kRed);

    TH1D * h_Bplus_M_all       = BmassHistogram("h_Bplus_M_all",       "m_{B^{+}} ",                    "m_{B^{+}}[MeV] preSel",              kBlack);

    //TMP empty histogram (filled when running over not pi0 TM)
    TH1D * h_Bplus_M_tmp = BmassHistogram("h_Bplus_M_tmp", "m_{B^{+}}", "m_{B^{+}}[MeV]",kBlack);

    //control PID plots
    TH1D * h_B_plus_TRUEID       = new TH1D("h_B_plus_TRUEID",       "B^{+} TRUEID",   100000, -1, 99999);
    TH1D * h_K_star_plus_TRUEID  = new TH1D("h_K_star_plus_TRUEID",  "K^{*+} TRUEID",  100000, -1, 99999);
    TH1D * h_K_plus_TRUEID       = new TH1D("h_K_plus_TRUEID",       "K^{+} TRUEID",   100000, -1, 99999);
    TH1D * h_pi_zero_TRUEID      = new TH1D("h_pi_zero_TRUEID",      "#pi^{0} TRUEID", 100000, -1, 99999);
    TH1D * h_gamma_TRUEID        = new TH1D("h_gamma_TRUEID",        "#gamma TRUEID",  100000, -1, 99999);

    TTree * treeTruthMatched = treeMC->CloneTree(0);
    treeTruthMatched->SetName("DecayTreeTruthMatched");

    //If needed to save the eents into a new file (when adding branches such as  TMed), uncomment
    //TMed branch    //0: didn't pass TM,1: passed TM
    Int_t iTMed = -1;
    treeTruthMatched->Branch("TMed", &iTMed, "iTMed/I");
    Int_t iTMedBKGCAT = -1;
    treeTruthMatched->Branch("TMedBKGCAT", &iTMedBKGCAT, "iTMedBKGCAT/I");
    Int_t iTMed_noPi0 = -1;
    if (!KshortChannel) treeTruthMatched->Branch("TMed_noPi0", &iTMed_noPi0, "iTMed_noPi0/I");
    Int_t iTMed_noPi0_noBplus = -1;
    if (!KshortChannel) treeTruthMatched->Branch("TMed_noPi0_noBplus", &iTMed_noPi0_noBplus, "iTMed_noPi0_noBplus/I");
    Int_t iTMed_noPi0_noKst_noBplus = -1;
    if (!KshortChannel) treeTruthMatched->Branch("TMed_noPi0_noKst_noBplus", &iTMed_noPi0_noKst_noBplus, "iTMed_noPi0_noKst_noBplus/I");

    //1: passed gammas (both), 2: one passed, one converted, 3: both converted, 4: one passed, one random, 5: one converted, one random, 6: both failed
    Int_t TM_gammas = -1;
    if (!KshortChannel) treeTruthMatched->Branch("TM_gammas", &TM_gammas, "TM_gammas/I");
    Int_t TM_gammas_noPi0 = -1;
    if (!KshortChannel) treeTruthMatched->Branch("TM_gammas_noPi0", &TM_gammas_noPi0, "TM_gammas_noPi0/I");
    Int_t TM_gammas_noPi0_noBplus = -1;
    if (!KshortChannel) treeTruthMatched->Branch("TM_gammas_noPi0_noBplus", &TM_gammas_noPi0_noBplus, "TM_gammas_noPi0_noBplus/I");
    Int_t TM_gammas_noPi0_noKst_noBplus= -1;
    if (!KshortChannel) treeTruthMatched->Branch("TM_gammas_noPi0_noKst_noBplus", &TM_gammas_noPi0_noKst_noBplus, "TM_gammas_noPi0_noKst_noBplus/I");

    //initialize VETO counters
    counters veto_counters;
    counters veto_counters_noPi0;
    counters veto_counters_noPi0_noBplus;
    counters veto_counters_noPi0_noKst_noBplus;
    //MC loop
    coutInfo("Truth matching of " + to_string(nEvents) + " MC events for " + year + magnet + ".");

    for(int i = 0; i < nEvents; i++){
        if(i%2000==0 && i != 0)coutDebug("Loading MC event "+to_string(i) +" / " +to_string(nEvents));
        treeMC->GetEntry(i);

        h_Bplus_M_all->Fill(B_plus_M);

        //Check the BKGCAT category and fill control histograms
        if (isBKGCAT(BKGCAT,ReferenceChannel)){
            BKGCAT_TRUE[BKGCAT/10]++;
            iTMedBKGCAT = 1;
            TRUE_events++;
            h_Bplus_M_TM_BKGCAT->Fill(B_plus_M);
        }
        else{
            BKGCAT_VETO[BKGCAT/10]++;
            iTMedBKGCAT = 0;
            h_Bplus_M_notTM_BKGCAT->Fill(B_plus_M);
        }

        //======================================
        //check all particles for correct TRUEID
        //======================================

        if (B0 || K1){
            //TODO
        }
        else{
            if (Inc){
                h_B_plus_TRUEID     ->Fill(TMath::Abs(B_plus_ID.B_plus_TRUEID));
                h_K_star_plus_TRUEID->Fill(TMath::Abs(K_star_ID.K_star_plus_TRUEID));
                h_K_plus_TRUEID     ->Fill(TMath::Abs(K_plus_ID.K_plus_TRUEID));
                h_pi_zero_TRUEID    ->Fill(TMath::Abs(pi_zero_ID.pi_zero_TRUEID));

            }
            iTMed = passAllTM(B_plus_ID,K_star_ID,J_psi_ID,mu_plus_ID,mu_minus_ID,
                              K_plus_ID,pi_zero_ID,
                              K_short_ID,pi_plus_ID,Ks_pi_plus_ID,Ks_pi_minus_ID,
                              &veto_counters, pi_zero_resolved_M, B_plus_M,
                              h_pi0_mismatched_mass,h_Bplus_M_pi0mismatched, h_Bplus_M_noPi0Constr, h_pi0_resolved_mass, h_Bplus_M_pi0random,
                              true, true, true, true, true, true, true);

            iTMed_noPi0 = passAllTM(B_plus_ID,K_star_ID,J_psi_ID,mu_plus_ID,mu_minus_ID,
                                    K_plus_ID,pi_zero_ID,
                                    K_short_ID,pi_plus_ID,Ks_pi_plus_ID,Ks_pi_minus_ID,
                                    &veto_counters_noPi0, pi_zero_resolved_M, B_plus_M,
                                    h_pi0_mismatched_mass,h_Bplus_M_pi0mismatched, h_Bplus_M_tmp, h_pi0_resolved_mass, h_Bplus_M_tmp,
                                    true, true, true, false, true, true, false);

            iTMed_noPi0_noBplus = passAllTM(B_plus_ID,K_star_ID,J_psi_ID,mu_plus_ID,mu_minus_ID,
                                            K_plus_ID,pi_zero_ID,
                                            K_short_ID,pi_plus_ID,Ks_pi_plus_ID,Ks_pi_minus_ID,
                                            &veto_counters_noPi0_noBplus, pi_zero_resolved_M, B_plus_M,
                                            h_pi0_mismatched_mass,h_Bplus_M_pi0mismatched, h_Bplus_M_tmp, h_pi0_resolved_mass, h_Bplus_M_tmp,
                                            false, true, true, false, true, true, true);
            iTMed_noPi0_noKst_noBplus= passAllTM(B_plus_ID,K_star_ID,J_psi_ID,mu_plus_ID,mu_minus_ID,
                                                 K_plus_ID,pi_zero_ID,
                                                 K_short_ID,pi_plus_ID,Ks_pi_plus_ID,Ks_pi_minus_ID,
                                                 &veto_counters_noPi0_noKst_noBplus, pi_zero_resolved_M, B_plus_M,
                                                 h_pi0_mismatched_mass,h_Bplus_M_pi0mismatched, h_Bplus_M_tmp, h_pi0_resolved_mass, h_Bplus_M_tmp,
                                                 false, false, true, false, true, true, true);
        }
        if (Kst2Kpluspi0Resolved){
            TM_gammas = passGammas(gamma1_ID,gamma2_ID,&veto_counters, true,true,true,true,true);
            TM_gammas_noPi0 = passGammas(gamma1_ID,gamma2_ID,&veto_counters, true,true,false,true,true);
            TM_gammas_noPi0_noBplus = passGammas(gamma1_ID,gamma2_ID,&veto_counters, false,true,false,true,true);
            TM_gammas_noPi0_noKst_noBplus =passGammas(gamma1_ID,gamma2_ID,&veto_counters, false,false,false,true,true);
        }

        if (iTMed) h_Bplus_M_TMed->Fill(B_plus_M);
        else        h_Bplus_M_notTMed->Fill(B_plus_M);
        if (iTMed_noPi0) h_Bplus_M_TMed_noPi0->Fill(B_plus_M);
        else             h_Bplus_M_notTMed_noPi0->Fill(B_plus_M);
        if (iTMed_noPi0_noBplus) h_Bplus_M_TMed_noPi0_noBplus->Fill(B_plus_M);
        else                    h_Bplus_M_notTMed_noPi0_noBplus->Fill(B_plus_M);
        if (iTMed_noPi0_noKst_noBplus) h_Bplus_M_TMed_noPi0_noKst_noBplus->Fill(B_plus_M);
        else                           h_Bplus_M_notTMed_noPi0_noKst_noBplus->Fill(B_plus_M);

        treeTruthMatched->Fill();

    }//loop over events end

    //output truthmatching results
    coutInfo("Resulting TruthMatching data: " + to_string(TRUE_events) + "/" + to_string(nEvents) + " are found to be correctly reconstructed!" );
    coutInfo("For the sample: " + year + " " + magnet + ":\t" + (ReferenceChannel ? "JpsiChannel" : (PHSP ? "PHSP" : "Signal")) + " MC events" );
    coutInfo("The following BKGCAT have been found in the sample:" );
    coutInfo("BKGCAT\tAccepted\tRejected\tTotal");
    for(int i = 0; i < 20; i++){
        if(BKGCAT_VETO[i] != 0 || BKGCAT_TRUE[i] != 0){
            coutInfo(to_string(i*10) + "\t" + to_string(BKGCAT_TRUE[i]) + "\t\t" + to_string(BKGCAT_VETO[i]) + "\t\t" + to_string(BKGCAT_TRUE[i]+BKGCAT_VETO[i]) );
        }
    }
    veto_counters.printCounters(B0,K1);
    veto_counters_noPi0.printCounters(false,false);
    veto_counters_noPi0_noBplus.printCounters(false,false);
    veto_counters_noPi0_noKst_noBplus.printCounters(false,false);

    if(!outputFile->IsOpen()){
        coutERROR("ile was not opened succesfully!");
        return 0;
    }

    outputFile->cd();
    treeTruthMatched->Write("",TObject::kWriteDelete);
    if (Kst2Kpluspi0Resolved){
        //Save hists into a file
        string path = GetInputFile(year,magnet,true,true,ReferenceChannel,PHSP, false);
        replace(path,".root", "_TMcontrolPlots.root");
        TFile *controlHists = new TFile(path.c_str(),"RECREATE");

        gROOT->SetBatch(kTRUE);
        controlHists->cd();

        h_Bplus_M_pi0mismatched->Write("",TObject::kWriteDelete);
        h_Bplus_M_pi0random->Write("",TObject::kWriteDelete);
        h_Bplus_M_noPi0Constr->Write("",TObject::kWriteDelete);

        h_Bplus_M_TM_BKGCAT->Write("",TObject::kWriteDelete);
        h_Bplus_M_notTM_BKGCAT->Write("",TObject::kWriteDelete);
        h_Bplus_M_TMed->Write("",TObject::kWriteDelete);
        h_Bplus_M_notTMed->Write("",TObject::kWriteDelete);
        h_Bplus_M_TMed_noPi0->Write("",TObject::kWriteDelete);
        h_Bplus_M_notTMed_noPi0->Write("",TObject::kWriteDelete);
        h_Bplus_M_TMed_noPi0_noBplus->Write("",TObject::kWriteDelete);
        h_Bplus_M_notTMed_noPi0_noBplus->Write("",TObject::kWriteDelete);
        h_Bplus_M_TMed_noPi0_noKst_noBplus->Write("",TObject::kWriteDelete);
        h_Bplus_M_notTMed_noPi0_noKst_noBplus->Write("",TObject::kWriteDelete);

        h_Bplus_M_all->Write("",TObject::kWriteDelete);

        if (Inc){
            h_B_plus_TRUEID->Write("",TObject::kWriteDelete);
            h_K_star_plus_TRUEID->Write("",TObject::kWriteDelete);
            h_K_plus_TRUEID->Write("",TObject::kWriteDelete);
            h_pi_zero_TRUEID->Write("",TObject::kWriteDelete);
            h_gamma_TRUEID->Write("",TObject::kWriteDelete);
        }

        h_pi0_mismatched_mass->Write("",TObject::kWriteDelete);
        h_pi0_resolved_mass->Write("",TObject::kWriteDelete);

        //Draw and save
        TM_canvas("Check_pi0",                h_pi0_resolved_mass,h_pi0_mismatched_mass,       h_pi0_mismatched_mass,                  year,ReferenceChannel,PHSP, B0, K1, Inc);
        TM_canvas("BKGCAT",                   h_Bplus_M_all,h_Bplus_M_TM_BKGCAT,               h_Bplus_M_notTM_BKGCAT,                 year,ReferenceChannel,PHSP, B0, K1, Inc);
        TM_canvas("TMed",                     h_Bplus_M_all,h_Bplus_M_TMed,                    h_Bplus_M_notTMed,                      year,ReferenceChannel,PHSP, B0, K1, Inc);
        TM_canvas("TMed_noPi0",               h_Bplus_M_all,h_Bplus_M_TMed_noPi0,              h_Bplus_M_notTMed_noPi0,                year,ReferenceChannel,PHSP, B0, K1, Inc);
        TM_canvas("TMed_noPi0_noBplus",       h_Bplus_M_all,h_Bplus_M_TMed_noPi0_noBplus,      h_Bplus_M_notTMed_noPi0_noBplus,        year,ReferenceChannel,PHSP, B0, K1, Inc);
        TM_canvas("TMed_noPi0_noKst_noBplus", h_Bplus_M_all,h_Bplus_M_TMed_noPi0_noKst_noBplus,h_Bplus_M_notTMed_noPi0_noKst_noBplus,  year,ReferenceChannel,PHSP, B0, K1, Inc);

        controlHists->Close();
    }

    delete h_Bplus_M_pi0mismatched;
    delete h_Bplus_M_pi0random;
    delete h_Bplus_M_noPi0Constr;
    delete h_Bplus_M_all;

    delete h_Bplus_M_TM_BKGCAT;
    delete h_Bplus_M_notTM_BKGCAT;
    delete h_Bplus_M_TMed;
    delete h_Bplus_M_notTMed;
    delete h_Bplus_M_TMed_noPi0;
    delete h_Bplus_M_notTMed_noPi0;
    delete h_Bplus_M_TMed_noPi0_noBplus;
    delete h_Bplus_M_notTMed_noPi0_noBplus;
    delete h_Bplus_M_TMed_noPi0_noKst_noBplus;
    delete h_Bplus_M_notTMed_noPi0_noKst_noBplus;

    delete h_B_plus_TRUEID;
    delete h_K_star_plus_TRUEID;
    delete h_K_plus_TRUEID;
    delete h_pi_zero_TRUEID;
    delete h_gamma_TRUEID;

    delete h_pi0_mismatched_mass;
    delete h_pi0_resolved_mass;

    delete treeMC;
    delete treeTruthMatched;
    outputFile->Close();
    string command = "mv " + string_input + " ";
    replace(string_input,".root","_TMbackup.root");
    command += string_input;
    system(command.c_str());
    replace(string_input,"_TMbackup.root",".root");
    command = "mv " + string_output + " " + string_input;
    system(command.c_str());

    coutInfo("Finished Truth-Matching for sample: " + year + " " + magnet + " : " + TheDecay + (ReferenceChannel ? " RefMC" : (PHSP ? " PHSP MC" : " ") ) );
 
    return 1;

}

bool restore_from_backup(string year = "2011", string magnet = "down", bool ReferenceChannel = false, bool PHSP = false, bool B0 = false, bool K1 = false, bool Inc = false){
    //In case one needs to rerun the TruthMatching
    string string_input = "";
    if (B0 || K1 || Inc)  string_input = GetInputFileBkg(year, magnet, true, ReferenceChannel, B0, K1, Inc);
    else  string_input = GetInputFile(year,magnet,true,true,ReferenceChannel,PHSP,false);

    replace(string_input,".root","_TMbackup.root");
    string command = "mv " + string_input + " ";
    replace(string_input,"_TMbackup.root",".root");
    command += string_input;
    coutDebug(command);
    system(command.c_str());

    coutInfo("Succesfully restored all MC for " + TheDecay + (ReferenceChannel ? " (Reference channel)" : "") + " in year " + year);
    return 1;
}

bool restore_allyears_from_backup(int Run = 1, bool Reference = false, bool PHSP = false, bool B0 = false, bool K1 = false, bool Inc = false){

    std::vector<string> years = (B0 || K1 || Inc) ? yearsBkgMC(Reference, B0, K1, Inc, Run) : yearsMC(Reference,PHSP,Run);
    if (years.empty()){
        coutERROR("Invalid Run number given: " + to_string(Run) + ". Exit program!");
        return 0;
    }
    for(unsigned int y = 0; y < years.size(); y++){
        if(restore_from_backup(years.at(y),"down", Reference, PHSP, B0, K1, Inc) == 0){
            coutERROR("Failed to restore " + years.at(y) + " magnet DOWN for " + TheDecay +  getDataTypeTag(true, Reference, PHSP, B0, K1, Inc) + " from backup");
            return 0;
        }
        if(restore_from_backup(years.at(y),"up", Reference, PHSP, B0, K1, Inc) == 0){
            coutERROR("Failed to restore " + years.at(y) + " magnet UP for " + TheDecay + getDataTypeTag(true, Reference, PHSP, B0, K1, Inc) +" from backup");
            return 0;
        }
        coutInfo("Sucesfully restored " + years.at(y) + " magnet UP for " + TheDecay + getDataTypeTag(true, Reference, PHSP, B0, K1, Inc) +" from backup");
    }
    return 1;
}

bool restore_all_from_backup(){
    if (restore_allyears_from_backup(12, false, false, false, false, false) ==0){
        coutERROR("Failed to restore signal MC from backup!");
        return 0;
    }
    if (restore_allyears_from_backup(12, true, false, false, false, false) ==0){
        coutERROR("Failed to restore reference MC from backup!");
        return 0;
    }
    if (restore_allyears_from_backup(12, false, true, false, false, false) ==0){
        coutERROR("Failed to restore PHSP MC from backup!");
        return 0;
    }

    coutInfo("Succesfully restored all MC for alls samples!");

    return 1;

}

Int_t TruthMatchAllBkg(bool Reference = false, Int_t Run = 1, bool B0 = false, bool K1 = false, bool Inc = false){

    std::vector<string> years = yearsBkgMC(Reference,B0,K1,Inc,Run);
    if (years.empty()){
        coutERROR("Invalid Run number given: " + to_string(Run) + ". Exit program!");
        return 0;
    }

    for(unsigned int y = 0; y < years.size(); y++){
        if(MCtruth(years.at(y),"down", Reference, false, B0, K1, Inc) == 0){
            coutERROR("Failed to process " + years.at(y) + " magnet DOWN for " + TheDecay+ getDataTypeTag(true, Reference, false, B0, K1, Inc) + " from backup");
            return 0;
        }

        if(MCtruth(years.at(y),"up", Reference, false, B0, K1, Inc) == 0){
            coutERROR("Failed to process " + years.at(y) + " magnet UP for " + TheDecay + getDataTypeTag(true, Reference, false, B0, K1, Inc) + " from backup");
            return 0;
        }
    }
    coutInfo("Succesfully finished all truth-matching for " + TheDecay + getDataTypeTag(true, Reference, false, B0, K1, Inc) + " in Run" + to_string(Run));
    return 1;
}

Int_t TruthMatchAll(bool Reference = false, bool PHSP = false, Int_t Run = 1){

    std::vector<string> years = yearsMC(Reference,PHSP,Run);
    if (years.empty()){
        coutERROR("Invalid Run number given: " + to_string(Run) + ". Exit program!");
        return 0;
    }

    for(unsigned int y = 0; y < years.size(); y++){
        if(MCtruth(years.at(y),"down", Reference, PHSP, false, false, false) == 0){
            coutERROR("Failed to process " + years.at(y) + " magnet DOWN for " + TheDecay + getDataTypeTag(true, Reference, PHSP, false, false, false) + " from backup");
            return 0;
        }
        if(MCtruth(years.at(y),"up", Reference, PHSP, false, false, false) == 0){
            coutERROR("Failed to process " + years.at(y) + " magnet UP for " + TheDecay + getDataTypeTag(true, Reference, PHSP, false, false, false) + " from backup");
            return 0;
        }
    }
    coutInfo("Succesfully finished all truth-matching for " + TheDecay + getDataTypeTag(true, Reference, PHSP, false, false, false) + " in Run" + to_string(Run));
    return 1;
}

Int_t TruthMatchAllAll(Int_t Run){
    if (TruthMatchAll(false, false, Run) ==0) return 0;
    if (TruthMatchAll(true,  false, Run) ==0) return 0;
    if (TruthMatchAll(false, true,  Run) ==0) return 0;
    return 1;
}



//This is a reliq from the TM issue we had back in the day in DaVinci, not really used now
Int_t CombinePolarityPi0MisID(string year ="2011", bool ReferenceChannel = false, bool PHSP = false){

    gROOT->SetBatch(kTRUE);
    string filePath = "";

    //Define combined TH1 and make it pretty
    TH1D * h_Bplus_M_pi0mismatched= new TH1D("h_Bplus_M_pi0mismatched" , "m_{B^{+}}, #pi^{0} not TMed", 100, cut_B_plus_M_low, cut_B_plus_M_high);
    TCanvas *canvas_B = new TCanvas("canvas", "#pi^0 TM check", 10,10, 1200,600);
    design_canvas_Bmass(canvas_B,"#pi^0 TM check");
    int counts = (cut_B_plus_M_high-cut_B_plus_M_low)/100;
    design_lines(h_Bplus_M_pi0mismatched,"m_{B^{+}}, #pi^{0} not TMed","Counts / " + to_string(counts)+"MeV",kRed+1);
    h_Bplus_M_pi0mismatched->GetYaxis()->SetTitleOffset(0.7);
    string Magnets[] = {"down", "up"};

    //Loop over polarities and load the histograms
    for(string & magnet : Magnets){ //loop over both magnet polaritys
        string path = GetInputFile(year,magnet,true,true,ReferenceChannel,PHSP, false);
        replace(path,".root", "_TMcontrolPlots.root");
        TFile *file = TFile::Open(path.c_str());
        TH1D *tmp = (TH1D*)file->Get("h_Bplus_M_pi0mismatched");
        h_Bplus_M_pi0mismatched->Add(tmp);
        file->Close();
    }

    //Save the combined TH1 in a canvas
    canvas_B->cd();
    h_Bplus_M_pi0mismatched->Draw();

    filePath = "./../../data/figures/pi0TM/" + year + (ReferenceChannel ? "_Ref":"") + (PHSP ? "_PHSP":"") + "_Bmass";

    string adr = filePath + ".eps";
    canvas_B->Print(adr.c_str(), "eps");
    adr = filePath + ".root";
    canvas_B->Print(adr.c_str(), "root");

    return 1;
}

//Also obsolete (see one function up)
Int_t CombineYearsPi0MisID(Int_t Run =1, bool ReferenceChannel = false, bool PHSP = false){

    gROOT->SetBatch(kTRUE);
    string filePath = "";

    //Define combined TH1 and make it pretty
    TH1D * h_Bplus_M_pi0mismatched= new TH1D("h_Bplus_M_pi0mismatched" , "m_{B^{+}}, #pi^{0} not TMed", 100, cut_B_plus_M_low, cut_B_plus_M_high);
    TCanvas *canvas_B = new TCanvas("canvas", "#pi^0 TM check", 10,10, 1200,600);
    design_canvas_Bmass(canvas_B,"#pi^0 TM check");
    int counts = (cut_B_plus_M_high-cut_B_plus_M_low)/100;
    design_lines(h_Bplus_M_pi0mismatched,"m_{B^{+}}, #pi^{0} not TMed","Counts / " + to_string(counts)+"MeV",kRed+1);
    h_Bplus_M_pi0mismatched->GetYaxis()->SetTitleOffset(0.7);

    //Loop over years and polarities, load the histograms
    std::vector<string> years;
    if(Run == 1){
      years.push_back("2011");
      years.push_back("2012");
    }
    if(Run == 2){
      years.push_back("2015");
      years.push_back("2016");
      if(Use2017 && !ReferenceChannel)years.push_back("2017");
      if(Use2018 && !ReferenceChannel)years.push_back("2018");
    }
    string Magnets[] = {"down", "up"};

    for(unsigned int y = 0; y < years.size(); y++){ //loop over all years
        for(string & magnet : Magnets){ //loop over both magnet polaritys
            string path = GetInputFile(years.at(y),magnet,true,true,ReferenceChannel,PHSP, false);
            replace(path,".root", "_TMcontrolPlots.root");
            TFile *file = TFile::Open(path.c_str());
            TH1D *tmp = (TH1D*)file->Get("h_Bplus_M_pi0mismatched");
            h_Bplus_M_pi0mismatched->Add(tmp);
            file->Close();
        }
    }

    //Save the combined TH1 in a canvas
    canvas_B->cd();
    h_Bplus_M_pi0mismatched->Draw();

    filePath = "./../../data/figures/pi0TM/Run" + to_string(Run) + (ReferenceChannel ? "_Ref":"") + (PHSP ? "_PHSP":"") + "_Bmass";
    string adr = filePath + ".eps";
    canvas_B->Print(adr.c_str(), "eps");
    adr = filePath + ".root";
    canvas_B->Print(adr.c_str(), "root");

    return 1;
}

