"""Cross-rank clock alignment using collective sync points.
Traces from different ranks have independent clocks. This module provides
algorithms for aligning timestamps across ranks using collective operations
as synchronization points.
The key insight is that collective operations must be called by all
participating ranks at approximately the same logical time. By identifying
matching collectives across ranks and computing the median offset between
their start times, we can estimate and correct for clock skew.
"""
from dataclasses import dataclass
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import Collective, CollectiveType
from wafer_core.lib.distributed_traces.models.rank_timeline import RankTimeline
from wafer_core.lib.distributed_traces.models.trace_session import TraceSession
@dataclass
class AlignmentResult:
    """Result of rank alignment.
    Attributes:
        offsets: Per-rank clock offsets in nanoseconds
        reference_rank: Rank used as reference (offset = 0)
        alignment_points: Number of collective operations used for alignment
        estimated_error_ns: Estimated alignment error in nanoseconds
        success: Whether alignment succeeded
        error: Error message if failed
    """
    offsets: dict[int, int]  # rank -> offset in nanoseconds
    reference_rank: int = 0
    alignment_points: int = 0
    estimated_error_ns: int = 0
    success: bool = True
    error: Optional[str] = None
class RankAligner:
    """Aligns timestamps across ranks using collective synchronization points.
    Algorithm:
    1. Identify collectives that appear across multiple ranks
    2. For each pair of ranks, compute offset from matching collectives
    3. Use median offset to reduce impact of outliers
    4. Apply offsets to align all ranks to a reference rank
    """
    def __init__(self, session: TraceSession) -> None:
        """Initialize the aligner with a trace session.
        Args:
            session: TraceSession containing timelines for multiple ranks
        """
        self.session = session
        self._offsets: dict[int, int] = {}
    def align(
        self,
        reference_rank: Optional[int] = None,
        min_alignment_points: int = 3,
    ) -> AlignmentResult:
        """Align all ranks to a reference rank.
        Args:
            reference_rank: Rank to use as reference. If None, uses rank 0.
            min_alignment_points: Minimum collective matches required for alignment.
        Returns:
            AlignmentResult with per-rank offsets and alignment quality metrics.
        """
        if self.session.num_ranks < 2:
            return AlignmentResult(
                offsets={r: 0 for r in self.session.ranks},
                reference_rank=self.session.ranks[0] if self.session.ranks else 0,
                success=True,
            )
        if reference_rank is None:
            reference_rank = min(self.session.ranks)
        if reference_rank not in self.session.timelines:
            return AlignmentResult(
                offsets={},
                reference_rank=reference_rank,
                success=False,
                error=f"Reference rank {reference_rank} not found in session",
            )
        offsets: dict[int, int] = {reference_rank: 0}
        total_alignment_points = 0
        errors: list[int] = []
        ref_timeline = self.session.timelines[reference_rank]
        for rank, timeline in self.session.timelines.items():
            if rank == reference_rank:
                continue
            offset, points, error = self._compute_offset(ref_timeline, timeline)
            if points < min_alignment_points:
                # Fall back to first collective alignment
                offset, points, error = self._fallback_alignment(ref_timeline, timeline)
            offsets[rank] = offset
            total_alignment_points += points
            if error is not None:
                errors.append(error)
        # Estimate alignment error as median of per-pair errors
        estimated_error = int(sum(errors) / len(errors)) if errors else 0
        return AlignmentResult(
            offsets=offsets,
            reference_rank=reference_rank,
            alignment_points=total_alignment_points,
            estimated_error_ns=estimated_error,
            success=True,
        )
    def _compute_offset(
        self,
        ref_timeline: RankTimeline,
        other_timeline: RankTimeline,
    ) -> tuple[int, int, Optional[int]]:
        """Compute clock offset between two timelines.
        Uses matching collectives to estimate the clock offset.
        Args:
            ref_timeline: Reference timeline
            other_timeline: Timeline to align
        Returns:
            Tuple of (offset_ns, num_alignment_points, error_estimate)
        """

        matches = self._find_matching_collectives(ref_timeline, other_timeline)
        if not matches:
            return 0, 0, None
        offsets = []
        for ref_coll, other_coll in matches:
            # Offset = ref_start - other_start
            # To align other to ref, add this offset to other's timestamps
            offset = ref_coll.start_time_ns - other_coll.start_time_ns
            offsets.append(offset)
        offsets.sort()
        median_offset = offsets[len(offsets) // 2]
        # Estimate error as deviation from median
        errors = [abs(o - median_offset) for o in offsets]
        median_error = sorted(errors)[len(errors) // 2]
        return median_offset, len(matches), median_error
    def _find_matching_collectives(
        self,
        timeline1: RankTimeline,
        timeline2: RankTimeline,
    ) -> list[tuple[Collective, Collective]]:
        """Find collectives that match between two timelines.
        Matching criteria:
        1. Same collective type
        2. Same message size (within tolerance)
        3. Same sequence in their respective timelines
        Args:
            timeline1: First timeline
            timeline2: Second timeline
        Returns:
            List of (coll1, coll2) pairs that match
        """
        matches = []
        colls1_by_type: dict[CollectiveType, list[Collective]] = {}
        colls2_by_type: dict[CollectiveType, list[Collective]] = {}
        for c in timeline1.collectives:
            colls1_by_type.setdefault(c.collective_type, []).append(c)
        for c in timeline2.collectives:
            colls2_by_type.setdefault(c.collective_type, []).append(c)
        for ctype, colls1 in colls1_by_type.items():
            colls2 = colls2_by_type.get(ctype, [])
            colls1 = sorted(colls1, key=lambda c: c.start_time_ns)
            colls2 = sorted(colls2, key=lambda c: c.start_time_ns)
            for i, c1 in enumerate(colls1):
                if i >= len(colls2):
                    break
                c2 = colls2[i]
                if c1.message_size_bytes > 0 and c2.message_size_bytes > 0:
                    ratio = c1.message_size_bytes / c2.message_size_bytes
                    if not (0.9 <= ratio <= 1.1):
                        continue
                matches.append((c1, c2))
        return matches
    def _fallback_alignment(
        self,
        ref_timeline: RankTimeline,
        other_timeline: RankTimeline,
    ) -> tuple[int, int, Optional[int]]:
        """Fallback alignment using first/last events.
        Used when not enough matching collectives are found.
        Args:
            ref_timeline: Reference timeline
            other_timeline: Timeline to align
        Returns:
            Tuple of (offset_ns, num_alignment_points, error_estimate)
        """
        # Try to align based on first event
        ref_start = ref_timeline.start_time_ns
        other_start = other_timeline.start_time_ns
        if ref_start > 0 and other_start > 0:
            offset = ref_start - other_start
            return offset, 1, None
        return 0, 0, None
    def apply_alignment(self, result: AlignmentResult) -> None:
        """Apply alignment offsets to the session's timelines.
        Modifies the timelines in place.
        Args:
            result: AlignmentResult from align()
        """
        for rank, offset in result.offsets.items():
            if rank in self.session.timelines:
                timeline = self.session.timelines[rank]
                timeline.clock_offset_ns = offset
                timeline.apply_clock_offset()
        self.session.aligned = True
def align_ranks(
    session: TraceSession,
    reference_rank: Optional[int] = None,
    apply: bool = True,
) -> AlignmentResult:
    """Convenience function to align ranks in a session.
    Args:
        session: TraceSession to align
        reference_rank: Optional reference rank
        apply: If True, apply offsets to timelines
    Returns:
        AlignmentResult with alignment details
    """
    aligner = RankAligner(session)
    result = aligner.align(reference_rank)
    if apply and result.success:
        aligner.apply_alignment(result)
    return result
