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


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() {

	vector<string> fileNames;
	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.;
			});
	
	std::cout << "Sum of weights: " << df2.Sum<double>("weightbysample").GetValue() << std::endl;

	// Define new column with calibrated data
	
	df2 = df2.Define("calibrated","column1 * weightbysample");

    // Plots
    
	auto h0 = df2.Histo1D({"h0", "weights", 3, 0, 3}, "weightbysample");
	TCanvas* c0 = new TCanvas();
    h0->DrawClone();
    h0->GetXaxis()->SetTitle("weights");
    h0->GetYaxis()->SetTitle("counts");

	auto h1 = df2.Histo1D({"h1", "calibrated data", 3, 0, 3}, "calibrated");
	TCanvas* c1 = new TCanvas();
    h1->DrawClone();
    h1->GetXaxis()->SetTitle("calibrated data");
    h1->GetYaxis()->SetTitle("counts");

	// Filter and check the weights
	df2 = df2.Filter("weightbysample >= 1");
	std::cout << "Weights larger than 1: " << df2.Count().GetValue() << endl;

}