#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++) {
            auto pass = [binsPt, binsEta, i, j](float pt, float eta, float probe) {
                return (pt > 25) && (pt >= binsPt[i]) && (pt < binsPt[i + 1]) && (eta >= binsEta[j]) && (eta < binsEta[j + 1]) && (probe > -1);

            };
            auto fail = [binsPt, binsEta, i, j](float pt, float eta, float probe) {
                return (pt > 25) && (pt >= binsPt[i]) && (pt < binsPt[i + 1]) && (eta >= binsEta[j]) && (eta < binsEta[j + 1]) && !(probe > -1);
            };
            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";
            h_tmp_pass.push_back(
                    rdf.Filter(pass, {"probe_lep_pt", "probe_lep_eta", "probe_muTrgPt"}).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, {"probe_lep_pt", "probe_lep_eta", "probe_muTrgPt"}).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();
}