import numpy as np

#c = 2.99792458e8
#h = 6.6260689633e-34
#joules_per_MeV = 1.602176487e-19/1e-6
kg_per_MeV = 1.782661758e-36/1e-6
pi0_mass = 134.9766*kg_per_MeV

def rocket_to_lab(energy, momentum, v):
    """
    Return the energy and momentum of a particle in the lab frame from its
    energy and momentum in a rocket frame traveling at a velocity `v` with
    respect to the lab frame.

    Args:
        - energy: float
            The energy of the particle in the rocket frame in kg.
        - momentum: array
            The momentum of the particle in the rocket frame in kg.
        - v: array
            The velocity of the rocket frame as seen in the lab frame in units
            of c.
    """
    e0 = float(energy)#/c**2
    p0 = np.asarray(momentum, float)#/c
    v = np.asarray(v, float)#/c

    try:
        assert e0**2 - p0.dot(p0) >= -1.0e-70
    except AssertionError:
        print e0**2 - p0.dot(p0)
        raise

    g = 1.0/np.sqrt(1.0-v.dot(v))

    x = np.dot(p0, v)/np.linalg.norm(v)
    p = p0 + ((g-1.0)*x + g*np.linalg.norm(v)*e0)*v/np.linalg.norm(v)
    e = np.sqrt(e0**2 - p0.dot(p0) + p.dot(p))

    #return e*c**2, p*c
    return e, p

def pi0_decay(energy, direction, theta, phi):
    """
    Return the energy and directions for two photons produced from a pi0
    decay in the lab frame given that one of the photons decayed at polar
    angles `theta` and `phi` in the pi0 rest frame.

    Args:
        - energy: float
            The total energy of the pi0 in MeV.
        - direction: array
            The direction of the pi0's velocity in the lab frame.
    """
    direction = np.asarray(direction)/np.linalg.norm(direction)
    pi0_e = float(energy)*kg_per_MeV#/c**2
    pi0_p = np.sqrt(pi0_e**2-pi0_mass**2)*direction
    pi0_v = pi0_p/pi0_e

    photon_e0 = pi0_mass/2.0
    photon_p0 = photon_e0*np.array([np.cos(theta)*np.sin(phi), np.sin(theta)*np.sin(phi), np.cos(phi)])

    #print photon_p0/np.linalg.norm(photon_p0)

    e1, p1 = rocket_to_lab(photon_e0, photon_p0, pi0_v)
    v1 = p1/np.linalg.norm(p1)
    e2, p2 = rocket_to_lab(photon_e0, -photon_p0, pi0_v)
    v2 = p2/np.linalg.norm(p2)

    return (e1/kg_per_MeV, v1), (e2/kg_per_MeV, v2)

if __name__ == '__main__':
    import sys

    npi0s = 10000

    pi0_e = 300.0*kg_per_MeV
    pi0_p = np.sqrt(pi0_e**2 - pi0_mass**2)

    print 'pi0 e: %f MeV' % (pi0_e/kg_per_MeV)
    print 'pi0 p: %f MeV' % (pi0_p/kg_per_MeV)

    e, cos_theta = [], []
    for i, (theta, phi) in enumerate(zip(np.random.rand(npi0s), np.arccos(np.random.uniform(-1.0, 1.0, size=npi0s)))):
        print '\r%i' % (i+1),
        sys.stdout.flush()

        (e1, v1), (e2, v2) = pi0_decay(pi0_e/kg_per_MeV, [0,0,1], theta, phi)
        e += [e1, e2]
        cos_theta += [v1.dot(v2)]

    import matplotlib.pyplot as plt

    plt.figure()
    plt.title('Energy Distribution of Photons from a %.0f MeV Pi0 Decay' % (pi0_e/kg_per_MeV))
    plt.hist(e, 100)
    plt.xlabel('Energy (MeV)')
    plt.figure()
    plt.title('Opening Angle between Photons in Lab from a %.0f MeV Pi0 Decay' % (pi0_e/kg_per_MeV))
    plt.hist(cos_theta, 100)
    plt.xlabel('Cos($\theta$)')
    plt.show()