#include "Math/GenVector/Boost.h"
#include "Math/LorentzVector.h"
#include "Math/PtEtaPhiE4D.h"
typedef ROOT::Math::LorentzVector<ROOT::Math::PtEtaPhiE4D<float>> LorentzVector;

double cosTheta(LorentzVector wP4, LorentzVector zP4, LorentzVector lP4) {
  ROOT::Math::Boost boostWcm;
  boostWcm.SetComponents(wP4.BoostToCM());
  ROOT::Math::Boost boostWZcm;
  boostWZcm.SetComponents((wP4+zP4).BoostToCM());

  auto lP4_w = boostWcm(lP4);
  auto wP4_wz = boostWZcm(wP4);
  double ct = lP4_w.Vect().Dot(wP4_wz.Vect()) / (lP4_w.P()*wP4_wz.P());
  return ct;
}

LorentzVector wP4estimate(LorentzVector lP4, double pTmiss, double phi_miss) {
  const double wMass = 80.385;
  double x = wMass*wMass/(2*lP4.Pt()*pTmiss) + std::cos(lP4.phi()-phi_miss);
  double eta_nu;
  if ( x < 1. ) {
    // Complex solution
    double complex_arg = std::atan2(std::sqrt(1-x*x), x);
    eta_nu = std::asinh(std::cos(complex_arg)*std::sinh(lP4.Eta()));
  }
  else {
    eta_nu = lP4.Eta() + (lP4.Eta() > 0. ? -1.:1) * std::acosh(x);
  }
  LorentzVector nuP4(pTmiss, eta_nu, phi_miss, pTmiss*cosh(eta_nu));
  return nuP4 + lP4;
}