#include <ROOT/RDataFrame.hxx>
#include <iostream>
#include <TFile.h>
#include <TTree.h>
#include <vector>
#include <string>
#include <TCanvas.h>

void generateData(const char *filename, int n, double value)
{

   TFile f(filename, "RECREATE", "file for testing");

   double var1;

   TTree *tree = new TTree("AnalysisTree", "AnalysisTree");
   auto branch1 = tree->Branch("column1", &var1, "column1/D");

   for (unsigned int i = 0; i < n; i++) {
      // var1 = rand() %100;
      var1 = value;
      tree->Fill();
   }

   tree->Write();
   f.Write();
   f.Close();
}

int main()
{

   std::vector<std::string> fileNames;
   std::vector<double> weights;

   // Create two files with trees for testing

   // 10 entries of 0.5
   generateData("test1.root", 10, 0.5);
   fileNames.push_back("test1.root");
   weights.push_back(2);

   // 10 entries of 2
   generateData("test2.root", 10, 2);
   fileNames.push_back("test2.root");
   weights.push_back(0.5);

   // Create Dataframe from files

   ROOT::RDataFrame df("AnalysisTree", fileNames);
   ROOT::RDF::RNode df2 = ROOT::RDataFrame(0);

   // Define weights depending on input file

   df2 = df.DefinePerSample("weightbysample", [&fileNames, &weights](unsigned int, const ROOT::RDF::RSampleInfo &id) {
      for (unsigned int i = 0; i < fileNames.size(); i++)
         if (id.Contains(fileNames[i]))
            return weights[i];
      return -1.;
   });

   // Book all operations without triggering the execution
   auto sum0 = df2.Sum<double>("weightbysample");
   df2 = df2.Define("calibrated", "column1 * weightbysample");
   auto h0 = df2.Histo1D({"h0", "weights", 3, 0, 3}, "weightbysample");
   auto h1 = df2.Histo1D({"h1", "calibrated data", 3, 0, 3}, "calibrated");
   df2 = df2.Filter("weightbysample >= 1");
   auto count0 = df2.Count();

   // Print sum of weights
   std::cout << "Sum of weights: " << sum0.GetValue() << std::endl;
   // Plots
   TCanvas *c0 = new TCanvas();
   h0->DrawClone();
   h0->GetXaxis()->SetTitle("weights");
   h0->GetYaxis()->SetTitle("counts");
   c0->SaveAs("h0.png");

   TCanvas *c1 = new TCanvas();
   h1->DrawClone();
   h1->GetXaxis()->SetTitle("calibrated data");
   h1->GetYaxis()->SetTitle("counts");
   c1->SaveAs("h1.png");

   std::cout << "Weights larger than 1: " << count0.GetValue() << std::endl;
}