"""Tests for correlation modules."""

import pytest

from wafer_core.lib.distributed_traces.correlation import (
    align_ranks,
    match_collectives_across_ranks,
)
from wafer_core.lib.distributed_traces.models import (
    Collective,
    CollectiveType,
    ProcessGroup,
    RankTimeline,
    TraceSession,
)


class TestRankAligner:
    """Tests for cross-rank clock alignment."""

    def test_single_rank_no_alignment(self) -> None:
        """Single rank session needs no alignment."""
        session = TraceSession(name="test")
        timeline = RankTimeline(rank=0)
        timeline.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=0,
            start_time_ns=1000000,
            end_time_ns=2000000,
        ))
        session.add_timeline(timeline)
        
        result = align_ranks(session)
        
        assert result.success
        assert result.offsets[0] == 0

    def test_aligns_two_ranks(self) -> None:
        """Should align two ranks using collective sync points."""
        session = TraceSession(name="test")
        
        # Rank 0 - reference
        t0 = RankTimeline(rank=0)
        t0.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=0,
            start_time_ns=1000000,
            end_time_ns=2000000,
            message_size_bytes=1000,
            sequence_number=1,
        ))
        session.add_timeline(t0)
        
        # Rank 1 - offset by 500000 ns
        t1 = RankTimeline(rank=1)
        t1.add_collective(Collective(
            collective_type=CollectiveType.ALL_REDUCE,
            rank=1,
            start_time_ns=1500000,  # 500000 later
            end_time_ns=2500000,
            message_size_bytes=1000,
            sequence_number=1,
        ))
        session.add_timeline(t1)
        
        result = align_ranks(session, reference_rank=0)
        
        assert result.success
        assert result.offsets[0] == 0
        # Offset should be approximately -500000 to align rank 1 to rank 0
        assert abs(result.offsets[1] - (-500000)) < 10000  # Allow some variance

    def test_alignment_with_multiple_collectives(self, sample_session: TraceSession) -> None:
        """Should use multiple collectives for better alignment."""
        result = align_ranks(sample_session)
        
        assert result.success
        assert result.alignment_points > 0


class TestCollectiveMatcher:
    """Tests for collective matching across ranks."""

    def test_matches_by_sequence_number(self) -> None:
        """Should match collectives with same sequence number."""
        session = TraceSession(name="test")
        pg = ProcessGroup(comm_id="0", world_size=2)
        
        for rank in range(2):
            t = RankTimeline(rank=rank)
            t.add_collective(Collective(
                collective_type=CollectiveType.ALL_REDUCE,
                rank=rank,
                start_time_ns=1000000,
                end_time_ns=2000000,
                sequence_number=1,
                process_group=pg,
            ))
            session.add_timeline(t)
        
        matches = match_collectives_across_ranks(session)
        
        assert len(matches) == 1
        assert matches[0].num_ranks == 2
        assert matches[0].collective_type == CollectiveType.ALL_REDUCE

    def test_matches_by_ordering_without_sequence(self) -> None:
        """Should match by ordering when sequence numbers unavailable."""
        session = TraceSession(name="test")
        
        for rank in range(2):
            t = RankTimeline(rank=rank)
            # Two collectives without sequence numbers
            t.add_collective(Collective(
                collective_type=CollectiveType.ALL_REDUCE,
                rank=rank,
                start_time_ns=1000000,
                end_time_ns=2000000,
            ))
            t.add_collective(Collective(
                collective_type=CollectiveType.ALL_REDUCE,
                rank=rank,
                start_time_ns=3000000,
                end_time_ns=4000000,
            ))
            session.add_timeline(t)
        
        matches = match_collectives_across_ranks(session)
        
        # Should match both AllReduce pairs
        assert len(matches) == 2

    def test_identifies_straggler(self, sample_session: TraceSession) -> None:
        """Should identify the straggler rank in matches."""
        matches = match_collectives_across_ranks(sample_session)
        
        assert len(matches) > 0
        # Rank 3 has longest duration in our fixture
        for match in matches:
            assert match.straggler_rank == 3

    def test_computes_straggler_impact(self, sample_session: TraceSession) -> None:
        """Should compute straggler impact correctly."""
        matches = match_collectives_across_ranks(sample_session)
        
        for match in matches:
            # Impact = (max - median) * (num_ranks - 1)
            assert match.straggler_impact_ns >= 0
