#include <vector>
#include <string>
#include "ROOT/RDataFrame.hxx"
#include "ROOT/RVec.hxx"
#include "ROOT/RDF/RInterface.hxx"
#include "TCanvas.h"
#include "TH1D.h"
#include "TLatex.h"
#include "TLegend.h"
#include "Math/Vector4Dfwd.h"
#include "TStyle.h"
using namespace ROOT::VecOps;
using RNode = ROOT::RDF::RNode;
using rvec_f = const RVec<float> &;
using rvec_i = const RVec<int> &;

const auto z_mass = 91.2;
    

// Select interesting events with four muons
RNode selection_4mu(RNode df)
{
   auto df_ge4m = df.Filter("Muon_size>=4", "At least four muons");
   //auto df_iso = df_ge4m.Filter("All(abs(Muon_pfRelIso04_all)<0.40)", "Require good isolation");
   auto df_kin = df_ge4m.Filter("(Muon_PT>5) && (abs(Muon_Eta)<2.4)", "Good muon kinematics");
   //auto df_ip3d = df_kin.Define("Muon_ip3d", "sqrt(Muon_dxy*Muon_dxy + Muon_dz*Muon_dz)");
   //auto df_sip3d = df_ip3d.Define("Muon_sip3d", "Muon_ip3d/sqrt(Muon_dxyErr*Muon_dxyErr + Muon_dzErr*Muon_dzErr)");
   //auto df_pv = df_sip3d.Filter("All(Muon_sip3d<4) && All(abs(Muon_dxy)<0.5) && All(abs(Muon_dz)<1.0)", "Track close to primary vertex with small uncertainty");
   //auto df_2p2n = df_kin.Filter("Muon_size==4 && Sum(Muon_Charge==1)==2 && Sum(Muon_Charge==-1)==2", "Two positive and two negative muons");
   return df_kin;
}

  //ROOT::EnableImplicitMT();
  //ROOT::RDataFrame df("Delphes;5", "tag_1_delphes_events.root");

  //NOTE the following:
  //My branches have been renamed from Muon.PT, Muon.Eta, Muon.Phi, and Muon.Charge 
  //to Muon_PT, Muon_Eta, Muon_Phi, and Muon_Charge 
  //to avoid errors that were arising from the "."
  //GetColumnNames shows that the change in the Branch Names have been implemented
  //Need to create a "Muon_Mass" branch since don't have one in my Delphes tree; the size 49254 was taken from the Muon_PT.size()
  //auto m = df.Define("muon_mass", [] { return ROOT::RVec<float>(49254, 0.1); });
  RVec<RVec<size_t>> reco_zz_to_4l(rvec_f Muon_PT, rvec_f Muon_Eta, rvec_f Muon_Phi, rvec_f m, rvec_i Muon_Charge)
{
    

   RVec<RVec<size_t>> idx(2);
   idx[0].reserve(2); idx[1].reserve(2);
   // Find first lepton pair with invariant mass closest to Z mass
   auto idx_cmb = Combinations(Muon_PT, 2);
   auto best_mass = -1;
   size_t best_i1 = 0; size_t best_i2 = 0;
   for (size_t i = 0; i < idx_cmb[0].size(); i++) {
      const auto i1 = idx_cmb[0][i];
      const auto i2 = idx_cmb[1][i];
      if (Muon_Charge[i1] != Muon_Charge[i2]) {
        ROOT::Math::PtEtaPhiMVector m1((Muon_PT)[i1], (Muon_Eta)[i1], (Muon_Phi)[i1], m[i1]);
        ROOT::Math::PtEtaPhiMVector m2((Muon_PT)[i2], (Muon_Eta)[i2], (Muon_Phi)[i2], m[i2]);
            
        auto this_mass = (m1 + m2).M();
        if (std::abs(z_mass - this_mass) < std::abs(z_mass - best_mass)) {
            best_mass = this_mass;
            best_i1 = i1;
            best_i2 = i2;
         }
      }
   }
   idx[0].emplace_back(best_i1);
   idx[0].emplace_back(best_i2);
   // Reconstruct second Z from remaining lepton pair
   for (size_t i = 0; i < 4; i++) {
      if (i != best_i1 && i != best_i2) {
         idx[1].emplace_back(i);
      }
   }
   // Return indices of the pairs building two Z bosons
   return idx;
}
// Compute Z masses from four leptons of the same kind and sort ascending in distance to Z mass
RVec<float> compute_z_masses_4l(const RVec<RVec<size_t>> &idx, rvec_f Muon_PT, rvec_f Muon_Eta, rvec_f Muon_Phi, rvec_f m)
{
   RVec<float> z_masses(2);
   for (size_t i = 0; i < 2; i++) {
      const auto i1 = idx[i][0]; const auto i2 = idx[i][1];
      ROOT::Math::PtEtaPhiMVector m1((Muon_PT)[i1], (Muon_Eta)[i1], (Muon_Phi)[i1], m[i1]);
      ROOT::Math::PtEtaPhiMVector m2((Muon_PT)[i2], (Muon_Eta)[i2], (Muon_Phi)[i2], m[i2]);
    z_masses[i] = (m1 + m2).M();
   }
   if (std::abs(z_masses[0] - z_mass) < std::abs(z_masses[1] - z_mass)) {
      return z_masses;
   } else {
      return Reverse(z_masses);
   }
}

// Compute mass of Higgs from four leptons of the same kind
float compute_higgs_mass_4l(const RVec<RVec<size_t>> &idx, rvec_f Muon_PT, rvec_f Muon_Eta, rvec_f Muon_Phi, rvec_f m)
{
   const auto i1 = idx[0][0]; const auto i2 = idx[0][1];
   const auto i3 = idx[1][0]; const auto i4 = idx[1][1];
   ROOT::Math::PtEtaPhiMVector p1(Muon_PT[i1], Muon_Eta[i1], Muon_Phi[i1], m[i1]);
   ROOT::Math::PtEtaPhiMVector p2(Muon_PT[i2], Muon_Eta[i2], Muon_Phi[i2], m[i2]);
   ROOT::Math::PtEtaPhiMVector p3(Muon_PT[i3], Muon_Eta[i3], Muon_Phi[i3], m[i3]);
   ROOT::Math::PtEtaPhiMVector p4(Muon_PT[i4], Muon_Eta[i4], Muon_Phi[i4], m[i4]);
   return (p1 + p2 + p3 + p4).M();
}

RNode reco_higgs_to_4mu(RNode df)
{
   //Filter interesting events
   auto df_base = selection_4mu(df);
   // Reconstruct Z systems
   auto df_z_idx = df.Define("Z_idx", reco_zz_to_4l, {"Muon_PT", "Muon_Eta", "Muon_Phi", "m", "Muon_Charge"});
   // Compute masses of Z systems
   auto df_z_mass = df_z_idx.Define("Z_mass", compute_z_masses_4l, {"Z_idx", "Muon_PT", "Muon_Eta", "Muon_Phi", "m"});
    // Reconstruct H mass
   auto df_h_mass = df_z_mass.Define("H_mass", compute_higgs_mass_4l, {"Z_idx", "Muon_PT", "Muon_Eta", "Muon_Phi", "m"});
   return df_h_mass;
};
void selectMuon1()
{
   ROOT::EnableImplicitMT();
   ROOT::RDataFrame df("Delphes;5", "tag_1_delphes_events.root");
   auto df_sig_4mu_reco = reco_higgs_to_4mu(df);
   auto df_h_sig_4mu = df_sig_4mu_reco.Histo1D({"h_sig_4mu", "", 3000, 0.25, 300}, "H_mass");
   df_h_sig_4mu->Draw();
}


//histZMass->Draw();



   //PLOTS

   // Make histogram of dimuon mass spectrum
   //auto h = df_z_mass.Histo1D({"Z_mass", "Z_mass", 30000, 0.25, 300}, "Z_mass");
   //h->Draw();
   // Request cut-flow report
   //auto report = df.Report();
   // Produce plot
   /*
   gStyle->SetOptStat(0); gStyle->SetTextFont(42);
   auto c = new TCanvas("c", "", 800, 700);
   c->SetLogx(); c->SetLogy();
   h->SetTitle("");
   h->GetXaxis()->SetTitle("m_{#mu#mu} (GeV)"); h->GetXaxis()->SetTitleSize(0.04);
   h->GetYaxis()->SetTitle("N_{Events}"); h->GetYaxis()->SetTitleSize(0.04);
   h->DrawClone();
   TLatex label; label.SetNDC(true);
   label.DrawLatex(0.175, 0.740, "#eta");
   label.DrawLatex(0.205, 0.775, "#rho,#omega");
   label.DrawLatex(0.270, 0.740, "#phi");
   label.DrawLatex(0.400, 0.800, "J/#psi");
   label.DrawLatex(0.415, 0.670, "#psi'");
   label.DrawLatex(0.485, 0.700, "Y(1,2,3S)");
   label.DrawLatex(0.755, 0.680, "Z");
   label.SetTextSize(0.040); label.DrawLatex(0.100, 0.920, "#bf{CMS Open Data}");
   label.SetTextSize(0.030); label.DrawLatex(0.630, 0.920, "#sqrt{s} = 8 TeV, L_{int} = 11.6 fb^{-1}");
   c->SaveAs("dimuon_spectrum.pdf");
   */

   // Print cut-flow report
   //report->Print();

 //}//end of void loop
    
     
int main()
{
   selectMuon1();
}