#include <TCanvas.h>
#include <TF1.h>
#include <TH1D.h>
#include <TH2D.h>
#include <TLegend.h>
#include <TMath.h>
#include <string.h>
#include "Math/Minimizer.h"

#include "TFile.h"
#include "TTree.h"
#include "TRandom.h"

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
// Class to create function to minimize, with dataframe parameters that can be set through a method 
//
class generateFunction
{
	private:
		string columnName;
	  	ROOT::RDF::RNode data1;
	  	ROOT::RDF::RNode data2;

	public:
		// set parameters
		void SetParameters(ROOT::RDF::RNode df1,ROOT::RDF::RNode df2, string name)
        	{
            	auto data1 = df1;
	            auto data2 = df2;
    	        auto columnName = name;
        	}

		// function that returns a double value after applying a cut on two dataframes, to be minimized
		double  calculateFOM(double l, double u){

			data1 = data1.Filter(std::to_string(l) + " < " + columnName + " < " + std::to_string(u));
			data2 = data2.Filter(std::to_string(l) + " < " + columnName + " < " + std::to_string(u));

		    auto temp1 = data2.Count();
		    int nSignal = temp1.GetValue();
		    auto temp2 = data1.Count();
		    int nBg = temp2.GetValue();

		    double FOM = -nSignal / sqrt(nBg);

		    return FOM;
		}
};

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
// function to generate root files for testing 
//
void generateData(const char *filename, int n){

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

	double var1;
	double var2;

	TTree *tree = new TTree("AnalysisTree","AnalysisTree");
   	auto branch1 = tree->Branch("column1",&var1,"column1/D");
	auto branch2 = tree->Branch("column2",&var2,"column2/D");	
	
	for (unsigned int i=0; i<n; i++){
		var1 = rand() %100;
		var2 = rand() %100;
		tree->Fill();
	}
	tree->Write();
	f.Write();
	f.Close();

}

///////////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
// Main function: sequential optimization in several iterations
//

Int_t optimizeCuts() {

    // ******************************
    // Instantiate the class
    // ******************************
	
	generateFunction func;

    // ******************************
    // Data
    // ******************************

	// Create two files with trees for testing

	generateData("test1.root",100);
	generateData("test2.root",100);

	// ******************************
	// Create RDataFrames 
	// ******************************
	
	ROOT::RDataFrame data1("AnalysisTree","test1.root");
	ROOT::RDataFrame data2("AnalysisTree","test2.root");

	// initialize filtered dataframe
	auto data_filtered1 = data1.Filter("column1 > 0");
	auto data_filtered2 = data2.Filter("column1 > 0");
	
	// ******************************
	// sequential optimization
	// ******************************
	
	// cut1:
	string columnName1 = "column1";
	double l1 = 0;
	double u1 = 1;
		
	func.SetParameters(data_filtered1, data_filtered2, columnName1);

	double FOM1 = func.calculateFOM(l1,u1);
	cout << FOM1 << endl;

	// here will be the minimization be done

	// apply cuts that result in best fit
	data_filtered1 = data1.Filter(std::to_string(l1) + "<" + columnName1 + "<" + std::to_string(u1));
	data_filtered2 = data2.Filter(std::to_string(l1) + "<" + columnName1 + "<" + std::to_string(u1));


	// *****	
	// cut2:
	// repeat with cut on different column
	string columnName2 = "column2";
	double l2 = 0;
	double u2 = 50;
		
	func.SetParameters(data_filtered1, data_filtered2, columnName2);
	
	double FOM2 = func.calculateFOM(l2,u2);
	cout << FOM2 << endl;	

return 0;
}