#include "math.h"
#include "stdlib.h"

#include "iostream"

const double pi = 3.1415926535;

double mass[4];
double vbeta, vgamma;
double t3_cm;
double t4_cm;
double p3_cm;
double p4_cm;

void SetBeamEnergy(double energy) {
    double umass = 931.494;
    double massHe4 = 4.00260325413 * umass;
    double massNe22 = 21.991385114 * umass;
    double massMg25 = 24.98583697 * umass;
    double massn = 939.56563;

    mass[0] = massHe4;
    mass[1] = massNe22;
    mass[2] = massMg25;
    mass[3] = massn;

    double eb = energy + mass[0];
    double pb2 = energy * energy + 2.0 * energy * mass[0];
    double pb = std::sqrt(pb2);
    double esum = eb + mass[1];
    vbeta = pb / esum;
    vgamma = 1.0 / std::sqrt(1.0 - vbeta * vbeta);
    double e_cm2 = esum * esum - pb2;
    double e_cm = std::sqrt(e_cm2);
    double t_cm = e_cm - mass[2] - mass[3];
    if (t_cm < 0.0) {
        std::cout << "No solution!" << std::endl;
        return;
    }
    double t_cm2 = t_cm * t_cm;
    t3_cm = (t_cm2 + 2.0 * mass[3] * t_cm) / e_cm / 2.0;
    t4_cm = (t_cm2 + 2.0 * mass[2] * t_cm) / e_cm / 2.0;
    double p3_cm2 = t3_cm * t3_cm + 2.0 * t3_cm * mass[2];
    double p4_cm2 = t4_cm * t4_cm + 2.0 * t4_cm * mass[3];
    p3_cm = std::sqrt(p3_cm2);
    p4_cm = std::sqrt(p4_cm2);
}

void GeneratePrimaries(double thetacmr) {
    double thetalabr;

    double tg_thetalabr =
        p4_cm * std::sin(thetacmr) /
        (vgamma * (p4_cm * std::cos(thetacmr) + vbeta * std::sqrt(p4_cm * p4_cm + mass[3] * mass[3])));
    if (tg_thetalabr >= 1.0e6)
        thetalabr = pi / 2.0;
    else
        thetalabr = std::atan(tg_thetalabr);
    if (thetalabr < 0.0) thetalabr = pi + thetalabr;

    double p4_cmx = p4_cm * std::sin(thetacmr);
    double p4_cmz = p4_cm * std::cos(thetacmr);
    double p4_labx = p4_cmx;
    double p4_labz = vgamma * (p4_cmz + vbeta * (t4_cm + mass[3]));
    double p4_lab = std::sqrt(p4_labx * p4_labx + p4_labz * p4_labz);
    double e4_lab = sqrt(p4_lab * p4_lab + mass[3] * mass[3]) - mass[3];

    // std::cout << thetalabr << ", " << e4_lab << std::endl;
    std::cout << e4_lab << std::endl;
}

int main() {
    SetBeamEnergy(1);
    for (int i = 0; i < 100; i++) GeneratePrimaries(i * pi / 100.0);
    return 0;
}