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.A * self.b.A * ke) / (self.B.A + self.b.A) * np.cos(theta)
        i2 = (self.B.A - self.a.A) / (self.B.A + self.b.A) + (self.a.A * self.b.A) / ((self.B.A + self.b.A) ** 2) * (np.cos(theta) ** 2)
        i3 = self.B.A * self.Q / (self.B.A + self.b.A)
        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.A / self.B.A * ke + self.b.A / self.B.A * eb
        i2 = 2 * np.sqrt(self.a.A * self.b.A * ke * eb) / self.B.A
        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)

theta = np.linspace(0, np.pi, 100)

r1 = Reaction(he_4, ne_22, neutron, mg_25)
res = r1.eb(theta, 1)

plt.plot(theta, res)
plt.title(r"$^{22}$Ne($\alpha$,n)$^{25}$Mg, $T_\alpha$=0.8 MeV")
plt.ylabel(r"$T_n$ (MeV)")
plt.xlabel(r"$\theta$ (rad)")
plt.xticks([0, np.pi / 4, np.pi / 2, np.pi * 3 / 4, np.pi], [r"$0$", r"$\pi/4$", r"$\pi/2$", r"$3\pi/4$", r"$\pi$"])
plt.xlim(0, np.pi)
plt.show()

exit(0)

energy = np.linspace(np.min(res), np.max(res), 100)

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

# res = r2.eb(theta, 0.0114)
# res2 = r2.eB(theta, 0.0114)

res = r2.eb(np.pi / 4, energy)
res2 = r2.eB(np.pi / 4, energy)

fig, ax = plt.subplots()

ax.plot(energy, res, color="C0")
ax.set_title(r"$^3$He(n,p)$^3$H, $\theta$=$\pi/4$")
ax.set_xlabel(r"$T_n$ (MeV)")
ax.set_ylabel(r"$T_p$ (MeV)")

ax1 = ax.twinx()
ax1.plot(energy, res2, color="C1")
ax1.set_ylabel(r"$T_t$ (MeV)")

plt.tight_layout()
plt.show()