#include <Math/Vector4D.h>
#include <ROOT/RVec.hxx>
#include <string>
using std::string;

using Vec_t = const ROOT::RVec<float> &;
using namespace ROOT::Math;
ROOT::VecOps::RVec<PxPyPzEVector>
TrueLep_Vecs(string variable, Vec_t Particle_PID, Vec_t Particle_E,
             Vec_t Particle_Px, Vec_t Particle_Py, Vec_t Particle_Pz,
             Vec_t Particle_Charge, Vec_t Particle_Status, Vec_t Particle_M1,
             Vec_t Particle_PT, Vec_t Particle_Eta) {

  ROOT::VecOps::RVec<PxPyPzEVector> TrueLep_Vec;
  ROOT::VecOps::RVec<float> TrueLepPT_Vec;
  Long64_t nne = Particle_PID.size();

  // filling the vectors with 4-momentum, charge and pT values
  for (int j = 0; j < nne; j++) {
    if (Particle_Status[j] == 1 && fabs(Particle_PID[j]) == 11 &&
        (fabs(Particle_PID[Particle_M1[j]]) == 24 ||
         fabs(Particle_PID[Particle_M1[j]]) == 23 ||
         fabs(Particle_PID[Particle_M1[j]]) == 9900012 ||
         fabs(Particle_PID[Particle_M1[j]]) == 15) &&
        Particle_PT[j] > 10.0 && fabs(Particle_Eta[j]) < 2.47) {
      if (variable == "TrueLep") {
        TrueLep_Vec.emplace_back(Particle_Px[j], Particle_Py[j], Particle_Pz[j],
                                 Particle_E[j]);
        TrueLepPT_Vec.emplace_back(Particle_PT[j]);
      }
      if (variable == "TrueEle") {
        TrueLep_Vec.emplace_back(Particle_Px[j], Particle_Py[j], Particle_Pz[j],
                                 Particle_E[j]);
        TrueLepPT_Vec.emplace_back(Particle_PT[j]);
      }
    }
    if (Particle_Status[j] == 1 && fabs(Particle_PID[j]) == 13 &&
        (fabs(Particle_PID[Particle_M1[j]]) == 24 ||
         fabs(Particle_PID[Particle_M1[j]]) == 23 ||
         fabs(Particle_PID[Particle_M1[j]]) == 9900012 ||
         fabs(Particle_PID[Particle_M1[j]]) == 15) &&
        Particle_PT[j] > 10.0 && fabs(Particle_Eta[j]) < 2.40) {
      if (variable == "TrueLep") {
        TrueLep_Vec.emplace_back(Particle_Px[j], Particle_Py[j], Particle_Pz[j],
                                 Particle_E[j]);
        TrueLepPT_Vec.emplace_back(Particle_PT[j]);
      }
      if (variable == "TrueMuon") {
        TrueLep_Vec.emplace_back(Particle_Px[j], Particle_Py[j], Particle_Pz[j],
                                 Particle_E[j]);
        TrueLepPT_Vec.emplace_back(Particle_PT[j]);
      }
    }
  }

  // arranging all vectors in descending pT order
  auto TrueLep_Vec_sorted = Take(TrueLep_Vec, Reverse(Argsort(TrueLepPT_Vec)));

  return TrueLep_Vec_sorted;
}

using namespace ROOT::Math;
ROOT::VecOps::RVec<int>
TrueLep_Charges(string variable, Vec_t Particle_PID, Vec_t Particle_E,
                Vec_t Particle_Px, Vec_t Particle_Py, Vec_t Particle_Pz,
                Vec_t Particle_Charge, Vec_t Particle_Status, Vec_t Particle_M1,
                Vec_t Particle_PT, Vec_t Particle_Eta) {

  ROOT::VecOps::RVec<int> TrueLepCharge_Vec;
  ROOT::VecOps::RVec<float> TrueLepPT_Vec;
  Long64_t nne = Particle_PID.size();

  // filling the vectors with 4-momentum, charge and pT values
  for (int j = 0; j < nne; j++) {
    if (Particle_Status[j] == 1 && fabs(Particle_PID[j]) == 11 &&
        (fabs(Particle_PID[Particle_M1[j]]) == 24 ||
         fabs(Particle_PID[Particle_M1[j]]) == 23 ||
         fabs(Particle_PID[Particle_M1[j]]) == 9910012 ||
         fabs(Particle_PID[Particle_M1[j]]) == 15) &&
        Particle_PT[j] > 10.0 && fabs(Particle_Eta[j]) < 2.47) {
      if (variable == "TrueEleCharge") {
        TrueLepCharge_Vec.emplace_back(Particle_Charge[j]);
        TrueLepPT_Vec.emplace_back(Particle_PT[j]);
      }
    }
    if (Particle_Status[j] == 1 && fabs(Particle_PID[j]) == 13 &&
        (fabs(Particle_PID[Particle_M1[j]]) == 24 ||
         fabs(Particle_PID[Particle_M1[j]]) == 23 ||
         fabs(Particle_PID[Particle_M1[j]]) == 9910012 ||
         fabs(Particle_PID[Particle_M1[j]]) == 15) &&
        Particle_PT[j] > 10.0 && fabs(Particle_Eta[j]) < 2.40) {
      if (variable == "TrueMuonCharge") {
        TrueLepCharge_Vec.emplace_back(Particle_Charge[j]);
        TrueLepPT_Vec.emplace_back(Particle_PT[j]);
      }
    }
  }

  // arranging all vectors in descending pT order
  auto TrueLepCharge_Vec_sorted =
      Take(TrueLepCharge_Vec, Reverse(Argsort(TrueLepPT_Vec)));

  return TrueLepCharge_Vec_sorted;
}