from ThreeVecRef import ThreeVec
import math

class LorentzVector(ThreeVec): # inherits from ThreeVec
    "Class for 4 Vector and operations"
    def __init__( self,t=0., x=0., y=0., z=0.):
        ThreeVec.__init__( self, x, y, z )  # initialize  ThreeVec
        self.t = t
    def Add( self, tv ):
        w = ThreeVec.Add( self, tv ) # call ThreeVec method first
        newvec = LorentzVector(  self.t+tv.t, w.x, w.y, w.z)
        return newvec
    def __add__( self, tv ):
        "method for overloading the + operator"
        newvec = self.Add(tv)
        return newvec

#    def Angle( self, tv ): no need to define here, available from ThreeVec
    def Mass( self ):
        return math.sqrt(self.t**2 - self.Length()**2)
    def InvMass( self, tv ):
        return math.sqrt((self.t+tv.t)**2 - ( (self.x+tv.x)**2 + (self.y+tv.y)**2 + (self.z+tv.z)**2 ))
    def __str__( self ):
        return "( %6.3f, %6.3f, %6.3f, %6.3f )" % ( self.t, self.x, self.y, self.z)
    def setXYZT( self,x, y, z, t):
        self.x = x
        self.y = y
        self.z = z
        self.t = t

    def setXYZM(self, x, y, z,  m):
        if ( m  >= 0 ): 
            self.setXYZT( x, y, z, math.sqrt(x*x+y*y+z*z+m*m))
        else :
            self.setXYZT( x, y, z, math.sqrt( max((x*x+y*y+z*z-m*m), 0. )))

    def setPtEtaPhiM(self, pt, eta, phi, m):
        pt = abs(pt)
        self.setXYZM(pt*math.cos(phi), pt*math.sin(phi), pt*math.sinh(eta) ,m)


if __name__ == "__main__" :
    a = LorentzVector(45.0002,0.,0.,45.0)
    b = LorentzVector(45.0002,31.8198,0.,31.8198)
    print ("Angle = " , a.Angle(b))
    print ("Mass a = " , a.Mass())
    print ("Mass b = " , b.Mass())
    print ("Mass a+b = " , b.InvMass(a))
    print (a, b)

    # Higgs example
    # https://twiki.cern.ch/twiki/bin/view/AtlasPublic/EventDisplaysFromHiggsSearches
    # and more infos: https://atlas.cern/updates/briefing/higgs-golden-channel 
    #                 https://atlas.cern/discover/physics
    # 
    # H->ZZ-> 4mu
    #m_4l=123.5 GeV. m_12=84 GeV, m_34=34.2 GeV. 
    #mu_1: pt=37.8 GeV, eta=0.61 phi=1.46. 
    #mu_2: pt=29.2 GeV, eta=-0.95, phi=-2.47. 
    #mu_3: pt=10.3 GeV, eta=0.62 phi=-1.41. 
    #mu_4: pt=32.6 GeV eta=-0.16, phi=2.85
    #
    # inv mass plot, animated gif: https://cds.cern.ch/record/2230893?ln=de
    #
    print ("=================")

    c = LorentzVector()
    d = LorentzVector()
    e = LorentzVector()
    f = LorentzVector()
    c.setPtEtaPhiM(37.8, 0.61, 1.46, 0.105)
    d.setPtEtaPhiM(29.2,-0.95,-2.47, 0.105)
    e.setPtEtaPhiM(10.3, 0.62,-1.41, 0.105)
    f.setPtEtaPhiM(32.6,-0.16, 2.85, 0.105)

    z1 = c+d
    z2 = e+f
    h = c+d+e+f

    print ("c Mass = %5.3f, %s " % (c.Mass(), c))
    print ("d Mass = %5.3f, %s " % (d.Mass(), d))
    print ("e Mass = %5.3f, %s " % (e.Mass(), e))
    print ("f Mass = %5.3f, %s " % (f.Mass(), f))


    print ("z1 = c + d = %5.3f, %s" %(c.InvMass(d), z1))
    print ("z2 = e + f = %5.3f, %s" %(e.InvMass(f), z2))
    print ("z1 = %5.3f" %( z1.Mass() ))
    print ("z2 = %5.3f" %( z2.Mass() ))

    print ("Mass h = z1+ z2 = c + d + e + f = " , h.Mass())