trueskill-tt/src/time_slice.rs

//! A single time step's worth of events.
//!
//! Renamed from `Batch` in T2.

use std::collections::HashMap;

use crate::{
    Index, N_INF,
    arena::ScratchArena,
    color_group::ColorGroups,
    drift::Drift,
    game::Game,
    gaussian::Gaussian,
    rating::Rating,
    storage::{CompetitorStore, SkillStore},
    time::Time,
};

#[derive(Debug)]
pub(crate) struct Skill {
    pub(crate) forward: Gaussian,
    backward: Gaussian,
    likelihood: Gaussian,
    pub(crate) elapsed: i64,
}

impl Skill {
    pub(crate) fn posterior(&self) -> Gaussian {
        self.likelihood * self.backward * self.forward
    }
}

impl Default for Skill {
    fn default() -> Self {
        Self {
            forward: N_INF,
            backward: N_INF,
            likelihood: N_INF,
            elapsed: 0,
        }
    }
}

#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub enum EventKind {
    Ranked,
    Scored { score_sigma: f64 },
}

#[derive(Clone, Debug)]
struct Item {
    agent: Index,
    /// This competitor's slot in the owning slice's `SkillStore`, resolved
    /// once at ingestion.
    ///
    /// The convergence loop reaches skills through this rather than through
    /// `agent`, which is what keeps `HashMap` hashing out of the hot path now
    /// that the store is compact rather than indexed by the global `Index`.
    slot: u32,
    likelihood: Gaussian,
}

impl Item {
    fn within_prior<T: Time, D: Drift<T>>(
        &self,
        forward: bool,
        skills: &SkillStore,
        agents: &CompetitorStore<T, D>,
    ) -> Rating<T, D> {
        let r = &agents[self.agent].rating;
        let skill = skills.at(self.slot);

        if forward {
            Rating::new(skill.forward, r.beta, r.drift).with_drift_scale(r.drift_scale)
        } else {
            Rating::new(skill.posterior() / self.likelihood, r.beta, r.drift)
                .with_drift_scale(r.drift_scale)
        }
    }
}

#[derive(Clone, Debug)]
struct Team {
    items: Vec<Item>,
    output: f64,
}

#[derive(Clone, Debug)]
pub(crate) struct Event {
    teams: Vec<Team>,
    log_evidence: f64,
    weights: Vec<Vec<f64>>,
    kind: EventKind,
}

impl Event {
    pub(crate) fn iter_agents(&self) -> impl Iterator<Item = Index> + '_ {
        self.teams
            .iter()
            .flat_map(|t| t.items.iter().map(|it| it.agent))
    }

    fn outputs(&self) -> Vec<f64> {
        self.teams
            .iter()
            .map(|team| team.output)
            .collect::<Vec<_>>()
    }

    pub(crate) fn within_priors<T: Time, D: Drift<T>>(
        &self,
        forward: bool,
        skills: &SkillStore,
        agents: &CompetitorStore<T, D>,
    ) -> Vec<Vec<Rating<T, D>>> {
        self.teams
            .iter()
            .map(|team| {
                team.items
                    .iter()
                    .map(|item| item.within_prior(forward, skills, agents))
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }

    /// Run inference for this event and return its per-item likelihoods.
    ///
    /// Reads `skills` immutably and does not touch `self`, so every event in
    /// a color group can run concurrently without any aliasing question —
    /// the mutation is deferred to `apply`.
    fn compute<T: Time, D: Drift<T>>(
        &self,
        skills: &SkillStore,
        agents: &CompetitorStore<T, D>,
        p_draw: f64,
        convergence: crate::ConvergenceOptions,
        arena: &mut ScratchArena,
    ) -> EventUpdate {
        let teams = self.within_priors(false, skills, agents);
        let result = self.outputs();
        let g = match self.kind {
            EventKind::Ranked => {
                Game::ranked_with_arena(teams, &result, &self.weights, p_draw, convergence, arena)
            }
            EventKind::Scored { score_sigma } => Game::scored_with_arena(
                teams,
                &result,
                &self.weights,
                score_sigma,
                convergence,
                arena,
            ),
        };

        EventUpdate {
            log_evidence: g.log_evidence,
            likelihoods: g.likelihoods,
        }
    }

    /// Fold a computed update into the skill store and cache it on the items.
    fn apply(&mut self, skills: &mut SkillStore, update: EventUpdate) {
        for (t, team) in self.teams.iter_mut().enumerate() {
            for (i, item) in team.items.iter_mut().enumerate() {
                let fresh = update.likelihoods[t][i];
                let old_likelihood = skills.at(item.slot).likelihood;
                let new_likelihood = (old_likelihood / item.likelihood) * fresh;
                skills.at_mut(item.slot).likelihood = new_likelihood;
                item.likelihood = fresh;
            }
        }

        self.log_evidence = update.log_evidence;
    }

    /// Compute and apply in one step — the sequential sweep.
    fn iteration_direct<T: Time, D: Drift<T>>(
        &mut self,
        skills: &mut SkillStore,
        agents: &CompetitorStore<T, D>,
        p_draw: f64,
        convergence: crate::ConvergenceOptions,
        arena: &mut ScratchArena,
    ) {
        let update = self.compute(skills, agents, p_draw, convergence, arena);
        self.apply(skills, update);
    }
}

/// The result of running inference for one event, before it is folded back
/// into the shared skill store.
#[derive(Debug)]
struct EventUpdate {
    log_evidence: f64,
    likelihoods: Vec<Vec<Gaussian>>,
}

/// One slice's worth of forward-only inference.
///
/// `posteriors` doubles as the outgoing forward message: the scratch sweep
/// never writes `backward`, so it stays `N_INF`, and `Skill::posterior()`
/// and `forward_prior_out` are then the same product.
#[derive(Debug)]
pub(crate) struct FilteredStep {
    pub(crate) log_evidence: f64,
    pub(crate) posteriors: Vec<(Index, Gaussian)>,
}

#[derive(Debug)]
pub struct TimeSlice<T: Time = i64> {
    pub(crate) events: Vec<Event>,
    pub(crate) skills: SkillStore,
    pub(crate) time: T,
    p_draw: f64,
    pub(crate) convergence: crate::ConvergenceOptions,
    arena: ScratchArena,
    pub(crate) color_groups: ColorGroups,
    /// Whether `color_groups` still reflects `events`.
    ///
    /// Coloring is rebuilt lazily, on the first full sweep after an append,
    /// rather than eagerly per append: the partition is thrown away and
    /// recomputed wholesale either way, so doing it per append made ingesting
    /// n events O(n^2) with no benefit — nothing reads the partition between
    /// an append and the next full sweep.
    color_groups_dirty: bool,
}

impl<T: Time> TimeSlice<T> {
    pub fn new(time: T, p_draw: f64, convergence: crate::ConvergenceOptions) -> Self {
        Self {
            events: Vec::new(),
            skills: SkillStore::new(),
            time,
            p_draw,
            convergence,
            arena: ScratchArena::new(),
            color_groups: ColorGroups::new(),
            color_groups_dirty: false,
        }
    }

    /// Recompute the color-group partition and reorder `self.events` into
    /// color-contiguous ranges. After this call, `self.color_groups.groups[c]`
    /// contains a contiguous ascending range of indices in `self.events`.
    pub(crate) fn recompute_color_groups(&mut self) {
        use crate::color_group::color_greedy;

        let n = self.events.len();
        if n == 0 {
            self.color_groups = ColorGroups::new();
            self.color_groups_dirty = false;
            return;
        }

        let cg = color_greedy(n, |ev_idx| {
            self.events[ev_idx].iter_agents().collect::<Vec<_>>()
        });

        let mut reordered: Vec<Event> = Vec::with_capacity(n);
        let mut new_groups: Vec<Vec<usize>> = Vec::with_capacity(cg.groups.len());
        let mut taken: Vec<Option<Event>> = self.events.drain(..).map(Some).collect();

        for group in &cg.groups {
            let mut new_indices: Vec<usize> = Vec::with_capacity(group.len());
            for &old_idx in group {
                let ev = taken[old_idx].take().expect("event already taken");
                new_indices.push(reordered.len());
                reordered.push(ev);
            }
            new_groups.push(new_indices);
        }

        self.events = reordered;
        self.color_groups = ColorGroups { groups: new_groups };
        self.color_groups_dirty = false;

        debug_assert!(
            self.color_groups.groups_are_contiguous(),
            "color groups must occupy contiguous event ranges"
        );
    }

    pub fn add_events<D: Drift<T>>(
        &mut self,
        composition: Vec<Vec<Vec<Index>>>,
        results: Option<Vec<Vec<f64>>>,
        weights: Option<Vec<Vec<Vec<f64>>>>,
        kinds: Vec<EventKind>,
        agents: &CompetitorStore<T, D>,
    ) {
        let mut unique = Vec::with_capacity(10);

        let this_agent = composition.iter().flatten().flatten().filter(|idx| {
            if !unique.contains(idx) {
                unique.push(*idx);

                return true;
            }

            false
        });

        for idx in this_agent {
            let elapsed = compute_elapsed(agents[*idx].last_time.as_ref(), &self.time);

            let forward = agents[*idx].receive(&self.time);

            if let Some(skill) = self.skills.get_mut(*idx) {
                skill.elapsed = elapsed;
                skill.forward = forward;
            } else {
                self.skills.insert(
                    *idx,
                    Skill {
                        forward,
                        backward: N_INF,
                        likelihood: N_INF,
                        elapsed,
                    },
                );
            }
        }

        let skills = &self.skills;

        let events = composition.iter().enumerate().map(|(e, event)| {
            let teams = event
                .iter()
                .enumerate()
                .map(|(t, team)| {
                    let items = team
                        .iter()
                        .map(|&agent| Item {
                            agent,
                            // Every participant was inserted into `skills`
                            // just above, so the slot always resolves.
                            slot: skills
                                .slot_of(agent)
                                .expect("participant must be present in the slice store"),
                            likelihood: N_INF,
                        })
                        .collect::<Vec<_>>();

                    Team {
                        items,
                        output: match &results {
                            Some(results) => results[e][t],
                            // No explicit result: rank by position, first team best.
                            None => (event.len() - (t + 1)) as f64,
                        },
                    }
                })
                .collect::<Vec<_>>();

            let weights = match &weights {
                Some(weights) => weights[e].clone(),
                None => teams
                    .iter()
                    .map(|team| vec![1.0; team.items.len()])
                    .collect::<Vec<_>>(),
            };

            Event {
                teams,
                log_evidence: 0.0,
                weights,
                kind: kinds[e],
            }
        });

        let from = self.events.len();

        self.events.extend(events);

        self.color_groups_dirty = true;

        self.iteration(from, agents);
    }

    pub(crate) fn posteriors(&self) -> HashMap<Index, Gaussian> {
        self.skills
            .iter()
            .map(|(idx, skill)| (idx, skill.posterior()))
            .collect::<HashMap<_, _>>()
    }

    /// Sweep this slice's events once, starting at index `from`.
    ///
    /// # Panics
    ///
    /// Panics if an event references a competitor with no entry in this
    /// slice's skill store. `add_events` inserts one for every participant, so
    /// this cannot happen for slices built through the public API.
    pub fn iteration<D: Drift<T>>(&mut self, from: usize, agents: &CompetitorStore<T, D>) {
        if from == 0 && self.color_groups_dirty {
            self.recompute_color_groups();
        }

        if from > 0 || self.color_groups.is_empty() {
            // Initial pass (add_events) or no color groups yet: simple sequential sweep.
            for event in self.events.iter_mut().skip(from) {
                let teams = event.within_priors(false, &self.skills, agents);
                let result = event.outputs();

                let g = match event.kind {
                    EventKind::Ranked => Game::ranked_with_arena(
                        teams,
                        &result,
                        &event.weights,
                        self.p_draw,
                        self.convergence,
                        &mut self.arena,
                    ),
                    EventKind::Scored { score_sigma } => Game::scored_with_arena(
                        teams,
                        &result,
                        &event.weights,
                        score_sigma,
                        self.convergence,
                        &mut self.arena,
                    ),
                };

                for (t, team) in event.teams.iter_mut().enumerate() {
                    for (i, item) in team.items.iter_mut().enumerate() {
                        let old_likelihood = self.skills.at(item.slot).likelihood;
                        let new_likelihood =
                            (old_likelihood / item.likelihood) * g.likelihoods[t][i];
                        self.skills.at_mut(item.slot).likelihood = new_likelihood;
                        item.likelihood = g.likelihoods[t][i];
                    }
                }

                event.log_evidence = g.log_evidence;
            }
        } else {
            self.sweep_color_groups(agents);
        }
    }

    /// Full event sweep using the color-group partition. Colors are processed
    /// sequentially; within each color the inner loop is parallel under rayon.
    ///
    /// Events in one color group touch disjoint agent sets, so none of them
    /// can observe another's writes. That makes the sweep separable: inference
    /// runs concurrently over shared `&self.skills`, and the resulting updates
    /// are folded in afterwards in index order. Splitting it this way needs no
    /// `unsafe` and no aliasing argument, and it keeps results bit-identical
    /// across thread counts because the apply order does not depend on which
    /// worker finished first.
    #[cfg(feature = "rayon")]
    fn sweep_color_groups<D: Drift<T>>(&mut self, agents: &CompetitorStore<T, D>) {
        use rayon::prelude::*;

        thread_local! {
            static ARENA: std::cell::RefCell<ScratchArena> =
                std::cell::RefCell::new(ScratchArena::new());
        }

        // Minimum color-group size to justify rayon's task-spawn overhead.
        // Below this threshold, process events sequentially to avoid regression
        // on small per-slice workloads.
        const RAYON_THRESHOLD: usize = 64;

        for color_idx in 0..self.color_groups.groups.len() {
            let group_len = self.color_groups.groups[color_idx].len();
            if group_len == 0 {
                continue;
            }

            let range = self.color_groups.color_range(color_idx);
            let p_draw = self.p_draw;
            let convergence = self.convergence;

            if group_len >= RAYON_THRESHOLD {
                let skills = &self.skills;
                let updates: Vec<EventUpdate> = self.events[range.clone()]
                    .par_iter()
                    .map(|ev| {
                        ARENA.with(|cell| {
                            let mut arena = cell.borrow_mut();
                            arena.reset();

                            ev.compute(skills, agents, p_draw, convergence, &mut arena)
                        })
                    })
                    .collect();

                for (ev, update) in self.events[range].iter_mut().zip(updates) {
                    ev.apply(&mut self.skills, update);
                }
            } else {
                for ev in &mut self.events[range] {
                    ev.iteration_direct(
                        &mut self.skills,
                        agents,
                        p_draw,
                        self.convergence,
                        &mut self.arena,
                    );
                }
            }
        }
    }

    /// Full event sweep using the color-group partition, sequential direct-write path.
    /// Events within each color group are updated inline — no EventOutput allocation —
    /// matching the T2 performance profile.
    #[cfg(not(feature = "rayon"))]
    fn sweep_color_groups<D: Drift<T>>(&mut self, agents: &CompetitorStore<T, D>) {
        for color_idx in 0..self.color_groups.groups.len() {
            if self.color_groups.groups[color_idx].is_empty() {
                continue;
            }
            let range = self.color_groups.color_range(color_idx);

            // Borrow self.events as a mutable slice for this color range.
            // self.skills and self.arena are separate fields — disjoint borrows are
            // allowed within a single method body.
            let p_draw = self.p_draw;
            for ev in &mut self.events[range] {
                ev.iteration_direct(
                    &mut self.skills,
                    agents,
                    p_draw,
                    self.convergence,
                    &mut self.arena,
                );
            }
        }
    }

    /// Iterate this slice alone until its posteriors stop moving, returning
    /// the number of iterations taken.
    ///
    /// Used by `filtered_step` to drive a scratch copy of the slice, and by
    /// tests. Production convergence across slices is driven by
    /// `History::converge`, which calls `iteration` directly.
    ///
    /// Honours `self.convergence`; it previously hard-coded an epsilon and a
    /// 20-iteration cap that matched neither `ConvergenceOptions` nor the
    /// schedule default.
    pub(crate) fn iterate_to_convergence<D: Drift<T>>(
        &mut self,
        agents: &CompetitorStore<T, D>,
    ) -> usize {
        use crate::{tuple_gt, tuple_max};

        let epsilon = self.convergence.epsilon;
        let max_iter = self.convergence.max_iter;

        let mut step = (f64::INFINITY, f64::INFINITY);
        let mut i = 0;

        while tuple_gt(step, epsilon) && i < max_iter {
            let old = self.posteriors();

            self.iteration(0, agents);

            let new = self.posteriors();

            step = old.iter().fold((0.0, 0.0), |step, (a, old)| {
                tuple_max(step, old.delta(new[a]))
            });

            i += 1;

            if !crate::step_is_finite(step) {
                break;
            }
        }

        i
    }

    pub(crate) fn forward_prior_out(&self, agent: &Index) -> Gaussian {
        let skill = self.skills.get(*agent).unwrap();
        skill.forward * skill.likelihood
    }

    pub(crate) fn backward_prior_out<D: Drift<T>>(
        &self,
        agent: &Index,
        agents: &CompetitorStore<T, D>,
    ) -> Gaussian {
        let skill = self.skills.get(*agent).unwrap();
        let n = skill.likelihood * skill.backward;
        n.forget(
            agents[*agent]
                .rating
                .drift_variance_for_elapsed(skill.elapsed),
        )
    }

    pub(crate) fn new_backward_info<D: Drift<T>>(&mut self, agents: &CompetitorStore<T, D>) {
        for (agent, skill) in self.skills.iter_mut() {
            skill.backward = agents[agent].message.unwrap_or(N_INF);
        }
        self.iteration(0, agents);
    }

    pub(crate) fn new_forward_info<D: Drift<T>>(&mut self, agents: &CompetitorStore<T, D>) {
        for (agent, skill) in self.skills.iter_mut() {
            skill.forward = agents[agent].receive_for_elapsed(skill.elapsed);
        }
        self.iteration(0, agents);
    }

    /// Run this slice's events on forward (filtering) information alone.
    ///
    /// `incoming` holds each competitor's forward message out of their
    /// previous appearance; a competitor absent from it starts at their
    /// configured prior. The sweep runs on a scratch copy, so the real slice
    /// is untouched — which is what makes the filtered estimates independent
    /// of whether `History::converge` has run.
    pub(crate) fn filtered_step<D: Drift<T>>(
        &self,
        incoming: &HashMap<Index, Gaussian>,
        agents: &CompetitorStore<T, D>,
    ) -> FilteredStep {
        let mut scratch = TimeSlice {
            events: self.events.clone(),
            skills: SkillStore::new(),
            time: self.time,
            p_draw: self.p_draw,
            convergence: self.convergence,
            arena: ScratchArena::new(),
            color_groups: ColorGroups::new(),
            color_groups_dirty: true,
        };

        for event in &mut scratch.events {
            for team in &mut event.teams {
                for item in &mut team.items {
                    item.likelihood = N_INF;
                }
            }

            event.log_evidence = 0.0;
        }

        for (agent, skill) in self.skills.iter() {
            let rating = &agents[agent].rating;

            let forward = match incoming.get(&agent) {
                Some(message) => message.forget(rating.drift_variance_for_elapsed(skill.elapsed)),
                None => rating.prior,
            };

            let slot = scratch.skills.insert(
                agent,
                Skill {
                    forward,
                    backward: N_INF,
                    likelihood: N_INF,
                    elapsed: skill.elapsed,
                },
            );

            // The cloned events carry slots resolved against the REAL store, so
            // the scratch must assign the same ones. It does because `iter()`
            // yields slot order and `insert` allocates slots in call order —
            // but that is a coupling between two types, so pin it here rather
            // than leave it to be rediscovered after it breaks.
            debug_assert_eq!(
                Some(slot),
                self.skills.slot_of(agent),
                "scratch slot must match the real slice's slot for {agent:?}"
            );
        }

        scratch.iterate_to_convergence(agents);

        FilteredStep {
            log_evidence: scratch.events.iter().map(|event| event.log_evidence).sum(),
            posteriors: scratch
                .skills
                .iter()
                .map(|(agent, skill)| (agent, skill.posterior()))
                .collect(),
        }
    }

    pub(crate) fn log_evidence<D: Drift<T>>(
        &self,
        targets: &[Index],
        forward: bool,
        agents: &CompetitorStore<T, D>,
    ) -> f64 {
        // Hashed once rather than scanned per player per event, so a
        // `log_evidence_for` with many keys is not quadratic.
        let target_set: std::collections::HashSet<Index> = targets.iter().copied().collect();
        // log_evidence is infrequent; a local arena avoids needing &mut self.
        let mut arena = ScratchArena::new();

        let run_event = |event: &Event, arena: &mut ScratchArena| -> f64 {
            let teams = event.within_priors(forward, &self.skills, agents);
            let result = event.outputs();
            match event.kind {
                EventKind::Ranked => {
                    Game::ranked_with_arena(
                        teams,
                        &result,
                        &event.weights,
                        self.p_draw,
                        self.convergence,
                        arena,
                    )
                    .log_evidence
                }
                EventKind::Scored { score_sigma } => {
                    Game::scored_with_arena(
                        teams,
                        &result,
                        &event.weights,
                        score_sigma,
                        self.convergence,
                        arena,
                    )
                    .log_evidence
                }
            }
        };

        if targets.is_empty() {
            if forward {
                self.events
                    .iter()
                    .map(|event| run_event(event, &mut arena))
                    .sum()
            } else {
                self.events.iter().map(|event| event.log_evidence).sum()
            }
        } else if forward {
            self.events
                .iter()
                .filter(|event| {
                    event
                        .teams
                        .iter()
                        .flat_map(|team| &team.items)
                        .any(|item| target_set.contains(&item.agent))
                })
                .map(|event| run_event(event, &mut arena))
                .sum()
        } else {
            self.events
                .iter()
                .filter(|event| {
                    event
                        .teams
                        .iter()
                        .flat_map(|team| &team.items)
                        .any(|item| target_set.contains(&item.agent))
                })
                .map(|event| event.log_evidence)
                .sum()
        }
    }

    pub fn get_composition(&self) -> Vec<Vec<Vec<Index>>> {
        self.events
            .iter()
            .map(|event| {
                event
                    .teams
                    .iter()
                    .map(|team| team.items.iter().map(|item| item.agent).collect::<Vec<_>>())
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }

    pub fn get_results(&self) -> Vec<Vec<f64>> {
        self.events
            .iter()
            .map(|event| {
                event
                    .teams
                    .iter()
                    .map(|team| team.output)
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>()
    }
}

/// Elapsed time from a competitor's previous appearance to `current`.
///
/// A negative elapsed means slices are being visited out of time order, which
/// would make drift *reduce* uncertainty. Release builds clamp to zero so a
/// bad timestamp degrades to "no drift" rather than corrupting the posterior;
/// debug builds trip instead, because reaching here is a bug in slice ordering
/// rather than something callers can cause with ordinary data.
pub(crate) fn compute_elapsed<T: Time>(last: Option<&T>, current: &T) -> i64 {
    let Some(last) = last else {
        return 0;
    };

    let elapsed = last.elapsed_to(current);

    debug_assert!(
        elapsed >= 0,
        "negative elapsed ({elapsed}) — slices visited out of time order"
    );

    elapsed.max(0)
}

#[cfg(test)]
mod tests {
    use approx::assert_ulps_eq;

    use super::*;
    use crate::{
        KeyTable, competitor::Competitor, drift::ConstantDrift, rating::Rating,
        storage::CompetitorStore,
    };

    #[test]
    fn test_one_event_each() {
        let mut index_map = KeyTable::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");
        let d = index_map.get_or_create("d");
        let e = index_map.get_or_create("e");
        let f = index_map.get_or_create("f");

        let mut agents: CompetitorStore<i64, ConstantDrift> = CompetitorStore::new();

        for agent in [a, b, c, d, e, f] {
            agents.insert(
                agent,
                Competitor {
                    rating: Rating::new(
                        Gaussian::from_ms(25.0, 25.0 / 3.0),
                        25.0 / 6.0,
                        ConstantDrift(25.0 / 300.0),
                    ),
                    ..Default::default()
                },
            );
        }

        let mut time_slice = TimeSlice::new(0i64, 0.0, crate::ConvergenceOptions::default());

        time_slice.add_events(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![c], vec![d]],
                vec![vec![e], vec![f]],
            ],
            Some(vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]]),
            None,
            vec![EventKind::Ranked; 3],
            &agents,
        );

        let post = time_slice.posteriors();

        assert_ulps_eq!(
            post[&a],
            Gaussian::from_ms(29.205220, 7.194481),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&b],
            Gaussian::from_ms(20.794779, 7.194481),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&c],
            Gaussian::from_ms(20.794779, 7.194481),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&d],
            Gaussian::from_ms(29.205220, 7.194481),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&e],
            Gaussian::from_ms(29.205220, 7.194481),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&f],
            Gaussian::from_ms(20.794779, 7.194481),
            epsilon = 1e-6
        );

        assert_eq!(time_slice.iterate_to_convergence(&agents), 1);
    }

    #[test]
    fn test_same_strength() {
        let mut index_map = KeyTable::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");
        let d = index_map.get_or_create("d");
        let e = index_map.get_or_create("e");
        let f = index_map.get_or_create("f");

        let mut agents: CompetitorStore<i64, ConstantDrift> = CompetitorStore::new();

        for agent in [a, b, c, d, e, f] {
            agents.insert(
                agent,
                Competitor {
                    rating: Rating::new(
                        Gaussian::from_ms(25.0, 25.0 / 3.0),
                        25.0 / 6.0,
                        ConstantDrift(25.0 / 300.0),
                    ),
                    ..Default::default()
                },
            );
        }

        let mut time_slice = TimeSlice::new(0i64, 0.0, crate::ConvergenceOptions::default());

        time_slice.add_events(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![a], vec![c]],
                vec![vec![b], vec![c]],
            ],
            Some(vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]]),
            None,
            vec![EventKind::Ranked; 3],
            &agents,
        );

        let post = time_slice.posteriors();

        assert_ulps_eq!(
            post[&a],
            Gaussian::from_ms(24.960978, 6.298544),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&b],
            Gaussian::from_ms(27.095590, 6.010330),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&c],
            Gaussian::from_ms(24.889681, 5.866311),
            epsilon = 1e-6
        );

        assert!(time_slice.iterate_to_convergence(&agents) > 1);

        let post = time_slice.posteriors();

        assert_ulps_eq!(
            post[&a],
            Gaussian::from_ms(25.000000, 5.419212),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&b],
            Gaussian::from_ms(25.000000, 5.419212),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&c],
            Gaussian::from_ms(25.000000, 5.419212),
            epsilon = 1e-6
        );
    }

    #[test]
    fn test_add_events() {
        let mut index_map = KeyTable::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");
        let d = index_map.get_or_create("d");
        let e = index_map.get_or_create("e");
        let f = index_map.get_or_create("f");

        let mut agents: CompetitorStore<i64, ConstantDrift> = CompetitorStore::new();

        for agent in [a, b, c, d, e, f] {
            agents.insert(
                agent,
                Competitor {
                    rating: Rating::new(
                        Gaussian::from_ms(25.0, 25.0 / 3.0),
                        25.0 / 6.0,
                        ConstantDrift(25.0 / 300.0),
                    ),
                    ..Default::default()
                },
            );
        }

        let mut time_slice = TimeSlice::new(0i64, 0.0, crate::ConvergenceOptions::default());

        time_slice.add_events(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![a], vec![c]],
                vec![vec![b], vec![c]],
            ],
            Some(vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]]),
            None,
            vec![EventKind::Ranked; 3],
            &agents,
        );

        time_slice.iterate_to_convergence(&agents);

        let post = time_slice.posteriors();

        assert_ulps_eq!(
            post[&a],
            Gaussian::from_ms(25.000000, 5.419212),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&b],
            Gaussian::from_ms(25.000000, 5.419212),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&c],
            Gaussian::from_ms(25.000000, 5.419212),
            epsilon = 1e-6
        );

        time_slice.add_events(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![a], vec![c]],
                vec![vec![b], vec![c]],
            ],
            Some(vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 0.0]]),
            None,
            vec![EventKind::Ranked; 3],
            &agents,
        );

        assert_eq!(time_slice.events.len(), 6);

        time_slice.iterate_to_convergence(&agents);

        let post = time_slice.posteriors();

        // These are convergence residuals, not exact values: by symmetry the
        // true mean is 25.0 and the iteration approaches it from above. The
        // previous expectation of 25.000003 was the residual after the
        // hard-coded 20-iteration cap; honouring `ConvergenceOptions` runs to
        // 30 and lands nearer the truth.
        assert_ulps_eq!(
            post[&a],
            Gaussian::from_ms(25.000001, 3.880150),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&b],
            Gaussian::from_ms(25.000001, 3.880150),
            epsilon = 1e-6
        );
        assert_ulps_eq!(
            post[&c],
            Gaussian::from_ms(25.000001, 3.880150),
            epsilon = 1e-6
        );
    }

    #[test]
    fn time_slice_color_groups_reorders_events() {
        // ev0: [a, b]; ev1: [c, d]; ev2: [a, c]
        // Greedy coloring: ev0→c0, ev1→c0 (disjoint), ev2→c1 (overlaps both).
        // After recompute_color_groups, physical order is [ev0, ev1, ev2]
        // and groups == [[0, 1], [2]].
        let mut index_map = KeyTable::new();

        let a = index_map.get_or_create("a");
        let b = index_map.get_or_create("b");
        let c = index_map.get_or_create("c");
        let d = index_map.get_or_create("d");

        let mut agents: CompetitorStore<i64, ConstantDrift> = CompetitorStore::new();

        for agent in [a, b, c, d] {
            agents.insert(
                agent,
                Competitor {
                    rating: Rating::new(
                        Gaussian::from_ms(25.0, 25.0 / 3.0),
                        25.0 / 6.0,
                        ConstantDrift(25.0 / 300.0),
                    ),
                    ..Default::default()
                },
            );
        }

        let mut ts = TimeSlice::new(0i64, 0.0, crate::ConvergenceOptions::default());

        ts.add_events(
            vec![
                vec![vec![a], vec![b]],
                vec![vec![c], vec![d]],
                vec![vec![a], vec![c]],
            ],
            Some(vec![vec![1.0, 0.0], vec![1.0, 0.0], vec![1.0, 0.0]]),
            None,
            vec![EventKind::Ranked; 3],
            &agents,
        );

        assert_eq!(ts.color_groups.n_colors(), 2);
        assert_eq!(ts.color_groups.groups[0], vec![0, 1]);
        assert_eq!(ts.color_groups.groups[1], vec![2]);

        assert_eq!(ts.color_groups.color_range(0), 0..2);
        assert_eq!(ts.color_groups.color_range(1), 2..3);

        // Events at positions 0 and 1 (color 0) must be disjoint — verify by
        // checking that the agent sets of self.events[0] and self.events[1] do
        // not include the agent at self.events[2].
        let agents_in_ev2: Vec<Index> = ts.events[2].iter_agents().collect();
        let agents_in_ev0: Vec<Index> = ts.events[0].iter_agents().collect();
        let agents_in_ev1: Vec<Index> = ts.events[1].iter_agents().collect();
        // ev0 and ev1 must be disjoint from each other (color-0 invariant).
        assert!(agents_in_ev0.iter().all(|ag| !agents_in_ev1.contains(ag)));
        // ev2 must share an agent with ev0 or ev1 (it needed its own color).
        let ev2_overlaps_ev0 = agents_in_ev2.iter().any(|ag| agents_in_ev0.contains(ag));
        let ev2_overlaps_ev1 = agents_in_ev2.iter().any(|ag| agents_in_ev1.contains(ag));
        assert!(ev2_overlaps_ev0 || ev2_overlaps_ev1);
    }
}