use rand::Rng;
use crate::{
    atom::{
        Atom,
        State,
        RadiationPattern,
    },
    newton::{
        ThreeVector,
        PhaseSpace,
        Axis,
    },
    phys::g as grav,
    trap::Trap,
};

pub fn recapture<S, T, R, G>(
    atom: &Atom<S, T, R>,
    tof: f64,
    N: usize,
    rng: &mut G,
) -> f64
where
    S: State,
    T: Trap,
    R: RadiationPattern,
    G: Rng + ?Sized,
{
    let mut count: usize = 0;
    let mut q: PhaseSpace;
    for _ in 0..N {
        q = atom.sample_phasespace(rng);
        q.pos
            += q.mom / atom.mass * tof
            - 0.5 * ThreeVector::from_axis(-grav, Axis::Z) * tof.powi(2);
        if atom.get_trap().is_trapped(atom.mass, q) {
            count += 1;
        }
    }
    return count as f64 / N as f64;
}