#include <ROOT/RDataFrame.hxx>
#include <TFile.h>
#include <TTree.h>

void total_cutBased() {
    // Create a ROOT::RDataFrame with the "Events" TTree from the input ROOT file
    ROOT::RDataFrame d("Events", "/Users/fljs825/JScode/root_projects/projects_2023/RPV/100000/26386219-82B4-E447-A286-680A5F6E8237.root");

    // Create a new column "selected_Electron_cutBased" containing only elements >= 3
    auto filtered = d.Define("selected_Electron_cutBased", [](const ROOT::RVec<int>& cutBased) {
        ROOT::RVec<int> selected;
        for (int value : cutBased) {
            if (value >= 3) {
                selected.push_back(value);
            }
        }
        return selected;
    }, {"Electron_cutBased"})
    .Filter([](const ROOT::RVec<int>& selected_Electron_cutBased) {
        // Keep the row if "selected_Electron_cutBased" contains at least one element
        return selected_Electron_cutBased.size() > 0;
    }, {"selected_Electron_cutBased"})
    .Snapshot("filtered_total", "filtered_total.root");

    // Load the previously saved snapshot
    ROOT::RDataFrame df("filtered_total", "filtered_total.root");

    // Filter out rows with empty "selected_Electron_cutBased"
    auto filtered_df = df.Filter([](const ROOT::RVec<int>& selected_Electron_cutBased) {
        return selected_Electron_cutBased.size() > 0;
    }, {"selected_Electron_cutBased"});

    // Get the list of all branches' names in the TTree
    auto branchNames = filtered_df.GetColumnNames();

    // Remove the branches you want to exclude
    branchNames.erase(std::remove(branchNames.begin(), branchNames.end(), "Electron_cutBased"), branchNames.end());

    // Create a new TTree with only the selected branches
    TFile* outputFile = TFile::Open("total.root", "RECREATE");
    filtered_df.Snapshot("Events", "total_cutBased.root", branchNames);

    // Save the TTree to the output file
    outputFile->Write();
    outputFile->Close();
}