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

using namespace std::chrono;

auto createTLV = [](double Pt, double Eta, double Phi, double E)
{
  TLorentzVector v;
  v.SetPtEtaPhiE(Pt, Eta, Phi, E);
  return v;
};

auto goodEntry = [](int n_jet, double metTST_pt){
  return n_jet > 2 && metTST_pt > 0.5;
};

auto getMJJ = [](TLorentzVector leadJet, TLorentzVector subleadJet){ 
  return (leadJet + subleadJet).M();
};

void convertTree(const char *inputFile, const char *inputTreeName, const char *outputFileName, const char *outputTreeName){
  
  // Create RDF 
  auto rdf = ROOT::RDataFrame(inputTreeName, inputFile);
  
  // Apply filters
  auto filtered = rdf.Filter(goodEntry, {"n_jet", "metTST_pt"});
  auto updatedDF = filtered.Define("leadJet", createTLV, {"jet_lead_pt", "jet_lead_eta", "jet_lead_phi", "jet_lead_E"})\
                           .Define("subleadJet", createTLV, {"jet_sublead_pt", "jet_sublead_eta", "jet_sublead_phi", "jet_sublead_E"});
  auto finalDF = updatedDF.Define("mJJ", getMJJ, {"leadJet", "subleadJet"});

  // Save file to disk
  finalDF.Snapshot<double>(outputTreeName, outputFileName, {"mJJ"});
}

int main() {
  auto filename = "testSpeed.root";
  auto treename = "output_tree";
  
  auto outputfile = "OutputFileRDFCpp.root";
  auto outputtree = "TMVA_input";

  high_resolution_clock::time_point t1 = high_resolution_clock::now();
  convertTree(filename, treename, outputfile, outputtree);
  high_resolution_clock::time_point t2 = high_resolution_clock::now();
  
  auto duration = duration_cast<milliseconds>( t2 - t1 ).count();
  std::cout << duration / 1000. << " seconds to convert a tree" << std::endl;

  return 0;
}