#include "../GlobalFunctions.hh"

//////////////////////////////////////////////////////////////////////
/// OptimizePreSelection()
/// loops over all .root-files created by Preselection() defined in BDTSelection.cc.
/// More kinematic, PID and some more cuts are applied and saved in a new file.
///


int OptimizePreSelection(std::string year = "2012", std::string magnet = "down", bool MC = true, bool ReferenceChannel = false, bool PHSP = false) {

    if (!MC && PHSP){
        cout << "[INFO]\tIf you want to process PHSP MC, MC needs to be set true as well as PHSP. Abort." << endl;
        return 0;
    }

    TFile* input = 0;
    TFile* output = 0;
    TTree *tree;
    TTree *treeInit; //save also non-truthmatched tree
        // Kst2Kpluspi0Resolved case
    if(Kst2Kpluspi0Resolved){
        if(!MC){
            input = TFile::Open(Form("%s/data/%s%s/%s%s_pi0Resolved_loose.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
            output = new TFile(Form("%s/data/%s%s/%s%s_pi0Resolved_new.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
        }
        else{
            if(ReferenceChannel){
                input = TFile::Open(Form("%s/data/MC/RefKplusPi0/%s%s/%s%s_pi0Resolved_loose.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/MC/RefKplusPi0/%s%s/%s%s_pi0Resolved_new.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
            else if(PHSP){
                input = TFile::Open(Form("%s/data/PHSP/KplusPi0/%s%s/%s%s_pi0Resolved_loose.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/PHSP/KplusPi0/%s%s/%s%s_pi0Resolved_new.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
            else{
                input = TFile::Open(Form("%s/data/MC/KplusPi0/%s%s/%s%s_pi0Resolved_loose.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/MC/KplusPi0/%s%s/%s%s_pi0Resolved_new.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
        }
    }

        // Kst2Kpluspi0Merged case
    if(Kst2Kpluspi0Merged){
        if(!MC){
            input = TFile::Open(Form("%s/data/%s%s/%s%s_pi0Merged.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
            output = new TFile(Form("%s/data/%s%s/%s%s_pi0Merged_new.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");

        }
        else{
            if(ReferenceChannel){
                input = TFile::Open(Form("%s/data/MC/RefKplusPi0/%s%s/%s%s_pi0Merged.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/MC/RefKplusPi0/%s%s/%s%s_pi0Merged_new.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
            else if (PHSP){
                input = TFile::Open(Form("%s/data/PHSP/KplusPi0/%s%s/%s%s_pi0Merged.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/PHSP/KplusPi0/%s%s/%s%s_pi0Merged_new.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
            else{
                input = TFile::Open(Form("%s/data/MC/KplusPi0/%s%s/%s%s_pi0Merged.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/MC/KplusPi0/%s%s/%s%s_pi0Merged_new.root",path_to_output_KplusPizero.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
        }
    }
        // Kst2Kspiplus case
    if(Kst2Kspiplus){
        if (!MC){
            input = TFile::Open(Form("%s/data/%s%s/%s%s_piplus.root",path_to_output_KshortPiplus.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
            output = new TFile(Form("%s/data/%s%s/%s%s_piplus_new.root",path_to_output_KshortPiplus.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
        }
        else{
            if(ReferenceChannel){
                input = TFile::Open(Form("%s/data/MC/RefKshortPiplus/%s%s/%s%s_piplus.root",path_to_output_KshortPiplus.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/MC/RefKshortPiplus/%s%s/%s%s_piplus_new.root",path_to_output_KshortPiplus.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
            else if(PHSP){
                input = TFile::Open(Form("%s/data/PHSP/KshortPiplus/%s%s/%s%s_piplus.root",path_to_output_KshortPiplus.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/PHSP/KshortPiplus/%s%s/%s%s_piplus_new.root",path_to_output_KshortPiplus.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
            else{
                input = TFile::Open(Form("%s/data/MC/KshortPiplus/%s%s/%s%s_piplus.root",path_to_output_KshortPiplus.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()));
                output = new TFile(Form("%s/data/MC/KshortPiplus/%s%s/%s%s_piplus_new.root",path_to_output_KshortPiplus.c_str(), year.c_str(),magnet.c_str(),year.c_str(),magnet.c_str()),"RECREATE");
            }
        }
    }

    string name = (MC ? "DecayTreeTruthMatched" : "DecayTree");
    input->cd();
    tree = (TTree*)input->Get(name.c_str());
    if (MC) treeInit =  (TTree*)input->Get("DecayTree");

    tree->SetBranchStatus("*",1);
    cout << "[INFO]\tOld tuple contains " << tree->GetEntries() << " entries." << endl;

    vector<std::string> Cuts;
    string tmp_string = "";

        // Kst2Kpluspi0Resolved Cuts
    if(Kst2Kpluspi0Resolved){
         //don't copy old cone variables
        //tree->SetBranchStatus("*_1.00",0);
        //tree->SetBranchStatus("*_0.90",0);
        //tree->SetBranchStatus("*_0.80",0);
        //tree->SetBranchStatus("*_0.70",0);
        //tree->SetBranchStatus("*_0.60",0);
        //tree->SetBranchStatus("*_0.50",0);

        Cuts.push_back("B_plus_M_DTF > 5000");
        //Cuts.push_back("nTracks < 200");
        //Cuts.push_back("abs(J_psi_M - 3096) < 80");
        //Cuts.push_back("totCandidates < 5");
        //Cuts.push_back("abs(K_star_plus_M - 892) < 100");
       // Cuts.push_back("K_star_plus_PT_DTF > 1350");
       //// Cuts.push_back("mu_minus_ProbNNmu > 0.25");
        //Cuts.push_back("mu_plus_ProbNNmu > 0.25");

       // tmp_string = (UseDTF ? "K_star_plus_M_DTF > 792" : "K_star_plus_M > 792");
       // Cuts.push_back(tmp_string.c_str());

       // tmp_string = (UseDTF ? "K_star_plus_M_DTF < 1050" : "K_star_plus_M < 1050");
       // Cuts.push_back(tmp_string.c_str());

       // tmp_string = (UseDTF ? "K_star_plus_PT_DTF > 300 " : "K_star_plus_PT > 300");
       // Cuts.push_back(tmp_string.c_str());

       // tmp_string = "abs(" + string(UseDTF ? "K_star_plus_M_DTF" : "K_star_plus_M") +  " - " + to_string(PDGMASS.K_STAR_PLUS) + ") < 150"; //remove non-resonant K+pi
       //cout << tmp_string << endl;
        //Cuts.push_back(tmp_string.c_str());

    }
         // Kst2Kpluspi0Merged cuts
    if(Kst2Kpluspi0Merged){
         //don't copy old cone variables
        tree->SetBranchStatus("*_1.00",0);
        tree->SetBranchStatus("*_0.90",0);
        tree->SetBranchStatus("*_0.80",0);
        tree->SetBranchStatus("*_0.70",0);
        tree->SetBranchStatus("*_0.60",0);
        tree->SetBranchStatus("*_0.50",0);

         //#TODO
    }
        // Kst2Kspiplus cuts
    if(Kst2Kspiplus){
        //#TODO
    }

    std::string AllCuts = "";

    if (Cuts.size()==0){
        cout << "[ERROR]\t No cuts selected, nothing to do." << endl;
        return 0;
    }

    AllCuts = AllCuts + *Cuts.begin();
    cout << "[INFO]\t\t Appling cuts:" << endl;
    cout << "\t\t\t" << AllCuts << endl;
    for (vector<std::string>::iterator t = Cuts.begin()+1;
                                       t != Cuts.end();
                                       ++t){
        AllCuts = AllCuts + " && " + *t ;
        cout << "\t\t\t" << *t << endl;
    }
    Cuts.clear();
    output->cd();
    TTree* NewTree = tree->CopyTree( AllCuts.c_str() );
    TTree* NewTreeInit;
    if (MC) NewTreeInit = treeInit->CopyTree( AllCuts.c_str() );
    output->Write();
    cout << "[INFO]\tNew tuple contains " << NewTree->GetEntries() << " entries." << endl;
    input->Close();
    output->Close();
    gROOT->Reset();

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



int OptimizePreSelectionAllData(Int_t Run = 1) {

    if(Run != 1 && Run != 2){
        std::cout << "[ERROR]\t\tInvalid Run number given: " << Run  << ". Exit program!" << std::endl;
        return 0;
    }

    if ( OptimizePreSelection((Run == 1 ? "2011" : "2015"),"down", false, false, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2011" : "2015"),"up"  , false, false, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2012" : "2016"),"down",false, false, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2012" : "2016"),"up"  ,false, false, false) == 0) return 0;
    return 1;

}

int OptimizePreSelectionAllMC(Int_t Run = 1) {

    if(Run != 1 && Run != 2){
        std::cout << "[ERROR]\t\tInvalid Run number given: " << Run  << ". Exit program!" << std::endl;
        return 0;
    }

    if ( OptimizePreSelection((Run == 1 ? "2011" : "2015"),"down",true, false, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2011" : "2015"),"up"  ,true, false, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2012" : "2016"),"down",true, false, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2012" : "2016"),"up"  ,true, false, false) == 0) return 0;
    return 1;

}

int OptimizePreSelectionAllRefChannel(Int_t Run = 1) {

    if(Run != 1){
        std::cout << "[ERROR]\t\tInvalid Run number given: " << Run  << "." << std::endl;
        std::cout << "[ERROR]\t\tInvalid RefChannel available only for Run 1 at the moment! Exit program!" << std::endl;
        return 0;
    }
    if ( OptimizePreSelection((Run == 1 ? "2011" : "2015"),"down",true, true, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2011" : "2015"),"up"  ,true, true, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2012" : "2016"),"down",true, true, false) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2012" : "2016"),"up"  ,true, true, false) == 0) return 0;
    return 1;

}


int OptimizePreSelectionAllPHSP(Int_t Run = 1) {

    if(Run != 1 && Run != 2){
        std::cout << "[ERROR]\t\tInvalid Run number given: " << Run  << ". Exit program!" << std::endl;
        return 0;
    }
    if ( OptimizePreSelection((Run == 1 ? "2011" : "2015"),"down",true, false, true) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2011" : "2015"),"up"  ,true, false, true) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2012" : "2016"),"down",true, false, true) == 0) return 0;
    if ( OptimizePreSelection((Run == 1 ? "2012" : "2016"),"up"  ,true, false, true) == 0) return 0;
    return 1;

}

int OptimizePreSelectionAll( Int_t Run = 1) {

    if (OptimizePreSelectionAllData(Run)==0) return 0;
    if (OptimizePreSelectionAllMC(Run)==0) return 0;
    if (OptimizePreSelectionAllPHSP(Run)==0) return 0;
    if (OptimizePreSelectionAllRefChannel(Run)==0) return 0;
    return 1;

}