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

#[derive(Error, Debug)]
pub enum RecaptureError {
    #[error("trap missing for state {0}")]
    TrapUndefined(String),
}
pub type RecaptureResult<T> = Result<T, RecaptureError>;

/// Perform a series of `N` release-recapture trials for fixed time of flight in
/// the trap for the current state of `atom`, returning the average recapture
/// probability.
pub fn recapture_rng<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(rng);
        q.pos
            += q.mom / atom.mass * tof
            - 0.5 * ThreeVector::from_axis(-grav, Axis::Z) * tof.powi(2);
        if atom.is_trapped(q) { count += 1; }
    }
    return count as f64 / N as f64;
}

/// Perform a series of `N` release-recapture trials for fixed time of flight in
/// the trap for the current state of `atom`, returning the average recapture
/// probability.
pub fn recapture<S, T, R>(atom: &Atom<S, T, R>, tof: f64, N: usize) -> f64
where
    S: State,
    T: Trap,
    R: RadiationPattern,
{
    let mut rng = rnd::thread_rng();
    return recapture_rng(atom, tof, N, &mut rng);
}

/// Perform a series of `N` release-recapture trials for fixed time of flight in
/// the trap for a given atomic state, returning the average recapture
/// probability.
pub fn recapture_for_rng<S, T, R, G>(
    atom: &Atom<S, T, R>,
    state: &S,
    tof: f64,
    N: usize,
    rng: &mut G,
) -> RecaptureResult<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_for_rng(state, rng)
            .ok_or_else(
                || RecaptureError::TrapUndefined(format!("{:?}", state))
            )?;
        q.pos
            += q.mom / atom.mass * tof
            - 0.5 * ThreeVector::from_axis(-grav, Axis::Z) * tof.powi(2);
        if atom.is_trapped_for(state, q)
            .ok_or_else(
                || RecaptureError::TrapUndefined(format!("{:?}", state))
            )?
        {
            count += 1;
        }
    }
    return Ok(count as f64 / N as f64);
}

/// Perform a series of `N` release-recapture trials for fixed time of flight in
/// the trap for a given atomic state, returning the average recapture
/// probability.
pub fn recapture_for<S, T, R>(
    atom: &Atom<S, T, R>,
    state: &S,
    tof: f64,
    N: usize
) -> RecaptureResult<f64>
where
    S: State,
    T: Trap,
    R: RadiationPattern,
{
    let mut rng = rnd::thread_rng();
    return recapture_for_rng(atom, state, tof, N, &mut rng);
}