use crate::{factor::VarStore, gaussian::Gaussian};

/// Reusable scratch buffers for `Game::likelihoods`.
///
/// A `TimeSlice` owns one arena; all events in the slice share it across
/// the convergence iterations. All Vecs are cleared (not dropped) on
/// `reset()` so their heap capacity is reused across games.
#[derive(Debug, Default)]
pub struct ScratchArena {
    pub(crate) vars: VarStore,
    pub(crate) sort_buf: Vec<usize>,
    pub(crate) inv_buf: Vec<usize>,
    pub(crate) team_prior: Vec<Gaussian>,
    pub(crate) lhood_lose: Vec<Gaussian>,
    pub(crate) lhood_win: Vec<Gaussian>,
}

impl ScratchArena {
    pub fn new() -> Self {
        Self::default()
    }

    #[inline]
    pub(crate) fn reset(&mut self) {
        self.vars.clear();
        self.sort_buf.clear();
        self.inv_buf.clear();
        self.team_prior.clear();
        self.lhood_lose.clear();
        self.lhood_win.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{N_INF, gaussian::Gaussian};

    #[test]
    fn reset_keeps_capacity() {
        let mut arena = ScratchArena::new();
        arena.vars.alloc(N_INF);
        arena.sort_buf.push(42);
        arena.team_prior.push(Gaussian::from_ms(0.0, 1.0));
        let var_cap = arena.vars.marginals.capacity();
        let sort_cap = arena.sort_buf.capacity();
        let prior_cap = arena.team_prior.capacity();
        arena.reset();
        assert_eq!(arena.vars.len(), 0);
        assert_eq!(arena.sort_buf.len(), 0);
        assert_eq!(arena.team_prior.len(), 0);
        assert_eq!(arena.vars.marginals.capacity(), var_cap);
        assert_eq!(arena.sort_buf.capacity(), sort_cap);
        assert_eq!(arena.team_prior.capacity(), prior_cap);
    }
}