use rand::Rng;
use thiserror::Error;
use crate::{
    atom::{
        Atom,
        AtomIter,
        RadiationPattern,
        State,
    },
    newton::{
        NewtonError,
        PhaseSpace,
        rka,
    },
    trap::Trap,
};

#[derive(Error, Debug)]
pub enum ScatterError {
    #[error("newton error {0}")]
    NewtonError(NewtonError),
}
pub type ScatterResult<T> = Result<T, ScatterError>;

#[derive(Clone, Debug)]
pub struct ScatterHistory<S>
where S: State,
{
    pub state: Vec<S>,
    pub time: Vec<f64>,
    pub trajectory: Vec<PhaseSpace>,
}

/// Calculate a single trajectory up to `t_final` *or* until the atom goes dark;
/// i.e. the final time is not guaranteed to be equal to `t_final`.
pub fn scatter_sim<S, T, R, G>(
    atom: &Atom<S, T, R>,
    t_final: f64,
    epsilon: f64,
    rng: &mut G,
) -> ScatterResult<(ScatterHistory<S>, Atom<S, T, R>)>
where
    S: State,
    T: Trap,
    R: RadiationPattern,
    G: Rng + ?Sized,
{
    let m: f64 = atom.mass;
    let mut state_history: Vec<S> = vec![atom.state];
    let mut t_cur: f64 = 0.0;
    let mut t: Vec<f64> = Vec::new();
    let (mut atom_iter, mut q_cur): (AtomIter<S, T, R>, PhaseSpace)
        = atom.to_state_iter();
    let mut trajectory: Vec<PhaseSpace> = Vec::new();
    // simulate until t_final or the atom goes dark
    while let Some((state, photon)) = atom_iter.next() {
        state_history.push(state);
        match photon.momentum_kick_lifetime(rng) {
            // absorb a photon
            (kick, None) => {
                q_cur.mom += kick;
            },
            // integrate position until photon emission
            (kick, Some((lifetime, tau))) => {
                let dt: f64
                    = atom_iter.get_trap().period(m)
                    .min(tau)
                    .min(lifetime)
                    .min(t_final - t_cur)
                    / 10.0;
                let t_int_end: f64 = (t_cur + lifetime).min(t_final);
                let rhs
                    = |_t: f64, q: PhaseSpace| -> PhaseSpace {
                        PhaseSpace {
                            pos: q.mom / m,
                            mom: -atom_iter.get_trap().gradient(q.pos),
                        }
                    };
                let (mut t_int, mut q_int): (Vec<f64>, Vec<PhaseSpace>)
                    = rka((t_cur, t_int_end), q_cur, dt, rhs, epsilon)
                    .map_err(ScatterError::NewtonError)?;
                t_cur = t_int.pop().unwrap();
                q_cur = q_int.pop().unwrap();
                t.append(&mut t_int);
                trajectory.append(&mut q_int);
                q_cur.mom += kick;
            },
        }
        if t_cur >= t_final { break; }
    }
    return Ok((
        ScatterHistory {
            state: state_history,
            time: t,
            trajectory,
        },
        atom_iter.dump(),
    ));
}

/// Calculate a single trajectory, recording times and positions/momenta only
/// when photons are emitted. Necessarily limited to when the atom goes dark.
pub fn scatter_sim_emission<S, T, R, G>(
    atom: &Atom<S, T, R>,
    t_final: f64,
    epsilon: f64,
    rng: &mut G,
) -> ScatterResult<(ScatterHistory<S>, Atom<S, T, R>)>
where
    S: State,
    T: Trap,
    R: RadiationPattern,
    G: Rng + ?Sized,
{
    let m: f64 = atom.mass;
    let mut state_history: Vec<S> = vec![atom.state];
    let mut t_cur: f64 = 0.0;
    let mut t: Vec<f64> = vec![0.0];
    let (mut atom_iter, mut q_cur): (AtomIter<S, T, R>, PhaseSpace)
        = atom.to_state_iter();
    let mut trajectory: Vec<PhaseSpace> = vec![q_cur];
    while let Some((state, photon)) = atom_iter.next() {
        state_history.push(state);
        match photon.momentum_kick_lifetime(rng) {
            (kick, None) => {
                q_cur.mom += kick;
            },
            (kick, Some((lifetime, tau))) => {
                let dt: f64
                    = atom_iter.get_trap().period(m)
                    .min(tau)
                    .min(lifetime)
                    .min(t_final - t_cur)
                    / 10.0;
                let t_int_end: f64 = (t_cur + lifetime).min(t_final);
                let rhs
                    = |_t: f64, q: PhaseSpace| -> PhaseSpace {
                        PhaseSpace {
                            pos: q.mom / m,
                            mom: -atom_iter.get_trap().gradient(q.pos),
                        }
                    };
                let (mut t_int, mut q_int): (Vec<f64>, Vec<PhaseSpace>)
                    = rka((t_cur, t_int_end), q_cur, dt, rhs, epsilon)
                    .map_err(ScatterError::NewtonError)?;
                t_cur = t_int.pop().unwrap();
                q_cur = q_int.pop().unwrap();
                t.push(t_cur);
                trajectory.push(q_cur);
                q_cur.mom += kick;
            },
        }
        if t_cur >= t_final { break; }
    }
    return Ok((
        ScatterHistory {
            state: state_history,
            time: t,
            trajectory,
        },
        atom_iter.dump(),
    ));
}

/// Calculate a single trajectory up to `t_final`, continuing past the point at
/// which the atom goes dark.
pub fn scatter_sim_dark<S, T, R, G>(
    atom: &Atom<S, T, R>,
    t_final: f64,
    epsilon: f64,
    rng: &mut G,
) -> ScatterResult<(ScatterHistory<S>, Atom<S, T, R>)>
where
    S: State,
    T: Trap,
    R: RadiationPattern,
    G: Rng + ?Sized,
{
    let (mut scatter_history, atom_sim): (ScatterHistory<S>, Atom<S, T, R>)
        = scatter_sim(atom, t_final, epsilon, rng)?;
    if scatter_history.time.last().unwrap_or(&0.0) < &t_final {
        let t_init: f64 = scatter_history.time.pop().unwrap_or(0.0);
        let q_init: PhaseSpace
            = scatter_history.trajectory.pop()
            .unwrap_or_else(|| atom_sim.sample_phasespace(rng));
        let dt: f64
            = atom_sim.get_trap().period(atom_sim.mass)
            .min(t_final - t_init)
            / 10.0;
        let rhs
            = |_t: f64, q: PhaseSpace| -> PhaseSpace {
                PhaseSpace {
                    pos: q.mom / atom_sim.mass,
                    mom: -atom_sim.get_trap().gradient(q.pos),
                }
            };
        let (mut t_int, mut q_int): (Vec<f64>, Vec<PhaseSpace>)
            = rka((t_init, t_final), q_init, dt, rhs, epsilon)
                .map_err(ScatterError::NewtonError)?;
        scatter_history.time.append(&mut t_int);
        scatter_history.trajectory.append(&mut q_int);
    }
    return Ok((scatter_history, atom_sim));
}