#include "SimpleInfo.hh"

ROOT::VecOps::RVec<SimpleInfo> getElement(const std::vector<std::vector<SimpleInfo>>& in) {
  return in[0];
}

void SimpleNtuple() {

  gInterpreter->GenerateDictionary("SimpleInfo;std::vector<SimpleInfo>;std::vector<std::vector<SimpleInfo>>;std::vector<std::vector<int>>", "SimpleInfo.hh");
  SimpleInfo single_obj;
  std::vector<SimpleInfo> vector_obj;
  std::vector<std::vector<int>> vector_vector_int;
  std::vector<std::vector<SimpleInfo>> vector_vector_obj;


  TFile* file = new TFile("simple_ntuple.root", "RECREATE");
  TTree* tree = new TTree("ntuple", "");
  int buffsize = 32000;
  int splitlevel = 99;
  tree->Branch("single_obj", &single_obj, buffsize, splitlevel);
  tree->Branch("vector_obj", &vector_obj, buffsize, splitlevel);
  tree->Branch("vector_vector_int", &vector_vector_int, buffsize, splitlevel);
  tree->Branch("vector_vector_obj.", &vector_vector_obj, buffsize, splitlevel);

  int n_events = 10;
  int n_elements = 2;
  int n_subelements = 3;
  for (int i_event = 0; i_event < n_events; ++i_event) {
    vector_obj.clear();

    single_obj.pdg = i_event+1; single_obj.time = 1.1*single_obj.pdg; // dummy pdg and time values

    for (int i_element = 0; i_element < n_elements; ++i_element) {
      SimpleInfo element; element.pdg = -i_event-1; element.time = 1.1*element.pdg - i_element*0.1;
      vector_obj.push_back(element);

      std::vector<SimpleInfo> vector;
      std::vector<int> vector_int;
      for (int i_subelement = 0; i_subelement < n_subelements; ++i_subelement) {
        SimpleInfo subelement; subelement.pdg = -i_event-10; subelement.time = 10.1*subelement.pdg - i_subelement*1.1;
        vector.push_back(subelement);

        vector_int.push_back(subelement.pdg);
      }
      vector_vector_obj.push_back(vector);
      vector_vector_int.push_back(vector_int);
    }
    tree->Fill();
  }
  file->Write();
  file->Close();


  TFile* infile = new TFile("simple_ntuple.root", "READ");

  ROOT::RDataFrame df("ntuple", infile);

  // single-object branches
  auto colType = df.GetColumnType("single_obj");
  std::cout << "single_obj: column type " << colType << std::endl;
  auto disp = df.Display({"single_obj", "single_obj.pdg", "single_obj.time"});
  disp->Print();  // OK

  // vector-object branches
  auto colType2 = df.GetColumnType("vector_obj");
  std::cout << "vector_obj: column type " << colType2 << std::endl;
  auto disp2 = df.Display({"vector_obj", "vector_obj.pdg", "vector_obj.time"});
  disp2->Print();  // OK

  // vector-vector-int branch
  auto colType3 = df.GetColumnType("vector_vector_int");
  std::cout << "vector_vector_int: column type " << colType3 << std::endl;
  auto disp3 = df.Display({"vector_vector_int"});//, "vector_vector_int.pdg", "vector_vector_int.time"});
  disp3->Print();  // OK

  // vector-vector-object
  auto colType4 = df.GetColumnType("vector_vector_obj");
  std::cout << "vector_vector_obj: column type " << colType4 << std::endl;
  auto disp4 = df.Display({"vector_vector_obj"});
  disp4->Print();  // not OK

  auto df2 = df.Define("vector_vector_obj_0", getElement, {"vector_vector_obj"});
  auto colType5 = df2.GetColumnType("vector_vector_obj_0");
  std::cout << "vector_vector_obj_0: column type " << colType5 << std::endl;
  auto disp5 = df2.Display({"vector_vector_obj", "vector_vector_obj_0.pdg", "vector_vector_obj_0.time"});
  disp5->Print();  // not OK
}