"""Communication/compute overlap analysis.
Analyzes the overlap between compute operations and communication
operations to measure how well training overlaps computation with
gradient synchronization.
High overlap is desirable as it hides communication latency behind
useful compute work. Low overlap indicates the training loop is
blocked waiting for communication.
"""
from dataclasses import dataclass
from typing import Optional
from wafer_core.lib.distributed_traces.models.rank_timeline import (
    RankTimeline,
    TimelineWindow,
    WindowType,
)
from wafer_core.lib.distributed_traces.models.trace_session import (
    IterationOverlap,
    OverlapMetrics,
    TraceSession,
)
@dataclass
class OverlapAnalysisResult:
    """Result of overlap analysis.
    Attributes:
        overall: Overall overlap metrics for the trace
        per_rank: Per-rank overlap metrics
        per_iteration: Per-iteration overlap if iterations detected
        windows: Time windows showing overlap regions
        recommendations: Suggestions for improving overlap
    """
    overall: OverlapMetrics
    per_rank: dict[int, OverlapMetrics]
    per_iteration: list[IterationOverlap]
    windows: list[TimelineWindow]
    recommendations: list[str]
class OverlapAnalyzer:
    """Analyzes communication/compute overlap in distributed training traces."""
    def __init__(self, session: TraceSession) -> None:
        """Initialize the analyzer.
        Args:
            session: TraceSession to analyze
        """
        self.session = session
    def analyze(self) -> OverlapAnalysisResult:
        """Perform full overlap analysis.
        Returns:
            OverlapAnalysisResult with overall and per-rank metrics
        """
        per_rank: dict[int, OverlapMetrics] = {}
        all_windows: list[TimelineWindow] = []
        all_iteration_metrics: list[IterationOverlap] = []
        for rank, timeline in self.session.timelines.items():
            metrics, windows = self._analyze_timeline(timeline)
            per_rank[rank] = metrics
            all_windows.extend(windows)
            # Try to detect iterations
            iterations = self._detect_iterations(timeline, windows)
            all_iteration_metrics.extend(iterations)
        # Aggregate to overall metrics
        overall = self._aggregate_metrics(list(per_rank.values()))
        # Generate recommendations
        recommendations = self._generate_recommendations(overall, per_rank)
        return OverlapAnalysisResult(
            overall=overall,
            per_rank=per_rank,
            per_iteration=all_iteration_metrics,
            windows=all_windows,
            recommendations=recommendations,
        )
    def _analyze_timeline(
        self,
        timeline: RankTimeline,
    ) -> tuple[OverlapMetrics, list[TimelineWindow]]:
        """Analyze overlap for a single timeline.
        Segments the timeline into windows and classifies each as:
        - compute_only: Only compute kernels active
        - communication_only: Only collective operations active
        - overlapped: Both compute and communication active
        - idle: Neither active
        Args:
            timeline: RankTimeline to analyze
        Returns:
            Tuple of (OverlapMetrics, list of TimelineWindows)
        """
        if not timeline.collectives and not timeline.compute_events:
            return self._empty_metrics(), []
        boundaries = set()
        for c in timeline.collectives:
            boundaries.add(c.start_time_ns)
            boundaries.add(c.end_time_ns)
        for e in timeline.compute_events:
            boundaries.add(e.start_time_ns)
            boundaries.add(e.end_time_ns)
        boundaries = sorted(boundaries)
        # Classify each window
        windows = []
        total_compute = 0
        total_communication = 0
        total_overlap = 0
        total_compute_only = 0
        total_comm_only = 0
        total_idle = 0
        for i in range(len(boundaries) - 1):
            start = boundaries[i]
            end = boundaries[i + 1]
            duration = end - start
            if duration <= 0:
                continue
            has_compute = self._has_active_compute(timeline, start, end)
            has_comm = self._has_active_communication(timeline, start, end)
            if has_compute and has_comm:
                window_type = WindowType.OVERLAPPED
                total_overlap += duration
                total_compute += duration
                total_communication += duration
            elif has_compute:
                window_type = WindowType.COMPUTE_ONLY
                total_compute += duration
                total_compute_only += duration
            elif has_comm:
                window_type = WindowType.COMMUNICATION_ONLY
                total_communication += duration
                total_comm_only += duration
            else:
                window_type = WindowType.IDLE
                total_idle += duration
            windows.append(TimelineWindow(
                window_type=window_type,
                start_time_ns=start,
                end_time_ns=end,
            ))
        overlap_ratio = (
            total_overlap / total_communication if total_communication > 0 else 0
        )
        metrics = OverlapMetrics(
            total_compute_time_ns=total_compute,
            total_communication_time_ns=total_communication,
            overlapped_time_ns=total_overlap,
            compute_only_time_ns=total_compute_only,
            communication_only_time_ns=total_comm_only,
            idle_time_ns=total_idle,
            overlap_ratio=overlap_ratio,
        )
        return metrics, windows
    def _has_active_compute(
        self,
        timeline: RankTimeline,
        start: int,
        end: int,
    ) -> bool:
        """Check if any compute event is active during [start, end)."""
        for e in timeline.compute_events:
            # Event is active if it overlaps with the window
            if e.start_time_ns < end and e.end_time_ns > start:
                return True
        return False
    def _has_active_communication(
        self,
        timeline: RankTimeline,
        start: int,
        end: int,
    ) -> bool:
        """Check if any collective is active during [start, end)."""
        for c in timeline.collectives:
            if c.start_time_ns < end and c.end_time_ns > start:
                return True
        return False
    def _detect_iterations(
        self,
        timeline: RankTimeline,
        windows: list[TimelineWindow],
    ) -> list[IterationOverlap]:
        """Try to detect training iterations and compute per-iteration overlap.
        Heuristic: Look for periodic patterns in the timeline, typically
        marked by barriers or regular AllReduce patterns.
        Args:
            timeline: Timeline to analyze
            windows: Pre-computed windows
        Returns:
            List of per-iteration overlap metrics
        """
        from wafer_core.lib.distributed_traces.models.collective import CollectiveType
        iterations: list[IterationOverlap] = []
        # Try to find iteration boundaries using barriers
        barriers = [
            c for c in timeline.collectives
            if c.collective_type == CollectiveType.BARRIER
        ]
        if len(barriers) >= 2:
            for i, (start_barrier, end_barrier) in enumerate(zip(barriers[:-1], barriers[1:])):
                iter_metrics = self._compute_iteration_metrics(
                    timeline,
                    start_barrier.end_time_ns,
                    end_barrier.start_time_ns,
                    i,
                )
                if iter_metrics:
                    iterations.append(iter_metrics)
        else:
            # Try to detect iterations from AllReduce patterns
            all_reduces = [
                c for c in timeline.collectives
                if c.collective_type == CollectiveType.ALL_REDUCE
            ]
            if len(all_reduces) >= 10:
                # Assume gap between groups indicates iteration boundary
                groups = self._group_by_gap(all_reduces)
                for i, group in enumerate(groups):
                    if not group:
                        continue
                    start = group[0].start_time_ns
                    end = group[-1].end_time_ns
                    iter_metrics = self._compute_iteration_metrics(timeline, start, end, i)
                    if iter_metrics:
                        iterations.append(iter_metrics)
        return iterations
    def _group_by_gap(
        self,
        collectives: list,
        gap_threshold_factor: float = 5.0,
    ) -> list[list]:
        """Group collectives by temporal gaps.
        Args:
            collectives: Sorted list of collectives
            gap_threshold_factor: Gap > median_gap * factor starts new group
        Returns:
            List of collective groups
        """
        if len(collectives) < 2:
            return [collectives]
        gaps = []
        for i in range(len(collectives) - 1):
            gap = collectives[i + 1].start_time_ns - collectives[i].end_time_ns
            gaps.append(gap)

        sorted_gaps = sorted(gaps)
        median_gap = sorted_gaps[len(sorted_gaps) // 2]
        threshold = median_gap * gap_threshold_factor
        groups: list[list] = [[collectives[0]]]
        for i, gap in enumerate(gaps):
            if gap > threshold:
                groups.append([])
            groups[-1].append(collectives[i + 1])
        return groups
    def _compute_iteration_metrics(
        self,
        timeline: RankTimeline,
        start_ns: int,
        end_ns: int,
        iteration: int,
    ) -> Optional[IterationOverlap]:
        """Compute overlap metrics for a single iteration.
        Args:
            timeline: Timeline to analyze
            start_ns: Iteration start time
            end_ns: Iteration end time
            iteration: Iteration number
        Returns:
            IterationOverlap metrics, or None if no activity
        """
        compute_time = 0
        comm_time = 0
        overlap_time = 0
        # Simple approximation: sum durations within iteration
        for c in timeline.collectives:
            if c.start_time_ns >= start_ns and c.end_time_ns <= end_ns:
                comm_time += c.duration_ns
        for e in timeline.compute_events:
            if e.start_time_ns >= start_ns and e.end_time_ns <= end_ns:
                compute_time += e.duration_ns
        if compute_time == 0 and comm_time == 0:
            return None
        # Estimate overlap (simplified - assumes events on different streams)
        total_time = end_ns - start_ns
        sum_time = compute_time + comm_time
        if sum_time > total_time:
            overlap_time = sum_time - total_time
        overlap_ratio = overlap_time / comm_time if comm_time > 0 else 0
        return IterationOverlap(
            iteration=iteration,
            start_time_ns=start_ns,
            end_time_ns=end_ns,
            compute_time_ns=compute_time,
            communication_time_ns=comm_time,
            overlapped_time_ns=overlap_time,
            overlap_ratio=overlap_ratio,
        )
    def _aggregate_metrics(
        self,
        metrics_list: list[OverlapMetrics],
    ) -> OverlapMetrics:
        """Aggregate overlap metrics across ranks."""
        if not metrics_list:
            return self._empty_metrics()
        n = len(metrics_list)
        total_compute = sum(m.total_compute_time_ns for m in metrics_list) // n
        total_comm = sum(m.total_communication_time_ns for m in metrics_list) // n
        overlapped = sum(m.overlapped_time_ns for m in metrics_list) // n
        compute_only = sum(m.compute_only_time_ns for m in metrics_list) // n
        comm_only = sum(m.communication_only_time_ns for m in metrics_list) // n
        idle = sum(m.idle_time_ns for m in metrics_list) // n
        overlap_ratio = overlapped / total_comm if total_comm > 0 else 0
        return OverlapMetrics(
            total_compute_time_ns=total_compute,
            total_communication_time_ns=total_comm,
            overlapped_time_ns=overlapped,
            compute_only_time_ns=compute_only,
            communication_only_time_ns=comm_only,
            idle_time_ns=idle,
            overlap_ratio=overlap_ratio,
        )
    def _empty_metrics(self) -> OverlapMetrics:
        """Return empty metrics."""
        return OverlapMetrics(
            total_compute_time_ns=0,
            total_communication_time_ns=0,
            overlapped_time_ns=0,
            compute_only_time_ns=0,
            communication_only_time_ns=0,
            idle_time_ns=0,
            overlap_ratio=0,
        )
    def _generate_recommendations(
        self,
        overall: OverlapMetrics,
        per_rank: dict[int, OverlapMetrics],
    ) -> list[str]:
        """Generate recommendations for improving overlap.
        Args:
            overall: Overall overlap metrics
            per_rank: Per-rank metrics
        Returns:
            List of recommendation strings
        """
        recommendations = []
        # Low overlap
        if overall.overlap_ratio < 0.3:
            recommendations.append(
                "Low overlap ratio ({:.1f}%). Consider:".format(
                    overall.overlap_percent
                )
            )
            recommendations.append(
                "  - Enable gradient bucketing in DDP (bucket_cap_mb parameter)"
            )
            recommendations.append(
                "  - Use smaller first bucket to start communication earlier"
            )
            recommendations.append(
                "  - Check if model uses torch.no_grad() blocks that delay sync"
            )
        # High communication-only time
        if overall.communication_only_time_ns > overall.compute_only_time_ns:
            recommendations.append(
                "Communication time exceeds compute time. Consider:"
            )
            recommendations.append(
                "  - Increasing batch size to get more compute per sync"
            )
            recommendations.append(
                "  - Using gradient accumulation to reduce sync frequency"
            )
        # Significant idle time
        if overall.idle_time_ns > overall.total_communication_time_ns * 0.1:
            recommendations.append(
                "Significant idle time detected ({:.1f}ms). Check for:".format(
                    overall.idle_time_ns / 1_000_000
                )
            )
            recommendations.append(
                "  - CPU bottlenecks in data loading"
            )
            recommendations.append(
                "  - Synchronization barriers in user code"
            )
        return recommendations
def analyze_overlap(session: TraceSession) -> OverlapAnalysisResult:
    """Convenience function to analyze overlap in a session.
    Args:
        session: TraceSession to analyze
    Returns:
        OverlapAnalysisResult
    """
    analyzer = OverlapAnalyzer(session)
    return analyzer.analyze()
