//! 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 {
    pub(crate) fn new() -> Self {
        Self::default()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.groups.is_empty()
    }

    /// Number of distinct colors in the partition. Test-only.
    #[cfg(test)]
    pub(crate) fn n_colors(&self) -> usize {
        self.groups.len()
    }

    /// Total event count across all colors. Test-only.
    #[cfg(test)]
    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`.
    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;

        debug_assert_eq!(
            end - start,
            group.len(),
            "color {color_idx} is not contiguous; its range would overlap other colors"
        );

        start..end
    }

    /// Whether every color occupies a contiguous, ascending range of event
    /// indices, and no two colors overlap.
    ///
    /// The parallel sweep derives one `&mut` sub-slice per color from these
    /// ranges and relies on them being disjoint. That disjointness is what
    /// makes concurrent writes to distinct skills sound, so it is checked
    /// rather than assumed.
    pub(crate) fn groups_are_contiguous(&self) -> bool {
        let mut expected_start = 0;

        for group in &self.groups {
            if group.is_empty() {
                continue;
            }

            let ascending_run = group
                .iter()
                .enumerate()
                .all(|(offset, &idx)| idx == group[0] + offset);

            if !ascending_run || group[0] != expected_start {
                return false;
            }

            expected_start += group.len();
        }

        true
    }
}

/// 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.
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);
    }
}

#[cfg(test)]
mod properties {
    use std::collections::HashSet;

    use proptest::prelude::*;

    use super::*;

    /// The property the whole parallel sweep rests on: two events sharing a
    /// competitor must never land in the same color, because a color group is
    /// run concurrently and two events touching one competitor would race.
    ///
    /// Hand-written cases cover the shapes someone thought of. This covers the
    /// ones nobody did — the correctness of `sweep_color_groups` depends on it
    /// holding for every input, not for five.
    fn check(events: &[Vec<usize>]) {
        let groups = color_greedy(events.len(), |ev| {
            events[ev]
                .iter()
                .copied()
                .map(Index::from)
                .collect::<Vec<_>>()
        });

        // Disjointness *between events* within a color. Deduplicated per
        // event, because one event legitimately naming a competitor twice is
        // not a collision — `color_greedy` collects each event's members into
        // a set for exactly that reason.
        for color in 0..groups.n_colors() {
            let mut seen: HashSet<usize> = HashSet::new();
            for &ev in &groups.groups[color] {
                let members: HashSet<usize> = events[ev].iter().copied().collect();
                for competitor in members {
                    assert!(
                        seen.insert(competitor),
                        "competitor {competitor} shared by two events in color {color}"
                    );
                }
            }
        }

        // Every event is assigned exactly once. Without this, a partition that
        // dropped events would satisfy disjointness trivially.
        let mut assigned: Vec<usize> = groups.groups.iter().flatten().copied().collect();
        assigned.sort_unstable();
        assert_eq!(assigned, (0..events.len()).collect::<Vec<_>>());
        assert_eq!(groups.total_events(), events.len());

        // No empty colors: one would waste a sweep and make `n_colors`
        // misleading.
        for (color, group) in groups.groups.iter().enumerate() {
            assert!(!group.is_empty(), "color {color} is empty");
        }

        // Contiguity is not a property of `color_greedy` — it holds only after
        // `recompute_color_groups` reorders the events so each color occupies
        // one range. What must always hold is that the reorder is *possible*:
        // relabelling events in group order yields contiguous groups. The
        // parallel sweep slices `&mut` sub-ranges from those, so if this ever
        // failed the reorder would produce overlapping ranges.
        let mut next = 0usize;
        let relabelled: Vec<Vec<usize>> = groups
            .groups
            .iter()
            .map(|group| {
                group
                    .iter()
                    .map(|_| {
                        let i = next;
                        next += 1;
                        i
                    })
                    .collect()
            })
            .collect();
        assert!(ColorGroups { groups: relabelled }.groups_are_contiguous());
    }

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(512))]

        /// Small competitor pool, so collisions are common and colors are
        /// forced to multiply.
        #[test]
        fn colors_are_disjoint_on_a_dense_pool(
            events in prop::collection::vec(
                prop::collection::vec(0usize..6, 1..4),
                0..20,
            )
        ) {
            check(&events);
        }

        /// Wide pool, so most events are independent and land in one color.
        #[test]
        fn colors_are_disjoint_on_a_sparse_pool(
            events in prop::collection::vec(
                prop::collection::vec(0usize..200, 1..6),
                0..30,
            )
        ) {
            check(&events);
        }

        /// Repeated competitors within one event must not confuse the
        /// member-set bookkeeping.
        #[test]
        fn colors_are_disjoint_with_repeated_members(
            events in prop::collection::vec(
                prop::collection::vec(0usize..3, 1..8),
                0..15,
            )
        ) {
            check(&events);
        }
    }
}