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

#include <algorithm> // accumulate
#include <vector>
#include <string>

void make_file_if_needed() {
    const auto fname = "read_bench.root";
    if (!gSystem->AccessPathName(fname)) // false if file is there
        return;
        
    ROOT::RDataFrame(1000000)
        .Define("x", [] { return float(rand()) / RAND_MAX; })
        .Snapshot("tree", fname);
}

std::vector<float> branch_getentry() {
    TFile f("read_bench.root");
    TTree* t = f.Get<TTree>("tree");
    TBranch* b = t->GetBranch("x");
    float x;
    b->SetAddress(&x);
    const auto n_entries = t->GetEntries();
    std::vector<float> v;
    v.reserve(n_entries);
    for (auto i = 0ll; i < n_entries; ++i) {
        b->GetEntry(i);
        v.push_back(x);
    }
    return v;
}

std::vector<float> tree_getentry() {
    TFile f("read_bench.root");
    TTree* t = f.Get<TTree>("tree");
    t->SetBranchStatus("*", 0);
    t->SetBranchStatus("x", 1);
    float x;
    t->SetBranchAddress("x", &x);
    const auto n_entries = t->GetEntries();
    std::vector<float> v;
    v.reserve(n_entries);
    for (auto i = 0ll; i < n_entries; ++i) {
        t->GetEntry(i);
        v.push_back(x);
    }
    return v;
}

std::vector<float> treereader() {
    TFile f("read_bench.root");
    TTreeReader r("tree", &f);
    std::vector<float> v;
    v.reserve(r.GetEntries());
    TTreeReaderValue<float> rv(r, "x");
    auto i = 0ull;
    while (r.Next()) {
        v.push_back(*rv);
        ++i;
    }
    return v;
}

std::vector<float> rdf_foreach() {
    const auto n_entries = TFile("read_bench.root").Get<TTree>("tree")->GetEntries();
    std::vector<float> v;
    v.reserve(n_entries);
    ROOT::RDataFrame("tree", "read_bench.root")
        .Foreach([&v](float x) { v.push_back(x); }, {"x"});
    return v;
}

std::vector<float> rdf_take() {
    return ROOT::RDataFrame("tree", "read_bench.root").Take<float>("x").GetValue();
}

void check_res(const std::vector<float>& v) {
#ifdef DEBUG
    const auto mean = std::accumulate(v.begin(), v.end(), 0.f) / v.size();
    std::cout << "mean: " << mean << std::endl;
#else
    (void)v;
#endif
}

void bench(const std::string& name, std::vector<float>(*f)(void))
{
    TStopwatch s;
    s.Start();
    for (int i = 0; i < 9; ++i)
        f();
    const auto v = f();
    s.Stop();
    std::cout << name << " " << s.RealTime() << std::endl;
    check_res(v);
}

int main() {
    make_file_if_needed();

    // read file once to warm filesystem cache
    branch_getentry();

    bench("branch\t\t", branch_getentry);
    bench("tree\t\t", tree_getentry);
    bench("treereader\t", treereader);
    bench("take\t\t", rdf_take);
    bench("foreach\t\t", rdf_foreach);

    return 0;
}