#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==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;
}


// Compute mass of two Z candidates from two electrons and two muons and sort ascending in distance to Z mass
RVec<float> compute_z_masses_2mu2mu(rvec_f mu_pt, rvec_f mu_eta, rvec_f mu_phi, rvec_f mu_mass)
{
   ROOT::Math::PtEtaPhiMVector p1(mu_pt[0], mu_eta[0], mu_phi[0], mu_mass[0]);
   ROOT::Math::PtEtaPhiMVector p2(mu_pt[1], mu_eta[1], mu_phi[1], mu_mass[1]);
   ROOT::Math::PtEtaPhiMVector p3(mu_pt[2], mu_eta[2], mu_phi[2], mu_mass[2]);
   ROOT::Math::PtEtaPhiMVector p4(mu_pt[3], mu_eta[3], mu_phi[3], mu_mass[3]);
   auto mu_z1 = (p1 + p2).M();
   auto mu_z2 = (p3 + p4).M();
   RVec<float> z_masses(2);
   if (std::abs(mu_z1 - z_mass) < std::abs(mu_z2 - z_mass)) {
      z_masses[0] = mu_z1;
      z_masses[1] = mu_z2;
   } else {
      z_masses[0] = mu_z1;
      z_masses[1] = mu_z2;
   }
   return z_masses;
}

// Compute Higgs mass from two electrons and two muons
float compute_xx_mass_2mu2mu( rvec_f mu_pt, rvec_f mu_eta, rvec_f mu_phi, rvec_f mu_mass)
{
   ROOT::Math::PtEtaPhiMVector p1(mu_pt[0], mu_eta[0], mu_phi[0], mu_mass[0]);
   ROOT::Math::PtEtaPhiMVector p2(mu_pt[1], mu_eta[1], mu_phi[1], mu_mass[1]);
   ROOT::Math::PtEtaPhiMVector p3(mu_pt[2], mu_eta[2], mu_phi[2], mu_mass[2]);
   ROOT::Math::PtEtaPhiMVector p4(mu_pt[3], mu_eta[3], mu_phi[3], mu_mass[3]);
   return (p1 + p2 + p3 + p4).M();
}

// Reconstruct Higgs from two electrons and two muons
RNode reco_XX_to_2mu2mu(RNode df)
{
   // Filter interesting events
   auto df_base = selection_4mu(df);
   // Compute masses of Z systems
   auto df_z_mass =
      df_base.Define("Z_mass", compute_z_masses_2mu2mu, {"Muon_PT", "Muon_Eta", "Muon_Phi", "m"});
   // Cut on mass of Z candidates
   //auto df_z_cut = filter_z_candidates(df_z_mass);
   // Reconstruct H mass
   auto df_h_mass = df_z_mass.Define("XX_mass", compute_xx_mass_2mu2mu, { "Muon_PT", "Muon_Eta", "Muon_Phi", "m"});
   return df_h_mass;
}



void selectXX1()
{
   ROOT::EnableImplicitMT();
   ROOT::RDataFrame df("Delphes;5", "tag_1_delphes_events.root");
   auto df_sig_2mu2mu_reco = reco_XX_to_2mu2mu(df);
   auto df_xx_sig_4mu = df_sig_2mu2mu_reco.Histo1D({"xx_sig_4mu", "", 3000, 0.25, 300}, "XX_mass");
   df_xx_sig_4mu->Draw();
}