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

#include "TFile.h"
#include "TRandom.h"
#include "TTree.h"
#include <ROOT/RDataFrame.hxx>

using namespace std;

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
// 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::RDataFrame(0);
  ROOT::RDF::RNode data2 = ROOT::RDataFrame(0);

public:

  generateFunction() = default;

  // set parameters
  void SetParameters(ROOT::RDF::RNode df1, ROOT::RDF::RNode df2, string name) {
    data1 = df1;
    data2 = df2;
    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([l,u](double c) { return l < c && c < u; }, {columnName});
    data2 = data2.Filter([l,u](double c) { return l < c && c < u; }, {columnName});

    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;
}

int main() { return optimizeCuts(); }