"""Link PyTorch operators to NCCL/RCCL collectives.
This module provides functionality to correlate high-level PyTorch operations
(like backward pass, DDP reducer, optimizer steps) with the low-level NCCL/RCCL
collective calls they trigger.
Linking Algorithm:
1. Use correlation IDs when available (from PyTorch profiler)
2. Fall back to timestamp proximity matching
3. Use call stack information to identify triggering operations
"""
from dataclasses import dataclass
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import Collective
from wafer_core.lib.distributed_traces.models.rank_timeline import ComputeEvent, RankTimeline
@dataclass
class PyTorchLink:
    """Link between a PyTorch operation and an NCCL collective.
    Attributes:
        collective: The NCCL/RCCL collective operation
        pytorch_op: Name of the PyTorch operator that triggered it
        pytorch_event: The compute event for the PyTorch operation (if available)
        confidence: Confidence score for this link (0.0-1.0)
        link_method: How the link was established
    """
    collective: Collective
    pytorch_op: str
    pytorch_event: Optional[ComputeEvent] = None
    confidence: float = 1.0
    link_method: str = "unknown"
class PyTorchNCCLLinker:
    """Links PyTorch operations to NCCL/RCCL collectives.
    This enables the UI to show "AllReduce triggered by DDP gradient sync
    for layer conv2.weight" rather than just showing a bare collective.
    """
    def __init__(self, timeline: RankTimeline) -> None:
        """Initialize the linker with a timeline.
        Args:
            timeline: RankTimeline containing both collectives and compute events
        """
        self.timeline = timeline
        self._links: list[PyTorchLink] = []
        self._correlation_map: dict[int, ComputeEvent] = {}
    def link_all(self) -> list[PyTorchLink]:
        """Link all collectives to their triggering PyTorch operations.
        Returns:
            List of PyTorchLink objects
        """
        self._links = []

        self._build_correlation_map()
        for collective in self.timeline.collectives:
            link = self._link_collective(collective)
            if link is not None:
                self._links.append(link)
        return self._links
    def _build_correlation_map(self) -> None:
        """Build map from correlation ID to compute event."""
        self._correlation_map = {}
        for event in self.timeline.compute_events:
            if event.correlation_id is not None:
                self._correlation_map[event.correlation_id] = event
    def _link_collective(self, collective: Collective) -> Optional[PyTorchLink]:
        """Link a single collective to its triggering operation.
        Tries multiple strategies in order of confidence:
        1. Correlation ID match
        2. Stack trace analysis
        3. Name inference from collective
        4. Timestamp proximity
        Args:
            collective: The collective to link
        Returns:
            PyTorchLink if a link is found, None otherwise
        """
        # Strategy 1: Use existing PyTorch op name if available
        if collective.pytorch_op_name:
            return PyTorchLink(
                collective=collective,
                pytorch_op=collective.pytorch_op_name,
                confidence=1.0,
                link_method="explicit",
            )
        # Strategy 2: Try correlation ID
        if collective.correlation_id is not None:
            event = self._correlation_map.get(collective.correlation_id)
            if event:
                return PyTorchLink(
                    collective=collective,
                    pytorch_op=event.pytorch_op_name or event.name,
                    pytorch_event=event,
                    confidence=0.95,
                    link_method="correlation_id",
                )
        # Strategy 3: Parse stack trace
        if collective.pytorch_stack_trace:
            op_name = self._parse_stack_trace(collective.pytorch_stack_trace)
            if op_name:
                return PyTorchLink(
                    collective=collective,
                    pytorch_op=op_name,
                    confidence=0.85,
                    link_method="stack_trace",
                )
        # Strategy 4: Infer from collective type and context
        inferred = self._infer_from_context(collective)
        if inferred:
            return PyTorchLink(
                collective=collective,
                pytorch_op=inferred,
                confidence=0.7,
                link_method="inferred",
            )
        # Strategy 5: Find nearest preceding compute event
        nearest = self._find_nearest_event(collective)
        if nearest:
            return PyTorchLink(
                collective=collective,
                pytorch_op=nearest.pytorch_op_name or nearest.name,
                pytorch_event=nearest,
                confidence=0.5,
                link_method="temporal_proximity",
            )
        return None
    def _parse_stack_trace(self, stack_trace: str) -> Optional[str]:
        """Extract meaningful operation name from stack trace.
        Looks for PyTorch framework functions that indicate what triggered
        the collective.
        Args:
            stack_trace: Stack trace string
        Returns:
            Operation name if found, None otherwise
        """
        # Common patterns in PyTorch stacks
        patterns = [
            ("DistributedDataParallel", "DDP reducer"),
            ("_ddp_reducer", "DDP gradient sync"),
            ("all_reduce", "AllReduce"),
            ("all_gather", "AllGather"),
            ("reduce_scatter", "ReduceScatter"),
            ("broadcast", "Broadcast"),
            ("backward", "Backward pass"),
            ("optimizer", "Optimizer step"),
            ("checkpoint", "Checkpoint"),
            ("autograd", "Autograd"),
        ]
        stack_lower = stack_trace.lower()
        for pattern, name in patterns:
            if pattern.lower() in stack_lower:
                return name
        # Try to extract function/module name
        import re
        # Look for patterns like "torch.distributed.all_reduce"
        match = re.search(r"torch\.([a-zA-Z_\.]+)", stack_trace)
        if match:
            return f"torch.{match.group(1)}"
        # Look for user code patterns
        match = re.search(r"File \"[^\"]+/([^/]+\.py)\", line \d+, in (\w+)", stack_trace)
        if match:
            return f"{match.group(1)}:{match.group(2)}"
        return None
    def _infer_from_context(self, collective: Collective) -> Optional[str]:
        """Infer triggering operation from collective context.
        Uses the collective type, message size, and extra metadata to
        infer what likely triggered it.
        Args:
            collective: The collective to analyze
        Returns:
            Inferred operation name, or None
        """
        from wafer_core.lib.distributed_traces.models.collective import CollectiveType
        extra = collective.extra or {}

        if "operation" in extra:
            return str(extra["operation"])
        if "op_name" in extra:
            return str(extra["op_name"])
        # Infer based on collective type and typical usage
        if collective.collective_type == CollectiveType.ALL_REDUCE:
            # AllReduce is typically used for gradient synchronization
            return "DDP gradient sync"
        if collective.collective_type == CollectiveType.ALL_GATHER:
            # AllGather typically used in FSDP
            return "FSDP parameter gather"
        if collective.collective_type == CollectiveType.REDUCE_SCATTER:
            # ReduceScatter typically used in FSDP/ZeRO
            return "FSDP gradient reduce-scatter"
        if collective.collective_type == CollectiveType.BROADCAST:
            # Broadcast often used for weight initialization or checkpoint
            return "Parameter broadcast"
        if collective.collective_type == CollectiveType.BARRIER:
            return "Synchronization barrier"
        return None
    def _find_nearest_event(self, collective: Collective) -> Optional[ComputeEvent]:
        """Find the compute event nearest in time to the collective.
        Looks for events that ended just before the collective started,
        as they are likely the triggering operation.
        Args:
            collective: The collective to find a trigger for
        Returns:
            Nearest preceding compute event, or None
        """
        # Look for events that ended within 10ms before the collective
        max_gap_ns = 10_000_000  # 10ms
        candidates = []
        for event in self.timeline.compute_events:
            # Event must end before collective starts
            if event.end_time_ns > collective.start_time_ns:
                continue
            gap = collective.start_time_ns - event.end_time_ns
            if gap <= max_gap_ns:
                candidates.append((event, gap))
        if not candidates:
            return None

        candidates.sort(key=lambda x: x[1])
        return candidates[0][0]
def link_pytorch_to_nccl(timeline: RankTimeline) -> list[PyTorchLink]:
    """Convenience function to link PyTorch operations to NCCL collectives.
    Args:
        timeline: RankTimeline to analyze
    Returns:
        List of PyTorchLink objects
    """
    linker = PyTorchNCCLLinker(timeline)
    return linker.link_all()
