"""Match collectives across ranks by type, size, ordering.
This module provides functionality to identify the same logical collective
operation across multiple ranks. This is essential for:
- Straggler detection (finding which rank was slowest)
- Cross-rank timing analysis
- Validating collective correctness
Matching Algorithm:
1. Group collectives by process group
2. Within each group, match by (collective_type, sequence_number, message_size)
3. Handle missing sequence numbers by falling back to ordering
"""
from dataclasses import dataclass
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import Collective, CollectiveType, ProcessGroup
from wafer_core.lib.distributed_traces.models.trace_session import CollectiveMatch, TraceSession
@dataclass
class MatchingConfig:
    """Configuration for collective matching.
    Attributes:
        size_tolerance: Tolerance for message size matching (0.0-1.0)
        require_all_ranks: If True, only match when all ranks participate
        max_time_diff_ns: Maximum time difference between matched collectives
    """
    size_tolerance: float = 0.1  # 10% tolerance
    require_all_ranks: bool = False
    max_time_diff_ns: Optional[int] = None  # No limit by default
class CollectiveMatcher:
    """Matches collective operations across ranks.
    Identifies the same logical collective operation on different ranks
    to enable cross-rank analysis and straggler detection.
    """
    def __init__(
        self,
        session: TraceSession,
        config: Optional[MatchingConfig] = None,
    ) -> None:
        """Initialize the matcher.
        Args:
            session: TraceSession with aligned timelines
            config: Optional matching configuration
        """
        self.session = session
        self.config = config or MatchingConfig()
        self._matches: list[CollectiveMatch] = []
    def match_all(self) -> list[CollectiveMatch]:
        """Match all collectives across ranks.
        Returns:
            List of CollectiveMatch objects, each representing one
            logical collective operation across multiple ranks.
        """
        self._matches = []
        process_groups = self._get_process_groups()
        # For each process group, match collectives
        for pg_id in process_groups:
            pg_matches = self._match_within_process_group(pg_id)
            self._matches.extend(pg_matches)
        return self._matches
    def _get_process_groups(self) -> set[str]:
        """Get all unique process group IDs in the session."""
        pg_ids: set[str] = set()
        for timeline in self.session.timelines.values():
            for coll in timeline.collectives:
                if coll.process_group:
                    pg_ids.add(coll.process_group.comm_id)
                else:
                    pg_ids.add("")
        return pg_ids
    def _match_within_process_group(self, pg_id: str) -> list[CollectiveMatch]:
        """Match collectives within a single process group.
        Args:
            pg_id: Process group ID to match within
        Returns:
            List of CollectiveMatch objects for this process group
        """

        colls_by_rank: dict[int, list[Collective]] = {}
        for rank, timeline in self.session.timelines.items():
            colls = []
            for c in timeline.collectives:
                coll_pg_id = c.process_group.comm_id if c.process_group else ""
                if coll_pg_id == pg_id:
                    colls.append(c)
            colls.sort(key=lambda c: c.start_time_ns)
            colls_by_rank[rank] = colls
        if not colls_by_rank:
            return []
        if self._has_sequence_numbers(colls_by_rank):
            return self._match_by_sequence(colls_by_rank, pg_id)
        else:
            return self._match_by_ordering(colls_by_rank, pg_id)
    def _has_sequence_numbers(self, colls_by_rank: dict[int, list[Collective]]) -> bool:
        """Check if collectives have sequence numbers."""
        for colls in colls_by_rank.values():
            for c in colls:
                if c.sequence_number is not None:
                    return True
        return False
    def _match_by_sequence(
        self,
        colls_by_rank: dict[int, list[Collective]],
        pg_id: str,
    ) -> list[CollectiveMatch]:
        """Match collectives by sequence number.
        Args:
            colls_by_rank: Collectives grouped by rank
            pg_id: Process group ID
        Returns:
            List of matched collectives
        """
        groups: dict[tuple[CollectiveType, int], dict[int, Collective]] = {}
        for rank, colls in colls_by_rank.items():
            for c in colls:
                if c.sequence_number is None:
                    continue
                key = (c.collective_type, c.sequence_number)
                if key not in groups:
                    groups[key] = {}
                groups[key][rank] = c
        matches = []
        for (ctype, seq), rank_colls in sorted(groups.items(), key=lambda x: x[0][1]):
            if len(rank_colls) < 2:
                continue
            # Validate message size similarity
            if not self._validate_sizes(list(rank_colls.values())):
                continue
            if not self._validate_time_diff(list(rank_colls.values())):
                continue
            if self.config.require_all_ranks:
                if len(rank_colls) != self.session.num_ranks:
                    continue
            msg_size = max(c.message_size_bytes for c in rank_colls.values())
            match = CollectiveMatch(
                collective_type=ctype,
                sequence_number=seq,
                process_group_id=pg_id,
                collectives=rank_colls,
                message_size_bytes=msg_size,
            )
            matches.append(match)
        return matches
    def _match_by_ordering(
        self,
        colls_by_rank: dict[int, list[Collective]],
        pg_id: str,
    ) -> list[CollectiveMatch]:
        """Match collectives by their ordering within each rank.
        Fallback method when sequence numbers are not available.
        Assumes collectives of the same type appear in the same order
        across all ranks.
        Args:
            colls_by_rank: Collectives grouped by rank
            pg_id: Process group ID
        Returns:
            List of matched collectives
        """
        by_type: dict[CollectiveType, dict[int, list[Collective]]] = {}
        for rank, colls in colls_by_rank.items():
            for c in colls:
                if c.collective_type not in by_type:
                    by_type[c.collective_type] = {}
                if rank not in by_type[c.collective_type]:
                    by_type[c.collective_type][rank] = []
                by_type[c.collective_type][rank].append(c)
        matches = []
        seq_counter = 0
        for ctype, rank_colls in by_type.items():

            min_count = min(len(colls) for colls in rank_colls.values())
            for i in range(min_count):
                matched = {}
                for rank, colls in rank_colls.items():
                    if i < len(colls):
                        matched[rank] = colls[i]
                if len(matched) < 2:
                    continue
                # Validate sizes and time
                if not self._validate_sizes(list(matched.values())):
                    continue
                if not self._validate_time_diff(list(matched.values())):
                    continue
                if self.config.require_all_ranks:
                    if len(matched) != self.session.num_ranks:
                        continue
                msg_size = max(c.message_size_bytes for c in matched.values())
                match = CollectiveMatch(
                    collective_type=ctype,
                    sequence_number=seq_counter,
                    process_group_id=pg_id,
                    collectives=matched,
                    message_size_bytes=msg_size,
                )
                matches.append(match)
                seq_counter += 1
        return matches
    def _validate_sizes(self, collectives: list[Collective]) -> bool:
        """Validate that message sizes are similar."""
        sizes = [c.message_size_bytes for c in collectives if c.message_size_bytes > 0]
        if len(sizes) < 2:
            return True  # Can't validate, assume OK
        min_size = min(sizes)
        max_size = max(sizes)
        if min_size == 0:
            return True
        ratio = max_size / min_size
        return ratio <= (1 + self.config.size_tolerance)
    def _validate_time_diff(self, collectives: list[Collective]) -> bool:
        """Validate that time differences are within limit."""
        if self.config.max_time_diff_ns is None:
            return True
        starts = [c.start_time_ns for c in collectives]
        time_diff = max(starts) - min(starts)
        return time_diff <= self.config.max_time_diff_ns
def match_collectives_across_ranks(
    session: TraceSession,
    config: Optional[MatchingConfig] = None,
) -> list[CollectiveMatch]:
    """Convenience function to match collectives in a session.
    Args:
        session: TraceSession with timelines
        config: Optional matching configuration
    Returns:
        List of CollectiveMatch objects
    """
    matcher = CollectiveMatcher(session, config)
    return matcher.match_all()
