#include <iostream>
#include <vector>
#include <ROOT/RDataFrame.hxx>
#include <ROOT/RDFHelpers.hxx>
#include <ROOT/RResultPtr.hxx>
#include <TH2D.h>

using ROOT::RDF::RNode;

template <typename T>
struct My2DObservablestruct {
  T values;
  ROOT::VecOps::RVec<int> categories;
};

template <typename T>
My2DObservablestruct<T> Build2DObservable(const T& values, const ROOT::VecOps::RVec<int> &categories) {
  return My2DObservablestruct<T>{values, categories};
};


template <typename T>
void fillscalar(TH2D &h, const My2DObservablestruct<T> &c) {
  for (auto cat : c.categories)
    {
      h.Fill(c.values,1.);
    }
};

// vector fill
template <typename T>
void fillvector(TH2D &h, const My2DObservablestruct<T> &c) {
  for (auto value : c.values)
	  {
	    for (auto cat : c.categories)
		    {
		      h.Fill(value, cat,1.);
		    }
	  }
};

auto add_2Dhistos = [] (TH2D h1, TH2D h2) { h1.Add(&h2); return h1;};

int main(int argc, char **argv)
{
  
  ROOT::RDataFrame df("tree2","testtree2.root");
  //this doesn't work

  //vecint is vec<int> -> categories
  //vecfloat - x1
  //float x2
  //aggregate is fill(x,categories)
  //
  ROOT::RDF::TH2DModel *m1 = new ROOT::RDF::TH2DModel("m1","m1",10,-5,5,10,-5,5);
  ROOT::RDF::TH2DModel *m2 = new ROOT::RDF::TH2DModel("m2","m2",10,-5,5,10,-5,5);
  std::vector<ROOT::RDF::TH2DModel*> models;
  models.push_back(m1);
  models.push_back(m2);
  std::vector<std::string> variables;
  variables.push_back("float"); //x2
  variables.push_back("vecfloat"); //x1
  std::vector<ROOT::RDF::RResultPtr<TH2D>> final2Dhistos;
  float a=5;
  auto df2 = df.Define("var_a", [a](){return a;});
 
  auto h1 = df2.Histo1D("var_a");
  auto latestDF = std::make_unique<RNode>(df2);
  //latestDF = std::make_unique<RNode>(latestDF->Define("r42",[](){return 42;}));
  
  for (int i=0; i<variables.size(); i++)
    {
      ROOT::RDF::RResultPtr<TH2D> final2Dhisto;
      std::cout <<i <<" variable: "<<variables[i]<<" Type: "<<latestDF->GetColumnType(variables[i])<<std::endl;
      std::string newname="v_"+variables[i];
      if (latestDF->GetColumnType(variables[i])=="Float_t")
	{
	  std::cout <<"Define v0 - float"<<std::endl;
	  final2Dhisto = std::make_unique<RNode>(latestDF->Define(newname,Build2DObservable<float>,{variables[0],"vecint"}))->Aggregate(fillscalar<float>,add_2Dhistos,newname,*(m1->GetHistogram()));
	}
      else if (latestDF->GetColumnType(variables[i])=="ROOT::VecOps::RVec<float>")
	{
	  std::cout <<"Define v1 - vector<float>"<<std::endl;
	  final2Dhisto = std::make_unique<RNode>(latestDF->Define(newname,Build2DObservable<ROOT::VecOps::RVec<float>>,{variables[1],"vecint"}))->Aggregate(fillvector<ROOT::VecOps::RVec<float>>,add_2Dhistos,newname,*(m2->GetHistogram()));
	}
      if (final2Dhisto) final2Dhistos.push_back(final2Dhisto);
    }

  TFile *fout = TFile::Open("output.root","RECREATE");
  h1->Write();
  int index =0;
  for (auto histo : final2Dhistos)
    {
      histo->Write();
      for (int i=1; i<histo->GetNbinsY(); i++)
	{
	  std::string name="proj_";
	  name+=std::to_string(index);
	  name+="_";
	  name+=std::to_string(i);
	  TH1D *h = (TH1D*) histo->ProjectionX(name.c_str(),i,i,"e");
	  h->Write();
	}
      index++;
    }
 
  fout->Close();

  ROOT::RDF::SaveGraph(df,"graph1.dot");
  return 0;
};