#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/Vector4D.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 == 6", "Events with 6 muons");
   
   auto df_kin = df_ge4m.Filter("(Muon.PT>5) && (abs(Muon.Eta)<2.4)", "Good muon kinematics");
   return df_kin;
}

  
  //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 pt, rvec_f eta, rvec_f phi, rvec_f mass, rvec_i 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(pt, 2); // Return the indices that represent all unique combinations of the elements of a given RVec
   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 (charge[i1] != charge[i2]) {
        ROOT::Math::PtEtaPhiMVector m1((pt)[i1], (eta)[i1], (phi)[i1], mass[i1]);
        ROOT::Math::PtEtaPhiMVector m2((pt)[i2], (eta)[i2], (phi)[i2], mass[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 pt, rvec_f eta, rvec_f phi, rvec_f 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((pt)[i1], (eta)[i1], (phi)[i1], mass[i1]);
      ROOT::Math::PtEtaPhiMVector m2((pt)[i2], (eta)[i2], (phi)[i2], mass[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_XX_mass_4l(const RVec<RVec<size_t>> &idx, rvec_f pt, rvec_f eta, rvec_f phi, rvec_f 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((pt)[i1], (eta)[i1], (phi)[i1], (mass)[i1]);
   ROOT::Math::PtEtaPhiMVector p2((pt)[i2], (eta)[i2], (phi)[i2], (mass)[i2]);
   ROOT::Math::PtEtaPhiMVector p3((pt)[i3], (eta)[i3], (phi)[i3], (mass)[i3]);
   ROOT::Math::PtEtaPhiMVector p4((pt)[i4], (eta)[i4], (phi)[i4], (mass)[i4]);
   return (p1 + p2 + p3 + p4).M();
}

RNode reco_XX_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("XX_mass", compute_XX_mass_4l, {"Z_idx", "Muon.PT", "Muon.Eta", "Muon.Phi", "m"});
   return df_h_mass;
}


void selectXXt()
{
   ROOT::EnableImplicitMT();
   ROOT::RDataFrame df("Delphes;5", "tag_1_delphes_events.root");
   auto df_sig_4mu_reco = reco_XX_to_4mu(df);
   auto df_xx_sig_4mu = df_sig_4mu_reco.Histo1D<int>({"xx_sig_4mu", "", 3000, 0.25, 300}, "XX_mass");
   df_xx_sig_4mu->Draw();
}


int main()
{
   selectXXt();
}