#include <iostream>
#include <vector>
#include <unordered_map>
#include <TROOT.h>
#include <ROOT/RDataFrame.hxx>
#include "yaml-cpp/yaml.h"

#include <stdio.h>  // for printf()
#include <stdlib.h> // for strtol()
#include <errno.h>  // for errno
#include <limits.h> // for INT_MIN and INT_MAX
#include <string.h>  // for strlen
#include <TSystem.h>  // for strlen

//Run with ./produce.exe photons
int main(int argc, char **argv) {
  std::string root_folder;
  std::string root_tree;
  root_tree = "HGCalTriggerNtuple";

  if (strlen(argv[1]) == 0)	return 1; // empty std::string

  //process_program_options(argc, argv);
  std::string particle = std::string(argv[1]);

  std::string infile = particle + "_0PU_bc_stc_hadd.root";
  std::string outfile = "skim_small_" + infile;

  std::string reachedEE="", deltarThreshold="", mipThreshold="";
  reachedEE = "2";
  deltarThreshold = "0.05";
  mipThreshold = "0.5";
  std::string vtmp = "tmp_good";


  // dataframe definition
  ROOT::RDataFrame dataframe(root_tree, infile);

  std::vector<std::string> gen_v = {"genpart_posx", "genpart_posy", "genpart_posz"};

  std::unordered_map<std::string,std::string> pmap = {
        {"photons", "22"},
        {"electrons", "11"},
        {"pions", "211"}};
  std::string condgen = "genpart_gen != -1 && ";
  condgen += "genpart_reachedEE == " + reachedEE;
  condgen += " && genpart_pid == abs(" + pmap[particle] + ")";
  condgen += " && genpart_exeta > 0";

  auto df = dataframe.Define(vtmp + "_gens", condgen);
  for(auto& v : gen_v) {  df = df.Define(vtmp + "_" + v, v + "[" + vtmp + "_gens]"); }

  auto dd2 = df;

  for(auto& var : gen_v) {
        dd2 = dd2.Define("good_" + var,
                                         [](const ROOT::VecOps::RVec<std::vector<float>> &v) {
                                           std::vector<std::vector<float>> vec(v.size());
                                           for(unsigned i=0; i<v.size(); ++i) {
                                                 vec[i] = std::vector<float>(v[i].begin(), v[i].end());
                                                 for (auto i : vec[i])
                                                     std::cout << i;
                                           std::cout << std::endl;
                                           }
                                           return vec;
                                         },
                                         {vtmp + "_" + var});
  }

  std::vector<std::string> allvars = gen_v;
  std::vector<std::string> good_allvars = {"event"};
  for(auto& v : allvars)
        good_allvars.push_back("good_" + v);

  dd2.Range(0,5).Snapshot(root_tree, outfile, good_allvars);

  return 0;
}