#include <iostream>
#include <chrono>
#include <ROOT/RDataFrame.hxx>

int main(int argc, char* argv[])
{
  uint32_t threads=0;
  if(argc==2)
    ROOT::EnableImplicitMT(std::atoi(argv[1]));
  

  ROOT::RDataFrame r(100000000);
  auto rr = r.Define("v", [](ULong64_t e) {return pow(1./(e+1),2);}, {"rdfentry_"});

  //
  // Write TTree in a subdirectory of TFile
  auto start = std::chrono::high_resolution_clock::now();
  rr.Snapshot("test/testTree","test.root",{"v"});
  auto finish = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> elapsed=finish-start;  
  std::cout << threads << "\t" << elapsed.count() << std::endl;

  //
  // Write TTree in root of TFile
  ROOT::RDF::RSnapshotOptions fopt;
  fopt.fMode="UPDATE";
  start = std::chrono::high_resolution_clock::now();
  rr.Snapshot("testTree","test.root",{"v"},fopt);
  finish = std::chrono::high_resolution_clock::now();
  elapsed=finish-start;  
  std::cout << threads << "\t" << elapsed.count() << std::endl;

  //
  // Print results
  TFile *f=TFile::Open("test.root");
  std::cout << "test/testTree entries = " << f->Get<TTree>("test/testTree")->GetEntries() << std::endl;
  std::cout << "testTree entries = " << f->Get<TTree>("testTree")->GetEntries() << std::endl;  
  f->Close();
}