#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 rvec_f = const RVec<float> &;
using rvec_i = const RVec<int> &;

const auto z_mass = 91.2;
    
void selectMuon1()
{
  ROOT::EnableImplicitMT();
  ROOT::RDataFrame df("Delphes;4", "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 muon_mass, 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], 0.1);
        ROOT::Math::PtEtaPhiMVector m2((Muon_PT)[i2], (Muon_Eta)[i2], (Muon_Phi)[i2], 0.1);
            
        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 muon_mass)
{
   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], 0.1);
      ROOT::Math::PtEtaPhiMVector m2((Muon_PT)[i2], (Muon_Eta)[i2], (Muon_Phi)[i2], 0.1);
    z_masses[i] = (p1 + p2).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 muon_mass)
{
   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], muon_mass[i1]);
   ROOT::Math::PtEtaPhiMVector p2(Muon_PT[i2], Muon_Eta[i2], Muon_Phi[i2], muon_mass[i2]);
   ROOT::Math::PtEtaPhiMVector p3(Muon_PT[i3], Muon_Eta[i3], Muon_Phi[i3], muon_mass[i3]);
   ROOT::Math::PtEtaPhiMVector p4(Muon_PT[i4], Muon_Eta[i4], Muon_Phi[i4], muon_mass[i4]);
   return (p1 + p2 + p3 + p4).M();
}

   // Reconstruct Z systems
   auto df_z_idx = df.Define("Z_idx", reco_zz_to_4l, {"Muon_PT", "Muon_Eta", "Muon_Phi", "muon_mass", "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", "muon_mass"});
    // 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", "muon_mass"});
   return df_h_mass;

 }//end of void loop