#include "TLorentzVector.h"

//snippet of the analysis code
//--------------------
TLorentzVector tlv_epos, tlv_eneg, tlv_ee, tlv_phot;
float Lplus, Lminus, Pplus, Pminus;
float cosThetaCS;

// 4-vector of the Z boson
if(el_charge > 0 ){
  tlv_epos.SetPtEtaPhiM(el_pt,  el_eta,  el_phi,  0.511);
  tlv_eneg.SetPtEtaPhiM(tag_pt, tag_eta, tag_phi, 0.511);
 } else{
  tlv_eneg.SetPtEtaPhiM(el_pt,  el_eta,  el_phi,  0.511);
  tlv_epos.SetPtEtaPhiM(tag_pt, tag_eta, tag_phi, 0.511);
 }
tlv_ee = tlv_epos + tlv_eneg;

//prepare components of the cosCS calculation
Lplus  = tlv_eneg.E()+tlv_eneg.Pz();
Lminus = tlv_eneg.E()-tlv_eneg.Pz();
Pplus  = tlv_epos.E()+tlv_epos.Pz();
Pminus = tlv_epos.E()-tlv_epos.Pz();

//does the cosCS calculation
cosThetaCS  = (Lplus*Pminus - Lminus*Pplus);
cosThetaCS *= fabs(tlv_ee.Pz());
cosThetaCS /= (tlv_ee.Mag()*tlv_ee.Pz());
cosThetaCS /= sqrt(tlv_ee.Mag2() + tlv_ee.Pt()*tlv_ee.Pt() );

//does the phiCS calculation
float  phiCS =  getPhi(tlv_ee, tlv_eneg, tlv_epos, 1);

// Z boson: mass, Y , pt
ee_M       = tlv_ee.M() ;
ee_Y       = tlv_ee.Rapidity();
ee_pt      = tlv_ee.Pt();
//--------------------




//------------------------------------------------------    
// boost event into the Collins Soper frame sets the frame axis
// input TLorentzVector Z, bool qDirOK
// output CSaxis, xAxis, yAxis
// return false if sanity checks fail, true otherwise
//------------------------------------------------------
bool ScanZeeAi_CutFlow::CSFrame(TLorentzVector Z, TVector3 &CSAxis, 
			        TVector3 &xAxis, TVector3 &yAxis, bool qDirOK = true){

  double ProtonMass = 938.272; // MeV
  double BeamEnergy = 4000000; // MeV

  double sign  = fabs(Z.Z())/Z.Z();
  bool isGood = true;
  TLorentzVector p1, p2;

  if (qDirOK){
    p1.SetPxPyPzE(0, 0, sign*BeamEnergy, 
		  TMath::Sqrt(BeamEnergy*BeamEnergy+ProtonMass*ProtonMass)); // quark
    p2.SetPxPyPzE(0, 0, -1*sign*BeamEnergy, 
		  TMath::Sqrt(BeamEnergy*BeamEnergy+ProtonMass*ProtonMass)); // antiquark
  } else {
    p1.SetPxPyPzE(0, 0, -1*sign*BeamEnergy, 
		  TMath::Sqrt(BeamEnergy*BeamEnergy+ProtonMass*ProtonMass)); // quark
    p2.SetPxPyPzE(0, 0, sign*BeamEnergy, 
		  TMath::Sqrt(BeamEnergy*BeamEnergy+ProtonMass*ProtonMass)); // antiquark
  }
  p1.Boost(-Z.BoostVector());
  p2.Boost(-Z.BoostVector());
  CSAxis = (p1.Vect().Unit()-p2.Vect().Unit()).Unit();
  yAxis = (p1.Vect().Unit()).Cross((p2.Vect().Unit()));
  yAxis = yAxis.Unit();
  xAxis = yAxis.Cross(CSAxis);
  xAxis = xAxis.Unit();

  return isGood;
}

//------------------------------------------------------
// Calculates cos(Theta*) in the CSFrame.
// boost muon in Z-rest frame, calls CSFrame and calculates angle between CSAxis and the boosted muon.
// return cos(Theta*)
//------------------------------------------------------
double ScanZeeAi_CutFlow::getCosTheta(TLorentzVector Z, TLorentzVector muM, 
                                      TLorentzVector muP, bool qDirOK = true)
{
  TVector3 CSAxis, yAxis, xAxis;
  TLorentzVector boostedMu = muM;
  boostedMu.Boost(-Z.BoostVector());

  if (!CSFrame(Z, CSAxis, xAxis, yAxis, qDirOK))
    return -999;
  return cos(boostedMu.Angle(CSAxis));
}

//------------------------------------------------------
// Calculates phi in the CSFrame.
// boosts the muon into te Z - rest frame, calls CSFrame and calculates 
// the angle of the boosted muon, projected onto the x,y plane.
// phi is shifted by PI to be consistent with standard plot of phi.
// retrun phi
//-------------------------------------------------------
double ScanZeeAi_CutFlow::getPhi(TLorentzVector Z, TLorentzVector muM, 
			         TLorentzVector muP, bool qDirOK = true)
{
  TLorentzVector boostedMu = muM;
  TVector3 CSAxis, xAxis, yAxis;

  boostedMu.Boost(-Z.BoostVector());

  CSFrame(Z, CSAxis, xAxis, yAxis, qDirOK);
  double phi = atan2((boostedMu.Vect()*yAxis),(boostedMu.Vect()*xAxis));
  if(phi<0) phi = phi + 2*kPI;

  return phi;
}

//-------------------------------------------------------
// Calculates cos(Theta*) directly, without explicit transformation into CSFrame.
// The formula is commonly used does an implicit transformation into the CSFrame.
// return cos(Theta*)
//-------------------------------------------------------
double ScanZeeAi_CutFlow::getCosTheta_direct(TLorentzVector Z, 
				             TLorentzVector muM, TLorentzVector muP, bool qDirOK = true){
  double cosTheta;
  double sign  = fabs(Z.Z())/Z.Z();
  if (!qDirOK)
    sign *= -1;
  double part1 = 2 / ( Z.M()*sqrt(Z.M()*Z.M() + Z.Pt()*Z.Pt() ));
  double part2 = (muM.Plus()*muP.Minus() - muP.Plus()*muM.Minus())/2;
  cosTheta = part1*part2;
  // add sign according to boost direktion in z
  cosTheta *= sign;

  return cosTheta;
}
//-------------------------------------------------------