r"""
Pauli's neutrino: the missing energy in beta decay
=====================================================

By 1930 beta decay looked like it broke energy conservation. If a
nucleus emits only an electron, :math:`A\to B + e^-`, two-body kinematics
gives the electron a single, fixed energy, essentially the whole
:math:`Q`-value. Instead, electrons come out with every energy from zero
up to :math:`Q`. Ellis and Wooster (1927) caught the decays of radium E
(:sup:`210`\ Bi, :math:`Q=1.16` MeV) in a calorimeter and found the heat
per decay matched the spectrum's *mean*, about a third of :math:`Q`, not
its endpoint. The rest of the energy was not being deposited at all.

Pauli's "desperate remedy" (1930) was a third, neutral, nearly massless
particle carrying off the difference. This example generates three-body
decays :math:`A\to B+e^-+\bar\nu` from pure phase space, built from two
chained :func:`~physicskit.particle.decays.two_body_decay` calls, and
checks with :func:`~physicskit.particle.kinematics.invariant_mass` that
energy and momentum balance only once the invisible particle is counted.
"""

# %%
import matplotlib.pyplot as plt
import numpy as np

from physicskit.particle.decays import two_body_decay, two_body_decay_momentum
from physicskit.particle.kinematics import boost_generic, invariant_mass

U = 931.494  # MeV
m_e = 0.511
m_B = 209.9829 * U - 84 * m_e  # Po-210 nucleus
Q = 1.162
M_A = m_B + m_e + Q  # Bi-210 nucleus, fixed by the Q-value

# %%
# If only the electron were emitted
# -------------------------------------
p_two = two_body_decay_momentum(M_A, m_B, m_e)
T_two = np.sqrt(p_two**2 + m_e**2) - m_e
print(f"two-body hypothesis: every electron has T = {T_two:.4f} MeV (Q = {Q} MeV)")

# %%
# Three-body phase space with a neutrino
# ------------------------------------------
# Pick the invariant mass :math:`m_{e\nu}` of the electron-neutrino pair,
# accept it with the phase-space weight :math:`p^*_1 p^*_2`, decay the
# nucleus to :math:`B + (e\nu)`, then decay the pair isotropically in its
# own rest frame and boost back. Pure phase space ignores the nuclear
# Coulomb field and radium E's unusual (forbidden) spectrum shape, so its
# mean electron energy is higher than the measured 0.34 MeV. The point
# stands: the electron gets a variable share, and the neutrino takes the
# rest. (Residuals of order 1e-6 MeV are floating-point rounding on
# nuclear masses near 2e5 MeV.)
rng = np.random.default_rng(1930)


def iso(rng):
    return rng.uniform(-1, 1), rng.uniform(0, 2 * np.pi)


w_max = two_body_decay_momentum(M_A, m_B, m_e) * two_body_decay_momentum(m_e + Q, m_e, 0.0)
events = []
while len(events) < 20000:
    m12 = rng.uniform(m_e, m_e + Q)
    w = two_body_decay_momentum(M_A, m_B, m12) * two_body_decay_momentum(m12, m_e, 0.0)
    if rng.uniform(0, w_max) > w:
        continue
    p_B, p_pair = two_body_decay(M_A, m_B, m12, *iso(rng))
    p_e, p_nu = two_body_decay(m12, m_e, 0.0, *iso(rng))
    beta = p_pair.p_vec / p_pair.E
    events.append((p_B, boost_generic(p_e, beta), boost_generic(p_nu, beta)))

T_e = np.array([e.E - m_e for _, e, _ in events])
E_nu = np.array([nu.E for _, _, nu in events])
T_B = np.array([B.E - m_B for B, _, _ in events])
print(f"\nthree-body: electron T from {T_e.min():.3f} to {T_e.max():.3f} MeV, mean {T_e.mean():.3f} MeV")
print(f"energy per decay seen by a calorimeter (electron + recoil): {np.mean(T_e + T_B):.3f} MeV")
print(f"energy carried off by the neutrino, on average:             {E_nu.mean():.3f} MeV")
print(f"electron + recoil + neutrino, every event: {np.max(np.abs(T_e + T_B + E_nu - Q)):.1e} MeV from Q")

# %%
# The books balance only with the neutrino
# --------------------------------------------
# Without the neutrino the visible products don't reconstruct the parent
# mass, and the electron and recoil nucleus aren't back to back as a
# two-body decay requires.
B0, e0, nu0 = events[0]
print(f"\ninvariant mass of B + e      = M_A - {M_A - invariant_mass([B0, e0]):.4f} MeV")
print(f"invariant mass of B + e + nu = M_A - {M_A - invariant_mass([B0, e0, nu0]):.1e} MeV")
cos_eB = np.array([np.dot(e.p_vec, B.p_vec) / (e.p_mag * B.p_mag) for B, e, _ in events])

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4))
ax1.hist(T_e, bins=50, range=(0, Q), density=True, color="steelblue", alpha=0.8, label="with a neutrino (three-body)")
ax1.axvline(T_two, color="firebrick", lw=3, label="electron alone (two-body)")
ax1.axvline(T_e.mean(), color="k", ls="--", label=f"mean {T_e.mean():.2f} MeV (what a calorimeter sees)")
ax1.set_xlabel("electron kinetic energy [MeV]")
ax1.set_ylabel("probability density")
ax1.set_title(r"$^{210}$Bi $\beta$ decay: a continuous spectrum")
ax1.legend(fontsize=8)
ax2.hist(cos_eB, bins=40, range=(-1, 1), density=True, color="darkorchid", alpha=0.8)
ax2.axvline(-1, color="firebrick", lw=3, label="two-body: always back to back")
ax2.set_xlabel(r"cos(angle between electron and recoil)")
ax2.set_title("Momentum balance needs a third body too")
ax2.legend(fontsize=8)
fig.tight_layout()

plt.show()

# %%
# Check
# -----
# A two-body decay would give every electron T ~ Q; with Pauli's neutrino the
# spectrum is continuous, energy and momentum balance in every event, and on
# average the neutrino carries the missing energy.
assert abs(T_two - Q) < 1e-5
assert T_e.min() < 0.05 * Q and T_e.max() > 0.95 * Q
assert np.max(np.abs(T_e + T_B + E_nu - Q)) < 1e-5
assert abs(M_A - invariant_mass([B0, e0, nu0])) < 1e-5