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

#include "dd4pod/Geant4ParticleCollection.h"

//namespace edm4hep {
//
//std::vector<float> pt (std::vector<MCParticleData> const& in){
//  std::vector<float> result;
//  for (size_t i = 0; i < in.size(); ++i) {
//    result.push_back(std::sqrt(in[i].momentum.x * in[i].momentum.x + in[i].momentum.y * in[i].momentum.y));
//  }
//  return result;
//}
//
//std::vector<float> eta(std::vector<MCParticleData> const& in){
//  std::vector<float> result;
//  ROOT::Math::PxPyPzMVector lv;
//  for (size_t i = 0; i < in.size(); ++i) {
//    lv.SetCoordinates(in[i].momentum.x, in[i].momentum.y, in[i].momentum.z, in[i].mass);
//    result.push_back(lv.Eta());
//  }
//  return result;
//}
//
//std::vector<float> cos_theta(std::vector<MCParticleData> const& in){
//  std::vector<float> result;
//  ROOT::Math::PxPyPzMVector lv;
//  for (size_t i = 0; i < in.size(); ++i) {
//    lv.SetCoordinates(in[i].momentum.x, in[i].momentum.y, in[i].momentum.z, in[i].mass);
//    result.push_back(cos(lv.Theta()));
//  }
//  return result;
//}
//
//}


std::vector<float> pt (std::vector<dd4pod::Geant4ParticleData> const& in){
  std::vector<float> result;
  for (size_t i = 0; i < in.size(); ++i) {
    result.push_back(std::sqrt(in[i].psx * in[i].psx + in[i].psy * in[i].psy));
  }
  return result;
}

int rdf_test() {

  ROOT::EnableImplicitMT();

  ROOT::RDataFrame df("events", "sim_barrel_clusters.root");

  auto df2 =
      df.Define("MCParticles_pt", pt, {"mcparticles"});
          //.Define("MCParticles_eta", edm4hep::eta, {"mcparticles"})
          //.Define("MCParticles_cosTheta", edm4hep::cos_theta, {"mcparticles"});

  std::string outfilename = "rdf_test.root";

  df2.Snapshot("events", outfilename,
               {
                   "MCParticles_pt",
                   "mcparticles"
               });
  return 0;
}