#-------------------------------------------------------
# example31.py
# X -> p pi- decay
# all values are in GeV
#-------------------------------------------------------
from ROOT import TLorentzVector

deg = 180 / 3.1415926

# given 4-vectors of decay products
proton = TLorentzVector()
pion = TLorentzVector()
X = TLorentzVector()
proton.SetXYZM(-0.49, -0.20, 2.11, 0.938)
pion.SetXYZM(-0.26, -0.05, 0.47, 0.140)

# setup mother 4-vector
X = proton + pion;
print("Mass of X = ",X.M())
print("Pt of X   = ",X.Pt())
print("Eta of X  = ",X.Eta())
print("Phi of X  = ",X.Phi())

# angle
angle = proton.Vect().Angle( pion.Vect() )  # lab frame
print("Angle between proton and pion in lab = " , angle * deg)
print("KE of proton in lab = ", proton.E()-proton.M())
print("KE of pion   in lab = ", pion.E()-pion.M())

# go to CM frame of X
betax = X.Px()/X.E()
betay = X.Py()/X.E()
betaz = X.Pz()/X.E()

# Boost() function modifies the original 4-vector components
proton.Boost(-betax, -betay, -betaz)
pion.Boost(-betax, -betay, -betaz)
angle_cm = proton.Vect().Angle( pion.Vect() ) # cm frame

print( "Angle between proton and pion in CM  = " , angle_cm * deg)
print( "KE of proton in lab = ", proton.E()-proton.M())
print( "KE of pion   in lab = ", pion.E()-pion.M())