#include <ROOT/RDataFrame.hxx>
#include <ROOT/TProcessExecutor.hxx>
#include <TFile.h>
#include <TTree.h>
#include <TStopwatch.h>
#include <vector>
#include <string>
#include <iostream>
using namespace std;

std::vector<double> ConvertTreeCpp(TTree &inputTree, const char *outputFileName, const char* outputTreeName, const double datasetWeight) {
	
	ROOT::RDataFrame inputDataFrame(inputTree);
	auto finalDF = inputDataFrame.Filter([](double x) { return x > 10; }, {"soft_term_pt"});

	finalDF = finalDF.Alias("nJets", "n_jet")
		.Alias("photonEta", "ph_eta")
		.Alias("subleadJetEta", "jet_sublead_eta")
		.Alias("metPt", "metTST_pt")
		.Alias("phCentrality", "g_centr")
		.Alias("nLeptons", "n_lep")
		.Alias("mJJ", "m_jj")
		.Alias("deltaYJJ", "dYjj")
		.Alias("deltaYJPh", "dYjg")
		.Alias("ptBalance", "pt_bal")
		.Alias("ptBalanceRed", "pt_bal_tr")
		.Define("sinDeltaPhiJJOver2", [](double dphi_jj) { return TMath::Sin(TMath::Abs(dphi_jj) / 2); }, {"dphi_jj"})
		.Alias("leadJetPt", "jet_lead_pt")
		.Define("weightModified", [datasetWeight](double weight) { return weight*datasetWeight; }, {"weight"})
		.Alias("nJets20", "n_jet20")
		.Alias("nJets25", "n_jet25")
		.Alias("nJets30", "n_jet30")
		.Alias("nJets35", "n_jet35")
		.Alias("nJets40", "n_jet40")
		.Alias("nJets45", "n_jet45");

	vector<string> branchList{"mJJ",
		"deltaYJJ",
		"ptBalance",
		"deltaYJPh",
		"phCentrality",
		"sinDeltaPhiJJOver2",
		"nJets",
		"weightModified",
		"ptBalanceRed",
		"leadJetPt",
		"subleadJetEta",
		"photonEta",
		"metPt",
		"nLeptons",
		"nJets20",
		"nJets25",
		"nJets30",
		"nJets35",
		"nJets40",
		"nJets45"};


	auto passedEntryNum = finalDF.Count();
	auto passedWeightSum = finalDF.Sum<double>("weight");
	auto passedWeightSumNorm = finalDF.Sum<double>("weightModified");
	finalDF.Snapshot<Double_t, Double_t, Double_t, Double_t, Double_t, double, Double_t, double, Double_t, Double_t, Double_t, Double_t, Double_t, Double_t, Double_t, Double_t, Double_t, Double_t, Double_t, Double_t>(outputTreeName, outputFileName, branchList);

	return {double(passedEntryNum.GetValue()), passedWeightSum.GetValue(), passedWeightSumNorm.GetValue()};
}

void run() {
   const std::vector<std::string> filenames{"Heavy_a.root", "Heavy_d.root", "Heavy_e.root"};
   const std::vector<int> luminosities{1,2,3};

   auto processOne = [&filenames, &luminosities](int idx) {
      TFile inputFile(("./input/" + filenames[idx]).c_str());

      auto *weightTree = inputFile.Get<TTree>("output_tree_sw");
      int weightSum;
      weightTree->SetBranchAddress("sum_of_weights_bk_xAOD", &weightSum);
      weightTree->GetEntry(0);

      auto *normKoeffTree = inputFile.Get<TTree>("norm_tree");
      double normKoeff;
      normKoeffTree->SetBranchAddress("koef", &normKoeff);
      normKoeffTree->GetEntry(0);

      const auto datasetWeight = luminosities[idx] * normKoeff / weightSum;

      auto *tree = inputFile.Get<TTree>("output_tree_sys_default");
      ConvertTreeCpp(*tree, ("./output/" + filenames[idx]).c_str(), "TMVA_input", datasetWeight);
      return true;
   };

   ROOT::TProcessExecutor pool(filenames.size());
   pool.Map(processOne, ROOT::TSeqI(filenames.size()));
}

int main() {
   std::vector<double> times;
   for (int i = 0; i < 3; ++i) {
      TStopwatch sw;
      sw.Start();
      run();
      sw.Stop();
      times.push_back(sw.RealTime());
   }
   TStatistic stat("", times.size(), times.data());
   std::cout << stat.GetMean() << " +/- " << stat.GetRMS() << std::endl;
   return 0;
}