#include <TLorentzVector.h>
#include <TMath.h>
#include <TVector3.h>
#include <iostream>
const double u_MeV = 931.49410372; // u in MeV

class Kinematics {
private:
    double m_A, m_d, m_p, m_B;         // Ground state masses in amu
    double E_beam;                     // Total kinetic energy of beam (MeV)
    double Q0;                         // Ground state Q-value

    // Cached outputs from the last calculation
    double fExc = -9999;
    double fBeta4 = -9999;
    double fTheta4 = -9999;
    double fP4 = -9999;

public:

    Kinematics() : m_A(0), m_d(0), m_p(0), m_B(0), E_beam(0), Q0(0) {}
    Kinematics(const Kinematics&) = default;
	Kinematics& operator=(const Kinematics&) = default;

    Kinematics(double m1, double m2, double m3, double m4, double ebeam_per_u) {
        setValues(m1, m2, m3, m4, ebeam_per_u);
    }

    void setValues(double m1, double m2, double m3, double m4, double ebeam_per_u) {
        m_A = m1; m_d = m2; m_p = m3; m_B = m4;
        E_beam = ebeam_per_u * m_A;
        Q0 = (m_A + m_d - m_p - m_B) * u_MeV;
    }

    void computeKinematics(double t3, double angle3_deg) {
        // 1. Convert ground state masses to energy units (MeV)
        double m1_MeV = m_A * u_MeV;
        double m2_MeV = m_d * u_MeV;
        double m3_MeV = m_p * u_MeV;

        // 2. Build Four-Vectors for Initial State
        double p1_mag = TMath::Sqrt(E_beam * E_beam + 2.0 * m1_MeV * E_beam);
        TLorentzVector P1(0, 0, p1_mag, m1_MeV + E_beam); // Beam along Z-axis
        TLorentzVector P2(0, 0, 0, m2_MeV);              // Target at rest

        // 3. Build Four-Vector for Detected Ejectile (Particle 3)
        double theta3_rad = angle3_deg * TMath::Pi() / 180.0;
        double p3_mag = TMath::Sqrt(t3 * t3 + 2.0 * m3_MeV * t3);
        
        TVector3 p3_vec(0, 0, 1);
        p3_vec.SetTheta(theta3_rad);
        p3_vec.SetMag(p3_mag);
        
        TLorentzVector P3(p3_vec, m3_MeV + t3);

        // 4. Solve for Recoil (Particle 4) using Four-Momentum Conservation
        TLorentzVector P4 = P1 + P2 - P3;

        // 5. Extract Recoil physical properties
        double m4_dynamic = P4.M(); // Invariant mass of the excited recoil nucleus
        fExc = m4_dynamic - (m_B * u_MeV); 
        fBeta4 = P4.Beta();
        fTheta4 = P4.Theta() * 180.0 / TMath::Pi();
        fP4 = P4.P();
    }

    // Getters that trigger computation
    double getExc(double t3, double angle3)   { computeKinematics(t3, angle3); return fExc; }
    double getBeta4(double t3, double angle3) { computeKinematics(t3, angle3); return fBeta4; }
    double getTheta4(double t3, double angle3){ computeKinematics(t3, angle3); return fTheta4; }
    double getBrho(double t3, double angle3, double charge_state) {
        computeKinematics(t3, angle3);
        return fP4 * 3.3359e-3 / charge_state;
    }
};

// Simple test function to run in ROOT
void KinematicsMacro() {
    // Example: d(16O, p)17O reaction with a 10 MeV/u 16O beam
    // Masses roughly: 16O (~15.995), d (~2.014), p (~1.007), 17O (~16.999)
    Kinematics k(15.9949, 2.0141, 1.0078, 16.9991, 10.0);
    
    double t3 = 12.0;       // Ejectile proton energy in MeV
    double angle3 = 45.0;   // Proton lab angle in degrees
    
    std::cout << "--- Kinematics Results ---" << std::endl;
    std::cout << "Excitation Energy: " << k.getExc(t3, angle3) << " MeV" << std::endl;
    std::cout << "Recoil Beta:        " << k.getBeta4(t3, angle3) << std::endl;
    std::cout << "Recoil Angle:       " << k.getTheta4(t3, angle3) << " deg" << std::endl;
}