#include "ROOT/RDataFrame.hxx"
#include <vector>
#include <string>
#include <sstream>
#include <iostream>

int main() {
    std::vector<float> binsPt;
    std::vector<float> binsEta;
    for (int i = 0; i < 10; i++) binsPt.push_back(25 + 5 * i);
    for (int i = 0; i < 10; i++) binsEta.push_back(-2.4 + 0.1 * i);

    ROOT::RDataFrame rdf("IDIsoToHLT/fitter_tree", "mdunser_data.root");

    std::vector<ROOT::RDF::RResultPtr<TH1D>> h_tmp_pass;
    std::vector<ROOT::RDF::RResultPtr<TH1D>> h_tmp_fail;
    int n = 0;
    for(int i = 0; i < binsPt.size() - 1; i++) {
        for(int j = 0; j < binsEta.size() - 1; j++) {
            std::stringstream ss2;
            ss2 << "bin" << i << j;
            const auto tmp_name = ss2.str();

            const auto tmp_name_pass = tmp_name + "_Pass";
            const auto tmp_name_fail = tmp_name + "_Fail";
            const auto pass_str = "(probe_lep_pt > 25) && (probe_lep_pt >= " + std::to_string(binsPt[i]) + ") && (probe_lep_pt < " + std::to_string(binsPt[i + 1]) + ") && (probe_lep_eta >= " + std::to_string(binsEta[j]) + ") && (probe_lep_eta <" + std::to_string(binsEta[j + 1]) + ") && (probe_muTrgPt > -1)";
            const auto fail_str = "(probe_lep_pt > 25) && (probe_lep_pt >= " + std::to_string(binsPt[i]) + ") && (probe_lep_pt < " + std::to_string(binsPt[i + 1]) + ") && (probe_lep_eta >= " + std::to_string(binsEta[j]) + ") && (probe_lep_eta < " + std::to_string(binsEta[j + 1]) + ") && !(probe_muTrgPt > -1)";
            h_tmp_pass.push_back(rdf.Filter(pass_str).Histo1D<float>({tmp_name_pass.c_str(), tmp_name_pass.c_str(), 60, 60., 120.}, "pair_mass"));
            h_tmp_fail.push_back(rdf.Filter(fail_str).Histo1D<float>({tmp_name_fail.c_str(), tmp_name_fail.c_str(), 60, 60., 120.}, "pair_mass"));
            n++;
        }
    }
    std::cout << "Created " << 2 * n << " histograms." << std::endl;

    auto f = TFile::Open("output.root", "RECREATE");
    for (auto h : h_tmp_pass) h->Write();
    for (auto h : h_tmp_fail) h->Write();
    f->Close();
}