
#include "../GlobalFunctions.hh"

//////////////////////////////////////////////////////////////////////
/// QuickChecker()
///
/// Epicly flexible tool to quickly create basic plots for checking various features.
///
/// All definitions happen in the AllBranches constructor, where used branches are defined.
/// Cuts are also defined in the constructor, the strings are transferred to CopyTree.
/// EpicOptimizer() doesn't need to be touched at all.
///
///
/// Options to create either standalone years plots (Run==0),
/// plots per Run (Run==1 || Run ==2) or both runs combined (Run==12)
///
/// Possibility to run on stripped/preselected data/MC/PHSP/Ref.
/// In the case of MC also possibility to run only on TruthMatched data.
///
/// TODO: should add an option tfor sweighted and BDT-cutted data
///
/// Standard output in data/MC/PHSP/Ref folder, year/run and preselection option in the name.
///
/// Disclaimer: Probably leaky and not super-safe to you, so proceed with caution.
///
///

/*******************************/
class Int_branch{
public:
    int BranchVar;      //where to save the branch
    const char*  BranchName;  //Name of the branch
    Int_branch(); //default constructor
    Int_branch(int var, const char* name); //default constructor
    ~Int_branch();   //destuctor
};
Int_branch::Int_branch(int var, const char* name){ //here add what branches you wanna load
    BranchVar = var;
    BranchName = name;
}
/*******************************/

class Float_branch{
public:
    float BranchVar;    //where to save the branch
    const char*  BranchName;  //Name of the branch
    Float_branch(); //default constructor
    ~Float_branch();   //destuctor
    Float_branch(float var, const char* name); //constructor
};
Float_branch::Float_branch(float var, const char* name){ //here add what branches you wanna load
    BranchVar = var;
    BranchName = name;
}

/*******************************/
class Float_arr_branch{
public:
    float *BranchVar;    //where to save the branch
    const char*  BranchName;  //Name of the branch
    Float_arr_branch(); //default constructor
    Float_arr_branch(float *var, const char* name); //constructor
    ~Float_arr_branch();   //destuctor
};
Float_arr_branch::Float_arr_branch(float *var, const char* name){ //here add what branches you wanna load
    BranchVar = var;
    BranchName = name;
}

/*******************************/
class Double_branch{
public:
    double BranchVar;   //where to save the branch
    const char*  BranchName;  //Name of the branch
    Double_branch(); //default constructor
    Double_branch(double var, const char* name); //constructor
    ~Double_branch();   //destuctor
};
Double_branch::Double_branch(double var, const char* name){ //here add what branches you wanna load
    BranchVar = var;
    BranchName = name;
}

/*******************************/
class Double_arr_branch{
public:
    double *BranchVar;    //where to save the branch
    const char*  BranchName;  //Name of the branch
    Double_arr_branch(); //default constructor
    Double_arr_branch(double *var, const char* name); //constructor
    ~Double_arr_branch();   //destuctor
};
Double_arr_branch::Double_arr_branch(double *var, const char* name){ //here add what branches you wanna load
    BranchVar = var;
    BranchName = name;
}

/*******************************/
class AllBranches{
public:
    vector <Int_branch*>     IntBranches;
    vector <Float_branch*>   FloatBranches;
    vector <Float_arr_branch*>   FloatArrBranches;
    vector <Double_branch*>  DoubleBranches;
    vector <Double_arr_branch*>  DoubleArrBranches;
    vector<std::string> Cuts;
    AllBranches(); //default constructor
    ~AllBranches();   //destuctor
};


AllBranches::AllBranches(){ //here add what branches you wanna load

   // int I_nTracks = 0;
   // IntBranches.push_back(new Int_branch(I_nTracks, "nTracks"));
/*
    double D_K_star_plus_PT = 0.0;
    DoubleBranches.push_back(new Double_branch(D_K_star_plus_PT, "K_star_plus_PT"));
    Cuts.push_back("K_star_plus_PT > 1300");
 //Cuts.push_back("K_star_plus_PT < 300");
 */

    double D_B_plus_ThetaL = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_ThetaL, "B_plus_ThetaL"));
    double D_B_plus_ThetaK = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_ThetaK, "B_plus_ThetaK"));
    double D_B_plus_Phi = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_Phi, "B_plus_Phi"));


    double D_B_plus_TRUEThetaL = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_TRUEThetaL, "B_plus_TRUEThetaL"));
    double D_B_plus_TRUEThetaK = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_TRUEThetaK, "B_plus_TRUEThetaK"));
    double D_B_plus_TRUEPhi = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_TRUEPhi, "B_plus_TRUEPhi"));


    double D_B_plus_M = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_M, "B_plus_M"));
    Cuts.push_back("B_plus_M < 600000000");

    /*
    double D_B_plus_M = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_M, "B_plus_M"));
    Cuts.push_back("B_plus_M > 5000");
    Cuts.push_back("B_plus_M < 6000");
    float K_star_plus_M = 0;
    DoubleBranches.push_back(new Double_branch(K_star_plus_M,"K_star_plus_M"));
    Cuts.push_back("K_star_plus_M > 792");
    Cuts.push_back("K_star_plus_M <1050");
    double D_K_plus_PT = 0.0;
    DoubleBranches.push_back(new Double_branch(D_K_plus_PT, "K_plus_PT"));
    Cuts.push_back("K_plus_PT > 300");

    double D_B_plus_PT = 0.0;
    DoubleBranches.push_back(new Double_branch(D_B_plus_PT, "B_plus_PT"));
    Cuts.push_back("B_plus_PT > 2000");

    double K_star_plus_FDCHI2_OWNPV = 0.0;
    DoubleBranches.push_back(new Double_branch(K_star_plus_FDCHI2_OWNPV, "K_star_plus_FDCHI2_OWNPV"));
    Cuts.push_back("K_star_plus_FDCHI2_OWNPV > -100000");

    double B_plus_FDCHI2_OWNPV = 0.0;
    DoubleBranches.push_back(new Double_branch(B_plus_FDCHI2_OWNPV, "B_plus_FDCHI2_OWNPV"));
    Cuts.push_back("B_plus_FDCHI2_OWNPV > 121");

    double K_star_plus_FDCHI2_OWNPV = 0.0;
    DoubleBranches.push_back(new Double_branch(K_star_plus_FDCHI2_OWNPV, "K_star_plus_FDCHI2_OWNPV"));
    Cuts.push_back("K_star_plus_FDCHI2_OWNPV > -100000");

    int B_plus_Hlt2Topo3BodyBBDTDecision_TOS = 0.0;
    int B_plus_Hlt2TopoMu2BodyBBDTDecision_TOS = 0.0;
    int B_plus_Hlt2TopoMu3BodyBBDTDecision_TOS = 0.0;
    int B_plus_Hlt2TopoMuMu2BodyBBDTDecision_TOS = 0.0;
    int B_plus_Hlt2TopoMuMu3BodyBBDTDecision_TOS = 0.0;


    double D_pi_zero_resolved_TRUEP_X = 0.0;
    DoubleBranches.push_back(new Double_branch(D_pi_zero_resolved_TRUEP_X, "pi_zero_resolved_TRUEP_X"));
    Cuts.push_back("pi_zero_resolved_TRUEP_X > -100000");

    double D_pi_zero_resolved_TRUEP_Y = 0.0;
    DoubleBranches.push_back(new Double_branch(D_pi_zero_resolved_TRUEP_Y, "pi_zero_resolved_TRUEP_Y"));
    Cuts.push_back("pi_zero_resolved_TRUEP_Y  > -100000");

    double D_pi_zero_resolved_TRUEP_Z = 0.0;
    DoubleBranches.push_back(new Double_branch(D_pi_zero_resolved_TRUEP_Z, "pi_zero_resolved_TRUEP_Z"));
    Cuts.push_back("pi_zero_resolved_TRUEP_Z > -100000");

    double D_pi_zero_resolved_ETA = 0.0;
    DoubleBranches.push_back(new Double_branch(D_pi_zero_resolved_ETA, "pi_zero_resolved_ETA"));
    Cuts.push_back("pi_zero_resolved_ETA > -100000");


    double D_pi_zero_resolved_PX = 0.0;
    DoubleBranches.push_back(new Double_branch(D_pi_zero_resolved_PX, "pi_zero_resolved_PX"));
    Cuts.push_back("pi_zero_resolved_PX > -100000");

    double D_pi_zero_resolved_PY = 0.0;
    DoubleBranches.push_back(new Double_branch(D_pi_zero_resolved_PY, "pi_zero_resolved_PY"));
    Cuts.push_back("pi_zero_resolved_PY > -100000");

    double D_pi_zero_resolved_PZ = 0.0;
    DoubleBranches.push_back(new Double_branch(D_pi_zero_resolved_PZ, "pi_zero_resolved_PZ"));
    Cuts.push_back("pi_zero_resolved_PZ> -100000");


    double D_K_star_plus_M_DTF = 0.0;
    DoubleBranches.push_back(new Double_branch(D_K_star_plus_M_DTF, "K_star_plus_M_DTF"));
    Cuts.push_back("K_star_plus_M_DTF > 792");
    Cuts.push_back("K_star_plus_M_DTF <1050");
*/
/*
    float K_star_plus_M = 0;
    DoubleBranches.push_back(new Double_branch(K_star_plus_M,"K_star_plus_M"));
    Cuts.push_back("K_star_plus_M > 792");
    Cuts.push_back("K_star_plus_M <1050");
*/
/*
       float F_B_plus_DTF_Kst_892_plus_M[20];
       FloatArrBranches.push_back(new Float_arr_branch(F_B_plus_DTF_Kst_892_plus_M,"B_plus_DTF_Kst_892_plus_M"));
       Cuts.push_back("B_plus_DTF_Kst_892_plus_M[0] > 792");
       Cuts.push_back("B_plus_DTF_Kst_892_plus_M[0] <1050");
*/
       //truthmatching
       //assign variables to TRUEID
     /*  int B_plus_TRUEID= 0;;
       IntBranches.push_back(new Int_branch(B_plus_TRUEID, "B_plus_TRUEID"));
       Cuts.push_back("abs(B_plus_TRUEID) == " + to_string(TRUEID.B_PLUS));

       int K_star_plus_TRUEID= 0;
       IntBranches.push_back(new Int_branch(K_star_plus_TRUEID, "K_star_plus_TRUEID"));
       Cuts.push_back("abs(K_star_plus_TRUEID) == " + to_string(TRUEID.K_STAR_PLUS));

       int K_plus_TRUEID= 0;
       IntBranches.push_back(new Int_branch(K_plus_TRUEID, "K_plus_TRUEID"));
       Cuts.push_back("abs(K_plus_TRUEID) == " + to_string(TRUEID.K_PLUS));

       int pi_zero_TRUEID= 0;
       IntBranches.push_back(new Int_branch(pi_zero_TRUEID, "pi_zero_resolved_TRUEID"));
       Cuts.push_back("abs(pi_zero_resolved_TRUEID) == " + to_string(TRUEID.PI_ZERO));

       int gamma1_TRUEID= 0;
       IntBranches.push_back(new Int_branch(gamma1_TRUEID, "gamma1_TRUEID"));
       Cuts.push_back("abs(gamma1_TRUEID) == " + to_string(TRUEID.GAMMA));

       int gamma2_TRUEID= 0;
       IntBranches.push_back(new Int_branch(gamma2_TRUEID, "gamma2_TRUEID"));
       Cuts.push_back("abs(gamma2_TRUEID) == " + to_string(TRUEID.GAMMA));

       //assign variables to TRUE_MOTHER_ID
       int K_star_plus_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(K_star_plus_MOTHER_ID, "K_star_plus_MC_MOTHER_ID"));
       Cuts.push_back("abs(K_star_plus_MC_MOTHER_ID) == " + to_string(TRUEID.B_PLUS));

       int K_plus_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(K_plus_MOTHER_ID, "K_plus_MC_MOTHER_ID"));
       Cuts.push_back("abs(K_plus_MC_MOTHER_ID) == " + to_string(TRUEID.K_STAR_PLUS));

       int pi_zero_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(pi_zero_MOTHER_ID, "pi_zero_resolved_MC_MOTHER_ID"));
       Cuts.push_back("abs(pi_zero_resolved_MC_MOTHER_ID) == " + to_string(TRUEID.K_STAR_PLUS));

       int gamma1_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(gamma1_MOTHER_ID, "gamma1_MC_MOTHER_ID"));
       Cuts.push_back("abs(gamma1_MC_MOTHER_ID) == " + to_string(TRUEID.PI_ZERO));

       int gamma2_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(gamma2_MOTHER_ID, "gamma2_MC_MOTHER_ID"));
       Cuts.push_back("abs(gamma2_MC_MOTHER_ID) == " + to_string(TRUEID.PI_ZERO));

       //assign variables to TRUE_GD_MOTHER_ID
       int K_plus_GD_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(K_plus_GD_MOTHER_ID, "K_plus_MC_GD_MOTHER_ID"));
       Cuts.push_back("abs(K_plus_MC_GD_MOTHER_ID) == " + to_string(TRUEID.B_PLUS));

       int pi_zero_GD_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(pi_zero_GD_MOTHER_ID, "pi_zero_resolved_MC_GD_MOTHER_ID"));
       Cuts.push_back("abs(pi_zero_resolved_MC_GD_MOTHER_ID) == " + to_string(TRUEID.B_PLUS));

       int gamma1_GD_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(gamma1_GD_MOTHER_ID, "gamma1_MC_GD_MOTHER_ID"));
       Cuts.push_back("abs(gamma1_MC_GD_MOTHER_ID) == " + to_string(TRUEID.K_STAR_PLUS));

       int gamma2_GD_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(gamma2_GD_MOTHER_ID, "gamma2_MC_GD_MOTHER_ID"));
       Cuts.push_back("abs(gamma2_MC_GD_MOTHER_ID) == " + to_string(TRUEID.K_STAR_PLUS));


       //assign variables to TRUE_GD_MOTHER_ID
       int gamma1_GD_GD_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(gamma1_GD_GD_MOTHER_ID, "gamma1_MC_GD_GD_MOTHER_ID"));
       Cuts.push_back("abs(gamma1_MC_GD_GD_MOTHER_ID) == " + to_string(TRUEID.B_PLUS));

       int gamma2_GD_GD_MOTHER_ID= 0;
       IntBranches.push_back(new Int_branch(gamma2_GD_GD_MOTHER_ID, "gamma2_MC_GD_GD_MOTHER_ID"));
       Cuts.push_back("abs(gamma2_MC_GD_GD_MOTHER_ID) == " + to_string(TRUEID.B_PLUS));
*/

}



int EpicOptimizer(std::string year = "2012", std::string magnet = "down", int Run = 0, bool MC = false,
                         bool ReferenceChannel = false, bool PHSP = false, bool preselected = true, bool truthMatched = false) {

    int MaxEntries = 1000000;
    //specify year(s)
    std::vector<std::string> years;
    if (!(Run == 0 || Run == 1 || Run == 2 || Run == 12)){
        cout << "[INFO]\t\tWrong run number, allowed options are 0 ,1, 2 or 12!" << endl;
        cout << "[INFO]\t\tUsing Run = 0, meaning year " << year << endl;
        Run = 0;
    }
    switch (Run){
    case (0): {
        years.push_back(year);
        break;
    }
    case (1): {
        years.push_back("2011");
        years.push_back("2012");
        break;
    }
    case (2): {
        years.push_back("2015");
        years.push_back("2016");
        break;
    }
    case (12): {
        years.push_back("2011");
        years.push_back("2012");
        years.push_back("2015");
        years.push_back("2016");
        break;
    }
    }

    if(ReferenceChannel && PHSP){
        std::cout << "[WARNING]\tCannot set boolean of reference channel and phase-space MC at the same time! Process Reference Channel!" << std::endl;
        PHSP = false;
    }

    TFile* output = 0;
    TChain* tree = 0;

    if (!preselected){
        if(Kst2Kpluspi0Resolved)  tree=new TChain("b2KstKpi0mumuResolvedTuple/DecayTree");
        if(Kst2Kpluspi0Merged)    tree=new TChain("b2KstKpi0mumuMergeddTuple/DecayTree");
        if(Kst2Kspiplus)          tree=new TChain("b2KstKs0pimumu_Tuple/DecayTree");
    }
    else{
        if ((MC||ReferenceChannel||PHSP) && truthMatched)   tree=new TChain("DecayTreeTruthMatched");
        else                                                tree=new TChain("DecayTree");
    }


    for (std::vector<string>::iterator y = years.begin(); y != years.end(); ++y){ //load data for given years
        if (!preselected){ //load stripped data
            //Kst2Kpluspi0Resolved case
            if(Kst2Kpluspi0Resolved){
                if(!MC){
                    if(smallSample) tree->Add(Form("%s/data/%s%s/*B2Kstmumu*13.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    else{
                        tree->Add(Form("%s/data/%s%s/*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                        std::cout << "Adding " << Form("%s/data/%s%s/*B2Kstmumu*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()) << endl;
                    }
                }
                else{
                    if(ReferenceChannel){
                        if(smallSample) tree->Add(Form("%s/data/MC/RefKplusPi0/%s%s/*B2KstJpsi*1.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                        else tree->Add(Form("%s/data/MC/RefKplusPi0/%s%s/*B2KstJpsi*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else if(PHSP){
                        if(smallSample) tree->Add(Form("%s/data/PHSP/KplusPi0/%s%s/*B2Kstmumu*1.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                        else tree->Add(Form("%s/data/PHSP/KplusPi0/%s%s/*B2Kstmumu*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else{
                        if(smallSample) tree->Add(Form("%s/data/MC/KplusPi0/%s%s/*B2Kstmumu*1.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                        else tree->Add(Form("%s/data/MC/KplusPi0/%s%s/*B2Kstmumu*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    }
                }
            }

            //Kst2Kpluspi0Merged case
            if(Kst2Kpluspi0Merged){
                if(!MC){
                    if(smallSample) tree->Add(Form("%s/data/%s%s/*B2Kstmumu*13.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    else tree->Add(Form("%s/data/%s%s/*B2Kstmumu*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                }
                else{
                    if(smallSample) tree->Add(Form("%s/data/MC/KplusPi0/%s%s/*B2Kstmumu*1.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    else tree->Add(Form("%s/data/MC/KplusPi0/%s%s/*B2Kstmumu*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                }
            }

            //Kst2Kspiplus case
            if(Kst2Kspiplus){
                tree=new TChain("b2KstKs0pimumu_Tuple/DecayTree");
                if(!MC){
                    if(smallSample) tree->Add(Form("%s/data/%s%s/*B2Kstmumu*23.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    else tree->Add(Form("%s/data/%s%s/*B2Kstmumu*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                }
                else{
                    if(ReferenceChannel){
                        if(smallSample) tree->Add(Form("%s/data/MC/RefKshortPiplus/%s%s/*B2KstJpsi*13.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                        else tree->Add(Form("%s/data/MC/RefKshortPiplus/%s%s/*B2KstJpsi*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else if(PHSP){
                        if(smallSample) tree->Add(Form("%s/data/PHSP/KshortPiplus/%s%s/*B2Kstmumu*1.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                        else tree->Add(Form("%s/data/PHSP/KshortPiplus/%s%s/*B2Kstmumu*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else{
                        if(smallSample) tree->Add(Form("%s/data/MC/KshortPiplus/%s%s/*B2Kstmumu*1.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                        else tree->Add(Form("%s/data/MC/KshortPiplus/%s%s/*B2Kstmumu*.root",path_to_data.c_str(), y->c_str(),magnet.c_str()));
                    }
                }
            }
        }
        else{ //load preselected data
            // Kst2Kpluspi0Resolved case
            if(Kst2Kpluspi0Resolved){
                if(!MC){
                    tree->Add(Form("%s/data/%s%s/%s%s_pi0Resolved.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    cout << Form("%s/data/%s%s/%s%s_pi0Resolved.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()) << endl;
                }
                else{
                    if(ReferenceChannel){
                        tree->Add(Form("%s/data/MC/RefKplusPi0/%s%s/%s%s_pi0Resolved.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else if (PHSP){
                        tree->Add(Form("%s/data/PHSP/KplusPi0/%s%s/%s%s_pi0Resolved.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else{
                        tree->Add(Form("%s/data/MC/KplusPi0/%s%s/%s%s_pi0Resolved.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                        cout << Form("%s/data/MC/KplusPi0/%s%s/%s%s_pi0Resolved.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()) << endl;
                    }
                }
            }

            // Kst2Kpluspi0Merged case
            if(Kst2Kpluspi0Merged){
                if(!MC){
                    tree->Add(Form("%s/data/%s%s/%s%s_pi0Merged.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                }
                else{
                    if(ReferenceChannel){
                        tree->Add(Form("%s/data/MC/RefKplusPi0/%s%s/%s%s_pi0Merged.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else if (PHSP){
                        tree->Add(Form("%s/data/PHSP/KplusPi0/%s%s/%s%s_pi0Merged.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else{
                        tree->Add(Form("%s/data/MC/KplusPi0/%s%s/%s%s_pi0Merged.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    }
                }
            }
            // Kst2Kspiplus case
            if(Kst2Kspiplus){
                if (!MC){
                    tree->Add(Form("%s/data/%s%s/%s%s_piplus.root",path_to_output_KshortPiplus.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                }
                else{
                    if(ReferenceChannel){
                        tree->Add(Form("%s/data/MC/RefKshortPiplus/%s%s/%s%s_piplus.root",path_to_output_KshortPiplus.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else if (PHSP){
                        tree->Add(Form("%s/data/PHSP/KshortPiplus/%s%s/%s%s_piplus.root",path_to_output_KplusPizero.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    }
                    else{
                        tree->Add(Form("%s/data/MC/KshortPiplus/%s%s/%s%s_piplus.root",path_to_output_KshortPiplus.c_str(), y->c_str(),magnet.c_str(), y->c_str(),magnet.c_str()));
                    }
                }
            }
        } //end of preselected data if
    } //end of the year loop

    // set branches here
    tree->SetBranchStatus("*",0);
    cout << "[INFO]\tOld tree entries: " << tree->GetEntries() << endl;


    AllBranches * AllB = NULL;
    AllB  = new AllBranches();


    for (std::vector<Int_branch*>::iterator IB = AllB->IntBranches.begin(); IB != AllB->IntBranches.end(); ++IB){
        tree->SetBranchStatus((*IB)->BranchName,1);
        tree->SetBranchAddress((*IB)->BranchName,&(*IB)->BranchVar);
    }
    for (std::vector<Float_branch*>::iterator FB = AllB->FloatBranches.begin(); FB != AllB->FloatBranches.end(); ++FB){
        tree->SetBranchStatus((*FB)->BranchName,1);
        tree->SetBranchAddress((*FB)->BranchName,&(*FB)->BranchVar);
    }
    for (std::vector<Float_arr_branch*>::iterator FaB = AllB->FloatArrBranches.begin(); FaB != AllB->FloatArrBranches.end(); ++FaB){
        tree->SetBranchStatus((*FaB)->BranchName,1);
        tree->SetBranchAddress((*FaB)->BranchName,&(*FaB)->BranchVar);
    }
    for (std::vector<Double_branch*>::iterator DB = AllB->DoubleBranches.begin(); DB != AllB->DoubleBranches.end(); ++DB){
        tree->SetBranchStatus((*DB)->BranchName,1);
        tree->SetBranchAddress((*DB)->BranchName,&(*DB)->BranchVar);
    }
    for (std::vector<Double_arr_branch*>::iterator DaB = AllB->DoubleArrBranches.begin(); DaB != AllB->DoubleArrBranches.end(); ++DaB){
        tree->SetBranchStatus((*DaB)->BranchName,1);
        tree->SetBranchAddress((*DaB)->BranchName,&(*DaB)->BranchVar);
    }

    //parse cuts
    std::string AllCuts = "";
    if (AllB->Cuts.size()==0){
        cout << "[ERROR]\t No cuts selected, just copying the tree." << endl;
        return 0;
    }

    AllCuts = AllCuts + *AllB->Cuts.begin();
    cout << "[INFO]\t\t MaxEntries: " << MaxEntries << endl;
    cout << "[INFO]\t\t Appling cuts:" << endl;
    cout << "\t\t\t" << AllCuts << endl;
    for (vector<std::string>::iterator t = AllB->Cuts.begin()+1; t != AllB->Cuts.end(); ++t){
        AllCuts = AllCuts + " && " + *t ;
        cout << "\t\t\t" << *t << endl;
    }
    AllB->Cuts.clear();
    TTree* NewTree = tree->CopyTree( AllCuts.c_str(),"",MaxEntries,0 );
    cout << "[INFO]\tNew tree entries: " << NewTree->GetEntries() << endl;

    //Set the appropriate outputfile
    string path = "";
    if (Run==0) Run = std::stoi( year );
    // Kst2Kpluspi0Resolved case
    if(Kst2Kpluspi0Resolved){
        path = path_to_output_KplusPizero + "/data/";
        if(!MC){
            path =path+ to_string(Run) + magnet + "_" +(preselected ? "preselected" : "") + "_quickcheck_pi0Resolved.root";
            output = new TFile(path.c_str(),"RECREATE");
        }
        else{
            if(ReferenceChannel)path =path +"MC/RefKplusPi0";
            else if (PHSP) path =path +"PHSP/KplusPi0/";
            else path =path +"MC/KplusPi0/";
            path =path + to_string(Run) + magnet +(preselected ? "_preselected" : "") + (truthMatched ? "_TruthMatched" : "")  + "_quickcheck_pi0Resolved.root";
            output = new TFile(path.c_str(),"RECREATE");
        }
    }

    // Kst2Kpluspi0Merged case
    if(Kst2Kpluspi0Merged){
        path = path_to_output_KplusPizero + "/data/";
        if(!MC){
            path =path+ to_string(Run) + magnet + "_" +(preselected ? "preselected" : "") + "_quickcheck_pi0Merged.root";
            output = new TFile(path.c_str(),"RECREATE");
        }
        else{
            if(ReferenceChannel)path =path +"MC/RefKplusPi0";
            else if (PHSP) path =path +"PHSP/KplusPi0/";
            else path =path +"MC/KplusPi0/";
            path =path + to_string(Run) + magnet +(preselected ? "_preselected" : "") + (truthMatched ? "_TruthMatched" : "")  + "_quickcheck_pi0Merged.root";
            output = new TFile(path.c_str(),"RECREATE");
        }
    }
    // Kst2Kspiplus case //TODO: needs a check
    if(Kst2Kspiplus){
        path = path_to_output_KshortPiplus + "/data/";
        if(!MC){
            path =path+ to_string(Run) + magnet + "_" +(preselected ? "preselected" : "") + "_quickcheck_piplus.root";
            output = new TFile(path.c_str(),"RECREATE");
        }
        else{
            if(ReferenceChannel)path =path +"MC/RefKshortPiplus/";
            else if (PHSP) path =path +"PHSP/KshortPiplus/";
            else path =path +"MC/KshortPiplus/";
            path =path + to_string(Run) + magnet +(preselected ? "_preselected" : "") + (truthMatched ? "_TruthMatched" : "")  + "_quickcheck_piplus.root";
            output = new TFile(path.c_str(),"RECREATE");
        }
    }

    output->cd();
    // Stupid arrays :)
    /*
     TCanvas *c1 = new TCanvas("c1", "", 10,10,800,800);
     c1->cd();
     NewTree->Draw("B_plus_DTF_Kst_892_plus_M[0]");
     output->cd();
     c1->Write();
    */
    NewTree->Write();
    output->Close();

    cout << "[INFO]\t\t New tree created." << endl;
    return 1;
}



int EpicOptimizerDataSeparate(string magnet = "down", bool preselected=false) {
    if ( EpicOptimizer( "2011" ,magnet,0, false, false,false,preselected,false) == 0) return 0;
    if ( EpicOptimizer( "2012" ,magnet,0, false, false,false,preselected,false) == 0) return 0;
    if ( EpicOptimizer( "2015" ,magnet,0, false, false,false,preselected,false) == 0) return 0;
    if ( EpicOptimizer( "2016" ,magnet,0, false, false,false,preselected,false) == 0) return 0;
    return 1;
}

int EpicOptimizerAllDataSeparate(bool preselected) {
    if ( EpicOptimizerDataSeparate("down",preselected) == 0) return 0;
    if ( EpicOptimizerDataSeparate("up",preselected) == 0) return 0;
    return 1;
}

int EpicOptimizerMCSeparate(string magnet = "down", bool preselected=false, bool truthMatched=false) {
    if ( EpicOptimizer( "2011" ,magnet,0, true, false,false,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2012" ,magnet,0, true, false,false,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2015" ,magnet,0, true, false,false,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2016" ,magnet,0, true, false,false,preselected,truthMatched) == 0) return 0;
    return 1;
}

int EpicOptimizerAllMCSeparate(bool preselected=false, bool truthMatched=false) {
    if ( EpicOptimizerMCSeparate("down",preselected, truthMatched) == 0) return 0;
    if ( EpicOptimizerMCSeparate("up",  preselected, truthMatched) == 0) return 0;
    return 1;
}


int EpicOptimizerRefSeparate(string magnet = "down", bool preselected=false, bool truthMatched=false) {
    if ( EpicOptimizer( "2011" ,magnet,0, true, true,false,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2012" ,magnet,0, true, true,false,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2015" ,magnet,0, true, true,false,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2016" ,magnet,0, true, true,false,preselected,truthMatched) == 0) return 0;
    return 1;
}

int EpicOptimizerAllRefSeparate(bool preselected=false, bool truthMatched=false) {
    if ( EpicOptimizerRefSeparate("down",preselected, truthMatched) == 0) return 0;
    if ( EpicOptimizerRefSeparate("up",  preselected, truthMatched) == 0) return 0;
    return 1;
}


int EpicOptimizerPHSPSeparate(string magnet = "down", bool preselected=false, bool truthMatched=false) {
    if ( EpicOptimizer( "2011" ,magnet,0, true, false,true,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2012" ,magnet,0, true, false,true,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2015" ,magnet,0, true, false,true,preselected,truthMatched) == 0) return 0;
    if ( EpicOptimizer( "2016" ,magnet,0, true, false,true,preselected,truthMatched) == 0) return 0;
    return 1;
}

int EpicOptimizerAllPHSPSeparate(bool preselected=false, bool truthMatched=false) {
    if ( EpicOptimizerPHSPSeparate("down",preselected, truthMatched) == 0) return 0;
    if ( EpicOptimizerPHSPSeparate("up",  preselected, truthMatched) == 0) return 0;
    return 1;
}

