//Tools to add kinematic and TRUE kinematic variables to tuples
//Renata Kopecna

#include "../GlobalFunctions.hh"
#include "../Paths.hpp"
#include "../Utils.hpp"

bool bkgSample = false;
std::string tmpAdress(string year, string magnet){
    //return "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/B0toKstJpsi/"+year+magnet+"/"+year+magnet+"_pi0Resolved.root"; //uncomment for B0->KstJpsi sample

    //return "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/B0toKstMuMu/"+year+magnet+"/"+year+magnet+"_pi0Resolved.root"; //uncomment for B0->KstMuMu sample
    //return "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoK1Jpsi/"+year+magnet+"/"+year+magnet+"_pi0Resolved.root"; //uncomment for B->K1Jpsi sample

    return "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoXJpsi/"+year+magnet+"/"+year+magnet+"_pi0Resolved.root";
    //return "/home/lhcb/kopecna/B2KstarMuMu/data/data/MC/BackgroundSamples/BtoK1MuMu/"+year+magnet+"/"+year+magnet+"_pi0Resolved.root"; //uncomment for B->K1Jpsi sample
}

int addPi0TRUEVariables(string year, int Run, string magnet, bool BDTed, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo){
    //sanity checks
    if (magnet == "Down")   magnet = "down";
    if (magnet == "Up")     magnet = "up";
    if (magnet == "Both" || magnet == "")   magnet = "both";

    if (magnet != "both" && magnet != "up" && magnet != "down"){
        coutWarning("Wrong magnet polarity used! Setting magnet to both!");
        magnet = "both";
    }
    if (magnet != "both" && BDTed){
        coutWarning("BDT is applied only per yer, not per polarity! Setting magnet to both!");
        magnet = "both";
    }
    if (!Kst2Kspiplus && KshortDecayInVelo){
        coutERROR("No LL/DD tracks in KplusPizero channel! Setting KshortDecaysInVelo to false!");
        KshortDecayInVelo = false;
    }

    gROOT->SetBatch(kTRUE);
    LHCbStyle();

    string treeAddress = returnFileAddress(year,Run,magnet,true,BDTed,true,ReferenceChannel,PHSP,KshortDecayInVelo);
    if (bkgSample) treeAddress = tmpAdress(year,magnet);
    if (BDTed) replace(treeAddress, "BDToutputSelection", "BDToutput");
    TFile *file = TFile::Open(treeAddress.c_str());
    if(file == 0){
        coutERROR("No file called " + treeAddress+ ". Exit program!");
        return 0;
    }
    coutInfo("Reading data from "+treeAddress);

    TTree *tree = (TTree*)file->Get("DecayTreeTruthMatched");

    int nOldEvts = tree->GetEntries();

    if(nOldEvts == 0){
        coutERROR("No entries found in TTree. Exit program!");
        return 0;
    }

    coutDebug("Old Tree entries:\t" + to_string(nOldEvts));

    //Save //NOW in a new file
    replace(treeAddress,".root","_test.root");
    coutDebug("Creating new file " + treeAddress);

    TFile *output = new TFile(treeAddress.c_str(),"RECREATE");
    TTree * newTree = tree->CloneTree();

    //Create new branches
    string particle = "pi_zero_resolved";

    TBranch *trueMomentum;
    double trueP = 0.0;
    string branchName = particle + "_TRUEP_NEW";
    string branchNameType =branchName+ "/D";
    trueMomentum = newTree->Branch(branchName.c_str(), &trueP, branchNameType.c_str() );

    TBranch *trueMomentumX;
    double truePX = 0.0;
    branchName = particle + "_TRUEP_X_NEW";
    branchNameType =branchName+ "/D";
    trueMomentumX = newTree->Branch(branchName.c_str(), &truePX, branchNameType.c_str() );

    TBranch *trueMomentumY;
    double truePY = 0.0;
    branchName = particle + "_TRUEP_Y_NEW";
    branchNameType =branchName+ "/D";
    trueMomentumY = newTree->Branch(branchName.c_str(), &truePY, branchNameType.c_str() );

    TBranch *trueMomentumZ;
    double truePZ = 0.0;
    branchName = particle + "_TRUEP_Z_NEW";
    branchNameType =branchName+ "/D";
    trueMomentumZ = newTree->Branch(branchName.c_str(), &truePZ, branchNameType.c_str() );

    TBranch *trueMomentumT;
    double truePT = 0.0;
    branchName = particle + "_TRUEPT_NEW";
    branchNameType =branchName+ "/D";
    trueMomentumT = newTree->Branch(branchName.c_str(), &truePT, branchNameType.c_str() );

    TBranch *trueMomentumE;
    double truePE = 0.0;
    branchName = particle + "_TRUEP_E_NEW";
    branchNameType =branchName+ "/D";
    trueMomentumE = newTree->Branch(branchName.c_str(), &truePE, branchNameType.c_str() );

    TBranch *trueEta;
    double trueETA = 0.0;
    branchName = particle + "_TRUEETA_NEW";
    branchNameType =branchName+ "/D";
    trueEta= newTree->Branch(branchName.c_str(), &trueETA, branchNameType.c_str() );

    TBranch *truePhi ;
    double truePHI = 0.0;
    branchName = particle + "_TRUEPHI_NEW";
    branchNameType =branchName+ "/D";
    truePhi = newTree->Branch(branchName.c_str(), &truePHI, branchNameType.c_str() );

    TBranch *trueMass;
    double trueM = 0.0;
    branchName = particle + "_TRUEM_NEW";
    branchNameType =branchName+ "/D";
    trueMass = newTree->Branch(branchName.c_str(), &trueM, branchNameType.c_str() );

    //Create a lorentz vector and load some branches form the old tree
    TLorentzVector LorVec[5];
    TLorentzVector LorVecTrue;
    TLorentzVector LorTmp,LorVecDiMuon;

    Double_t PX[5];
    Double_t PY[5];
    Double_t PZ[5];
    Double_t PT[5];
    Double_t PE[5];

    string particles[5] = {"B_plus","K_star_plus","K_plus","mu_plus","mu_minus"};

    for (int i = 0; i < 5; i++){
        newTree->SetBranchAddress( Form("%s_TRUEP_X",particles[i].c_str()), &PX[i]);
        newTree->SetBranchAddress( Form("%s_TRUEP_Y",particles[i].c_str()), &PY[i]);
        newTree->SetBranchAddress( Form("%s_TRUEP_Z",particles[i].c_str()), &PZ[i]);
        newTree->SetBranchAddress( Form("%s_TRUEPT",particles[i].c_str()), &PT[i]);
        newTree->SetBranchAddress( Form("%s_TRUEP_E",particles[i].c_str()), &PE[i]);
    }

    //Loop over the old tree
    coutInfo("Starting the loop!");
    for (int i = 0; i < nOldEvts; i++){
        newTree->GetEntry(i);
        if (0ul == (i % 10000ul) || nOldEvts == i + 1) coutDebug("Read event " + to_string(i) + "/" + to_string(nOldEvts));

        //Get LorVec
        for (int i = 0; i < 5; i++){
            LorVec[i].SetPxPyPzE(PX[i],PY[i],PZ[i],PE[i]);
        }
        LorVecDiMuon = LorVec[3]+LorVec[4];
        LorTmp = LorVec[0] - LorVecDiMuon;
        LorVecTrue = LorTmp - LorVec[2];

        //Calculate the momenta and eta and phi
        trueP  = LorVecTrue.P();
        truePX = LorVecTrue.Px();
        truePY = LorVecTrue.Py();
        truePZ = LorVecTrue.Pz();
        truePT = LorVecTrue.Pt();
        truePE = LorVecTrue.E();
        trueM  = LorVecTrue.M();
        truePHI = LorVecTrue.Phi();
        trueETA = LorVecTrue.Eta();
        //Fill
        trueMomentum->Fill();
        trueMomentumX->Fill();
        trueMomentumY->Fill();
        trueMomentumZ->Fill();
        trueMomentumT->Fill();
        trueMomentumE->Fill();
        trueEta->Fill();
        truePhi->Fill();
        trueMass->Fill();
    }

    if (nOldEvts != newTree->GetEntries()){
        coutERROR("Something went wrong with filling the tree! Exit without saving the new tree.");
        return 0;
    }

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

    output->cd();
    coutDebug("Writting into the new file " + treeAddress );
    newTree->Write("",TObject::kWriteDelete);

    //tree->Clear();
    output->Close();
    file->Close();

    string command = "mv " + treeAddress + " ";
    replace(treeAddress,"_test.root",".root");

    command = "rm " + treeAddress + "; " + command + treeAddress;
    coutDebug("Running: " + command);
    system(command.c_str());

    return 1;

}

int addPi0TRUEVariablesAllYearMC(string year, bool ReferenceChannel, bool PHSP){

    if (addPi0TRUEVariables(year,getRunID(year),"down",false,ReferenceChannel,PHSP,false)==0) return 0;
    if (addPi0TRUEVariables(year,getRunID(year),"up",  false,ReferenceChannel,PHSP,false)==0) return 0;
   // if (addPi0TRUEVariables(year,getRunID(year),"both", false,ReferenceChannel,PHSP,false)==0) return 0;
   // if (addPi0TRUEVariables(year,getRunID(year),"both", true, ReferenceChannel,PHSP,false)==0) return 0;
    return 1;
}

int addPi0TRUEVariablesAllMC(int Run, bool ReferenceChannel, bool PHSP){

    for (auto &year: yearsData(Run)){
        addPi0TRUEVariablesAllYearMC(year,ReferenceChannel,PHSP);
    }
    return 1;
}

int addPi0TRUEVariablesAllMC(int Run){

    if (addPi0TRUEVariablesAllMC(Run,false,false)==0) return 0;
    if (addPi0TRUEVariablesAllMC(Run,true, false)==0) return 0;
    if (addPi0TRUEVariablesAllMC(Run,false,true )==0) return 0;
    return 1;
}

int addVariable(string year, int Run, string magnet, bool Preselected, bool BDTed, bool TM, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo,
                string particle){
    //sanity checks
    if (magnet == "Down")   magnet = "down";
    if (magnet == "Up")     magnet = "up";
    if (magnet == "Both" || magnet == "")   magnet = "both";

    if (magnet != "both" && magnet != "up" && magnet != "down"){
        coutWarning("Wrong magnet polarity used! Setting magnet to both!");
        magnet = "both";
    }
    if (magnet != "both" && BDTed){
        coutWarning("BDT is applied only per yer, not per polarity! Setting magnet to both!");
        magnet = "both";
    }
    if (TM && !Preselected){
        coutWarning("Stripped MC is not TruthMatched! Setting TM to false!");
        TM = false;
    }
    if (!Kst2Kspiplus && KshortDecayInVelo){
        coutERROR("No LL/DD tracks in KplusPizero channel! Setting KshortDecaysInVelo to false!");
        KshortDecayInVelo = false;
    }
    if (!BDTed && !Preselected){ //if TMVAcut == -1, no TMVA cut is applied
        coutWarning("TMVA cut can be only aplied on preselected data! Setting Preselected to true!");
        Preselected = true; //Cannot do BDT cut on stripped data
    }

    string treeAddress = returnFileAddress(year,Run,magnet,Preselected,BDTed,true,ReferenceChannel,PHSP,KshortDecayInVelo);
    if (bkgSample) treeAddress = tmpAdress(year,magnet);
    if (BDTed) replace(treeAddress, "BDToutputSelection", "BDToutput");
    TFile *file = TFile::Open(treeAddress.c_str());
    if(file == 0){
        coutERROR("No file called " + treeAddress+ ". Exit program!");
        return 0;
    }
    coutInfo("Reading data from TTree... "+treeAddress);

    TTree *tree;
    if (Preselected) tree = (TTree*)file->Get("DecayTreeTruthMatched");
    else tree = (TTree*)file->Get("DecayTree");


    int nOldEvts = tree->GetEntries();

    if(nOldEvts == 0){
        coutERROR("No entries found in TTree. Exit program!");
        return 0;
    }

    coutDebug("Old Tree entries:\t" + to_string(nOldEvts));

    //Save //NOW in a new file
    replace(treeAddress,".root","_test.root");
    coutDebug("Creating new file " + treeAddress);

    TFile *output = new TFile(treeAddress.c_str(),"RECREATE");
    TTree * newTree = tree->CloneTree();


    //Create new branches
    string branchName = "";
    string branchNameType = "";

    //Create new branches for MC
    TBranch *trueMass;
    TBranch *trueMomentum;
    TBranch *trueEta;
    TBranch *truePhi ;

    double trueM = 0.0;
    double trueP = 0.0;
    double trueETA = 0.0;
    double truePHI = 0.0;

        branchName = particle + "_TRUEM";
        branchNameType =branchName+ "/D";
        trueMass = newTree->Branch(branchName.c_str(), &trueM, branchNameType.c_str() );

        branchName = particle + "_TRUEP";
        branchNameType =branchName+ "/D";
        trueMomentum = newTree->Branch(branchName.c_str(), &trueP, branchNameType.c_str() );

        branchName = particle + "_TRUEETA";
        branchNameType =branchName+ "/D";
        trueEta= newTree->Branch(branchName.c_str(), &trueETA, branchNameType.c_str() );

        branchName = particle + "_TRUEPHI";
        branchNameType =branchName+ "/D";
        truePhi = newTree->Branch(branchName.c_str(), &truePHI, branchNameType.c_str() );

    //Create a lorentz vector and load some branches form the old tree
    TLorentzVector LorVecTrue;

    Double_t TRUEP_X;
    Double_t TRUEP_Y;
    Double_t TRUEP_Z;
    Double_t TRUEP_E;

    newTree->SetBranchAddress( Form("%s_TRUEP_X",particle.c_str()), &TRUEP_X );
    newTree->SetBranchAddress( Form("%s_TRUEP_Y",particle.c_str()), &TRUEP_Y );
    newTree->SetBranchAddress( Form("%s_TRUEP_Z",particle.c_str()), &TRUEP_Z );
    newTree->SetBranchAddress( Form("%s_TRUEP_E",particle.c_str()), &TRUEP_E );


    //Loop over the old tree
    coutInfo("Starting the loop!");
    for (int i = 0; i < nOldEvts; i++){
        newTree->GetEntry(i);
        if (0ul == (i % 10000ul) || nOldEvts == i + 1) coutDebug("Read event " + to_string(i) + "/" + to_string(nOldEvts));

        //Get TRUE LorVec
        LorVecTrue.SetPxPyPzE(TRUEP_X,TRUEP_Y,TRUEP_Z,TRUEP_E);

        //Calculate the momenta and eta and phi
        trueM = LorVecTrue.M();
        trueP = LorVecTrue.P();
        trueETA = LorVecTrue.Eta();
        truePHI = LorVecTrue.Phi();


        //Fill
        trueMass->Fill();
        trueMomentum->Fill();
        trueEta->Fill();
        truePhi->Fill();

    }

    if (nOldEvts != newTree->GetEntries()){
        coutERROR("Something went wrong with filling the tree! Exit without saving the new tree.");
        return 0;
    }

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

    output->cd();
    coutDebug("Writting into the new file " + treeAddress );
    newTree->Write("",TObject::kWriteDelete);

    output->Close();
    file->Close();

    string command = "mv " + treeAddress + " ";
    replace(treeAddress,"_test.root",".root");

    command = "rm " + treeAddress + "; " + command + treeAddress;
    coutDebug("Running: " + command);
    system(command.c_str());


    return 1;

}

int addAllVariables(string year, int Run, string magnet, bool Preselected, bool BDTed, bool TM, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo){

    vector <string> particles = {"B_plus", "K_star_plus","pi_zero_resolved","K_plus", "gamma1","gamma2","mu_minus","mu_plus"} ;//TODO

    for (auto& particle: particles){
        coutInfo("Running particle " + particle);
        if (addVariable(year,Run,magnet,Preselected,BDTed,TM,ReferenceChannel,PHSP,KshortDecayInVelo,particle)==0) return 0;
    }
    return 1;

}


int addAllVariablesAllYearMC(string year, bool ReferenceChannel, bool PHSP){
    //if (addAllVariables(year,getRunID(year),"down",true,false,false,ReferenceChannel,PHSP,false)==0) return 0;
    //if (addAllVariables(year,getRunID(year),"up",  true,false,false,ReferenceChannel,PHSP,false)==0) return 0;
    //if (addAllVariables(year,getRunID(year),"both",true,false,true,true,ReferenceChannel,PHSP,false)==0) return 0;
    if (addAllVariables(year,getRunID(year),"both",true,true,true,ReferenceChannel,PHSP,false)==0) return 0;
    if (Kst2Kspiplus){
        if (addAllVariables(year,getRunID(year),"down",true,false,false,ReferenceChannel,PHSP,true)==0) return 0;
        if (addAllVariables(year,getRunID(year),"up",  true,false,false,ReferenceChannel,PHSP,true)==0) return 0;
     //   if (addAllVariables(year,getRunID(year),"both",true,false,true,true,ReferenceChannel,PHSP,true)==0) return 0;
        if (addAllVariables(year,getRunID(year),"both",true,true,true,ReferenceChannel,PHSP,true)==0) return 0;
    }
    return 1;
}

int addAllVariablesAllMC(int Run,bool ReferenceChannel, bool PHSP){
    for (auto &year: yearsMC(ReferenceChannel,PHSP,Run)){
        addAllVariablesAllYearMC(year,ReferenceChannel,PHSP);
    }
    return 1;
}

int addAllVariablesAllMCSamples(int Run){
    if (addAllVariablesAllMC(Run,false,false)==0) return 0;
    if (addAllVariablesAllMC(Run,true, false)==0) return 0;
    if (addAllVariablesAllMC(Run,false,true )==0) return 0;
    return 1;
}

int testFunction(){
    setVerboseLevel(0);
    return 1;

}

int addXMuMuMass(bool Kplus, bool DTF, string year, string magnet, int Run, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo){
    //sanity checks

    if (!Kst2Kspiplus && KshortDecayInVelo){
        coutERROR("No LL/DD tracks in KplusPizero channel! Setting KshortDecaysInVelo to false!");
        KshortDecayInVelo = false;
    }

    gROOT->SetBatch(kTRUE);
    LHCbStyle();


    string treeAddress = returnFileAddress(year,Run,magnet,true,false, MC,ReferenceChannel,PHSP,KshortDecayInVelo);
    if (bkgSample) treeAddress = tmpAdress(year,magnet);
    TFile *file = TFile::Open(treeAddress.c_str());
    if(file == 0){
        coutERROR("No file called " + treeAddress+ ". Exit program!");
        return 0;
    }
    coutInfo("Reading data from "+treeAddress);

    TTree *tree = (TTree*)file->Get(treeName(MC,true).c_str());
    int nOldEvts = tree->GetEntries();

    if(nOldEvts == 0){
        coutERROR("No entries found in TTree. Exit program!");
        return 0;
    }

    coutDebug("Old Tree entries:\t" + to_string(nOldEvts));

    //Save //NOW in a new file
    replace(treeAddress,".root","_test.root");
    coutDebug("Creating new file " + treeAddress);

    TFile *output = new TFile(treeAddress.c_str(),"RECREATE");
    TTree * newTree = tree->CloneTree();

    //Create new branches

    TBranch *M_XMuMu;
    double new_M = 0.0;
    string branchName = Kplus ? ("M_KplusMuMu") : (DTF ? "M_PizMuMu" : "M_PizMuMu_noDTF");
    string branchNameType =branchName+ "/D";
    M_XMuMu = newTree->Branch(branchName.c_str(), &new_M, branchNameType.c_str() );


    //Create a lorentz vector and load some branches form the old tree
    const int nParticles = 3;

    TLorentzVector LorVec[nParticles];
    TLorentzVector LorTmp,LorVecDiMuon;

    Double_t PX[nParticles];
    Double_t PY[nParticles];
    Double_t PZ[nParticles];
    Double_t PE[nParticles];


    string particles[nParticles] = {Kplus ? "K_plus" : "pi_zero_resolved","mu_plus","mu_minus"};

    string s_DTF = DTF ? "_DTF" : "";
    for (int i = 0; i < nParticles; i++){
        newTree->SetBranchAddress( Form("%s_PX%s",particles[i].c_str(),s_DTF.c_str()), &PX[i]);
        newTree->SetBranchAddress( Form("%s_PY%s",particles[i].c_str(),s_DTF.c_str()), &PY[i]);
        newTree->SetBranchAddress( Form("%s_PZ%s",particles[i].c_str(),s_DTF.c_str()), &PZ[i]);
        newTree->SetBranchAddress( Form("%s_PE%s",particles[i].c_str(),s_DTF.c_str()), &PE[i]);
    }

    //Loop over the old tree
    coutInfo("Starting the loop!");
    for (int i = 0; i < nOldEvts; i++){
        newTree->GetEntry(i);
        if (0ul == (i % 10000ul) || nOldEvts == i + 1) coutDebug("Read event " + to_string(i) + "/" + to_string(nOldEvts));

        //Get LorVec
        for (int i = 0; i < nParticles; i++){
            LorVec[i].SetPxPyPzE(PX[i],PY[i],PZ[i],PE[i]);
        }
        LorVecDiMuon = LorVec[1]+LorVec[2];
        LorTmp = LorVec[0] + LorVecDiMuon;
        new_M = LorTmp.M();

        //Fill
        M_XMuMu->Fill();
    }

    if (nOldEvts != newTree->GetEntries()){
        coutERROR("Something went wrong with filling the tree! Exit without saving the new tree.");
        return 0;
    }

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

    output->cd();
    coutDebug("Writting into the new file " + treeAddress );
    newTree->Write("",TObject::kWriteDelete);

    //tree->Clear();
    output->Close();
    file->Close();

    string command = "mv " + treeAddress + " ";
    replace(treeAddress,"_test.root",".root");

    command = "rm " + treeAddress + "; " + command + treeAddress;
    coutDebug("Running: " + command);
    system(command.c_str());

    return 1;

}

int addAllXMuMuMass(bool Kplus, bool DTF, int Run, bool MC, bool ReferenceChannel, bool PHSP){
    for (auto &year: yearsMC(ReferenceChannel,PHSP,Run)){
        addXMuMuMass(Kplus, DTF,year, "up",getRunID(year), MC, ReferenceChannel,PHSP,false);
        addXMuMuMass(Kplus, DTF,year, "down",getRunID(year), MC, ReferenceChannel,PHSP,false);
    }
    return 1;
}

int addAllXMuMuMass(bool Kplus, bool DTF,int Run){
    addAllXMuMuMass(Kplus, DTF, Run, false, false, false);
    addAllXMuMuMass(Kplus, DTF, Run, true,  false, false);
    addAllXMuMuMass(Kplus, DTF, Run, true,  true,  false);
    addAllXMuMuMass(Kplus, DTF, Run, true,  false, true);
    return 1;
}


int applyVetoKplusMuMuMass(string year, string magnet, int Run, bool MC, bool ReferenceChannel, bool PHSP, bool KshortDecayInVelo){
    //sanity checks

    if (!Kst2Kspiplus && KshortDecayInVelo){
        coutERROR("No LL/DD tracks in KplusPizero channel! Setting KshortDecaysInVelo to false!");
        KshortDecayInVelo = false;
    }
    gROOT->SetBatch(kTRUE);
    LHCbStyle();

    string treeAddress = returnFileAddress(year,Run,magnet,true,false, MC,ReferenceChannel,PHSP,KshortDecayInVelo);
    if (bkgSample) treeAddress = tmpAdress(year,magnet);

    TFile *file = TFile::Open(treeAddress.c_str());
    if(file == 0){
        coutERROR("No file called " + treeAddress+ ". Exit program!");
        return 0;
    }

    coutInfo("Reading data from "+treeAddress);

    TTree *tree = (TTree*)file->Get(treeName(MC,true).c_str());

    //Save //NOW in a new file
    replace(treeAddress,".root","_test.root");
    coutDebug("Creating new file " + treeAddress);

    TFile *output   = new TFile(treeAddress.c_str(),"RECREATE");
    TTree * newTree = tree->CopyTree("M_KplusMuMu<5179 || M_KplusMuMu>5379","", tree->GetEntries());
    output->cd();
    coutDebug("Writting into the new file " + treeAddress );
    newTree->Write("",TObject::kWriteDelete);

    //tree->Clear();
    output->Close();
    file->Close();

    string command = "mv " + treeAddress + " ";
    replace(treeAddress,"_test.root",".root");

    command = "rm " + treeAddress + "; " + command + treeAddress;
    coutDebug("Running: " + command);
    system(command.c_str());

    return 1;
}

int applyAllVetoKplusMuMuMass(int Run, bool MC, bool ReferenceChannel, bool PHSP){
    for (auto &year: yearsMC(ReferenceChannel,PHSP,Run)){
        applyVetoKplusMuMuMass(year, "up",getRunID(year), MC, ReferenceChannel,PHSP,false);
        applyVetoKplusMuMuMass(year, "down",getRunID(year), MC, ReferenceChannel,PHSP,false);
    }
    return 1;
}

int applyAllVetoKplusMuMuMass(int Run){
    applyAllVetoKplusMuMuMass(Run, false, false, false);
    applyAllVetoKplusMuMuMass(Run, true,  false, false);
    applyAllVetoKplusMuMuMass(Run, true,  true,  false);
    applyAllVetoKplusMuMuMass(Run, true,  false, true);
    return 1;
}

int applyVetoPizMuMuMass(string year, string magnet, int Run, bool MC){
    //sanity checks

    gROOT->SetBatch(kTRUE);
    LHCbStyle();

    //MC only do Jpsi region
    string treeAddress = returnFileAddress(year,Run,magnet,true,false,MC,MC,false,false);

    TFile *file = TFile::Open(treeAddress.c_str());
    if(file == 0){
        coutERROR("No file called " + treeAddress+ ". Exit program!");
        return 0;
    }

    coutInfo("Reading data from "+treeAddress);

    TTree *tree = (TTree*)file->Get(treeName(MC,true).c_str());

    //Save //NOW in a new file
    replace(treeAddress,".root","_test.root");
    coutDebug("Creating new file " + treeAddress);

    TFile *output   = new TFile(treeAddress.c_str(),"RECREATE");

    coutDebug("Writting into the new file " + treeAddress );
    if (MC){
        TTree * newTree = tree->CopyTree("M_PizMuMu>3700","", tree->GetEntries());
        newTree->Write("",TObject::kWriteDelete);
    }
    else{
        string cut = "!(" + getJpsicut() + ") || ( M_PizMuMu>3700)";
        TTree *newTree = tree->CopyTree(cut.c_str(),"", tree->GetEntries());
        newTree->Write("",TObject::kWriteDelete);
    }

    //tree->Clear();
    output->Close();
    file->Close();

    string command = "mv " + treeAddress + " ";
    replace(treeAddress,"_test.root",".root");

    command = "rm " + treeAddress + "; " + command + treeAddress;
    coutDebug("Running: " + command);
    system(command.c_str());

    return 1;
}

int applyVetoPizMuMuMass(int Run){
    for (auto &year: yearsData(Run)){
        applyVetoPizMuMuMass(year, "up",getRunID(year), true);
        applyVetoPizMuMuMass(year, "up",getRunID(year), false);
        applyVetoPizMuMuMass(year, "down",getRunID(year), true);
        applyVetoPizMuMuMass(year, "down",getRunID(year), false);
    }
    return 1;
}