//! Greedy graph coloring for within-slice event independence.
//!
//! Events sharing no `Index` can be processed in parallel under async-EP
//! semantics. This module partitions a list of events into "colors" such
//! that events of the same color touch disjoint index sets.
//!
//! The algorithm is greedy: for each event in ingestion order, place it in
//! the lowest-numbered color whose existing members share no `Index`. If
//! no existing color accepts the event, open a new color.
//!
//! Complexity: O(n × c × m) where n is events, c is colors (small, ≤ 5 in
//! practice), and m is average team size.

use std::collections::HashSet;

use crate::Index;

/// Partition of event indices into color groups.
///
/// Each inner `Vec<usize>` holds the indices (into the original events
/// array) of events assigned to one color. Colors are iterated in ascending
/// order by convention.
#[derive(Clone, Debug, Default)]
pub(crate) struct ColorGroups {
    pub(crate) groups: Vec<Vec<usize>>,
}

impl ColorGroups {
    #[allow(dead_code)]
    pub(crate) fn new() -> Self {
        Self::default()
    }

    #[allow(dead_code)]
    pub(crate) fn n_colors(&self) -> usize {
        self.groups.len()
    }

    #[allow(dead_code)]
    pub(crate) fn is_empty(&self) -> bool {
        self.groups.is_empty()
    }

    /// Total event count across all colors.
    #[allow(dead_code)]
    pub(crate) fn total_events(&self) -> usize {
        self.groups.iter().map(|g| g.len()).sum()
    }

    /// Contiguous index range for one color after events have been reordered
    /// into color-contiguous positions by `TimeSlice::recompute_color_groups`.
    #[allow(dead_code)]
    pub(crate) fn color_range(&self, color_idx: usize) -> std::ops::Range<usize> {
        let group = &self.groups[color_idx];
        if group.is_empty() {
            return 0..0;
        }
        let start = *group.first().unwrap();
        let end = *group.last().unwrap() + 1;
        start..end
    }
}

/// Compute color groups greedily.
///
/// `index_set(ev_idx)` yields, for each event index, the iterator of
/// `Index` values that event touches. The returned `ColorGroups` has one
/// inner `Vec<usize>` per color, containing event indices in the order
/// they were assigned.
#[allow(dead_code)]
pub(crate) fn color_greedy<I, F>(n_events: usize, index_set: F) -> ColorGroups
where
    F: Fn(usize) -> I,
    I: IntoIterator<Item = Index>,
{
    let mut groups: Vec<Vec<usize>> = Vec::new();
    let mut members: Vec<HashSet<Index>> = Vec::new();

    for ev_idx in 0..n_events {
        let ev_members: HashSet<Index> = index_set(ev_idx).into_iter().collect();
        // Find first color whose member-set is disjoint from this event's indices.
        let chosen = members.iter().position(|m| m.is_disjoint(&ev_members));
        let color_idx = match chosen {
            Some(c) => c,
            None => {
                groups.push(Vec::new());
                members.push(HashSet::new());
                groups.len() - 1
            }
        };
        groups[color_idx].push(ev_idx);
        members[color_idx].extend(ev_members);
    }

    ColorGroups { groups }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn idx(i: usize) -> Index {
        Index::from(i)
    }

    #[test]
    fn single_event_gets_one_color() {
        let cg = color_greedy(1, |_| vec![idx(0), idx(1)]);
        assert_eq!(cg.n_colors(), 1);
        assert_eq!(cg.groups[0], vec![0]);
    }

    #[test]
    fn disjoint_events_share_a_color() {
        let cg = color_greedy(2, |i| match i {
            0 => vec![idx(0), idx(1)],
            1 => vec![idx(2), idx(3)],
            _ => unreachable!(),
        });
        assert_eq!(cg.n_colors(), 1);
        assert_eq!(cg.groups[0], vec![0, 1]);
    }

    #[test]
    fn overlapping_events_need_separate_colors() {
        let cg = color_greedy(2, |i| match i {
            0 => vec![idx(0), idx(1)],
            1 => vec![idx(1), idx(2)],
            _ => unreachable!(),
        });
        assert_eq!(cg.n_colors(), 2);
        assert_eq!(cg.groups[0], vec![0]);
        assert_eq!(cg.groups[1], vec![1]);
    }

    #[test]
    fn three_events_two_colors() {
        // Event 0: {0, 1}; event 1: {2, 3}; event 2: {0, 2}.
        // Greedy: ev0→c0, ev1→c0 (disjoint), ev2 overlaps both→c1.
        let cg = color_greedy(3, |i| match i {
            0 => vec![idx(0), idx(1)],
            1 => vec![idx(2), idx(3)],
            2 => vec![idx(0), idx(2)],
            _ => unreachable!(),
        });
        assert_eq!(cg.n_colors(), 2);
        assert_eq!(cg.groups[0], vec![0, 1]);
        assert_eq!(cg.groups[1], vec![2]);
    }

    #[test]
    fn total_events_counts_correctly() {
        let cg = color_greedy(4, |_| vec![idx(0)]);
        // All events touch index 0 → 4 distinct colors.
        assert_eq!(cg.n_colors(), 4);
        assert_eq!(cg.total_events(), 4);
    }
}
