import numpy as np
from matplotlib import pyplot as plt

mu = 931.4940954  # MeV/c^2


class Particle(object):
    """Define a Particle

    Parameters
    ----------
    A : int
        mass number
    m : float
        mass
    """

    def __init__(self, A, m) -> None:
        self.A = A
        self.m = m


class Reaction(object):
    """A Nuclear Reaction A(a,b)B
    target nuclei (A) is assumed to be quiescent

    Parameters
    ----------
    a/A/b/B : Particle
        Particle a, A, b, B
    """

    def __init__(self, a, A, b, B) -> None:
        self.a: Particle = a
        self.A: Particle = A
        self.b: Particle = b
        self.B: Particle = B

    @property
    def Q(self):
        """Returns Q value of the reaction"""
        return (self.A.m + self.a.m - self.B.m - self.b.m) * mu

    @property
    def threshold(self):
        """Returns threshold energy of the reaction"""
        if self.Q < 0:
            return -(self.A.m + self.a.m) / self.A.m * self.Q
        return 0

    def eb(self, theta, ke):
        """Returns energy of b in the lab frame

        Parameters
        ----------
        theta : float
            angle between a and b in the lab frame in rad
        ke : float
            kinetic energy of a in MeV
        """
        i1 = np.sqrt(self.a.m * self.b.m * ke) / (self.B.m + self.b.m) * np.cos(theta)
        i2 = (self.B.m - self.a.m) / (self.B.m + self.b.m) + (self.a.m * self.b.m) / ((self.B.m + self.b.m) ** 2) * (np.cos(theta) ** 2)
        i3 = self.B.m * self.Q / (self.B.m + self.b.m)
        return (i1 + np.sqrt(i2 * ke + i3)) ** 2

    def eB(self, theta, ke):
        """Returns energy of B in the lab frame

        Parameters
        ----------
        theta : float
            angle between a and b in the lab frame in rad
        ke : float
            kinetic energy of a in MeV
        """
        eb = self.eb(theta, ke)
        i1 = self.a.m / self.B.m * ke + self.b.m / self.B.m * eb
        i2 = 2 * np.sqrt(self.a.m * self.b.m * ke * eb) / self.B.m
        return i1 - i2 * np.cos(theta)


neutron = Particle(1, 1.008665)
proton = Particle(1, 1.007825)
tritium = Particle(3, 3.016049)
he_3 = Particle(3, 3.016029)
he_4 = Particle(4, 4.002603)
ne_22 = Particle(22, 21.991385)
mg_25 = Particle(25, 24.985837)

r1 = Reaction(he_3, neutron, proton, tritium)

# En = 0.1
# theta = np.pi / 2
# Ep = r1.eb(theta, En)
# Et = r1.eB(theta, En)
# alpha = np.arcsin(np.sqrt(proton.m * Ep / (tritium.m * Et)) * np.sin(theta))
# alpha = alpha % (2 * np.pi)

# print(theta / np.pi * 180, alpha / np.pi * 180)
# print(Ep + Et - r1.Q - En)
# print(np.sqrt(proton.m * Ep) * np.sin(theta) - np.sqrt(tritium.m * Et) * np.sin(alpha))
# print(np.sqrt(neutron.m * En))
# print(np.sqrt(proton.m * Ep) * np.cos(theta) + np.sqrt(tritium.m * Et) * np.cos(alpha))

theta = np.random.rand(1000) * np.pi
# energy = np.random.rand(1000) * (0.8 - r1.Q)
energy = 0

Ep = r1.eb(theta, energy)
Et = r1.eB(theta, energy)
alpha = np.arcsin(np.sqrt(proton.m * Ep / (tritium.m * Et)) * np.sin(theta))
alpha = alpha % (2 * np.pi)

data = np.hstack((Ep.reshape(-1, 1), Et.reshape(-1, 1), (theta + alpha).reshape(-1, 1)))
np.savetxt('random_pt.txt', data, fmt='%.6f')

fig = plt.figure()

ax = plt.subplot(2, 1, 1)
ax.plot(theta, Ep, '.')
ax.set_xlabel(r'$\theta$')
ax.set_xticks([0, np.pi / 4, np.pi / 2, 3 * np.pi / 4, np.pi], [r'$0$', r'$\pi/4$', r'$\pi$/2', r'$3\pi/4$', r'$\pi$'])
ax.set_ylabel(r'$E_p$ (MeV)')

ax = plt.subplot(2, 1, 2)
ax.plot(alpha, Et, '.')
ax.set_xlabel(r'$\alpha$')
ax.set_xticks([0, np.pi / 4, np.pi / 2, 3 * np.pi / 4, np.pi], [r'$0$', r'$\pi/4$', r'$\pi$/2', r'$3\pi/4$', r'$\pi$'])
ax.set_ylabel(r'$E_t$ (MeV)')

plt.tight_layout()
plt.subplots_adjust()
plt.show()