"""Collective operation representation.
Defines data structures for representing collective communication operations
from NCCL and RCCL traces.
"""
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
class CollectiveType(Enum):
    """Type of collective operation."""
    ALL_REDUCE = "all_reduce"
    ALL_GATHER = "all_gather"
    REDUCE_SCATTER = "reduce_scatter"
    BROADCAST = "broadcast"
    ALL_TO_ALL = "all_to_all"
    REDUCE = "reduce"
    GATHER = "gather"
    SCATTER = "scatter"
    BARRIER = "barrier"
    SEND = "send"
    RECV = "recv"
    UNKNOWN = "unknown"
    @classmethod
    def from_string(cls, s: str) -> "CollectiveType":
        """Parse collective type from string (case-insensitive).
        Handles various naming conventions:
        - nccl:all_reduce, ncclAllReduce, all_reduce, AllReduce
        - rccl:all_gather, rcclAllGather, all_gather, AllGather
        """
        normalized = s.lower().replace("-", "_")
        # Strip nccl/rccl prefix
        for prefix in ["nccl:", "rccl:", "nccl", "rccl"]:
            if normalized.startswith(prefix):
                normalized = normalized[len(prefix) :]
                break

        # ncclAllReduce -> all_reduce
        import re
        normalized = re.sub(r"([a-z])([A-Z])", r"\1_\2", normalized).lower()
        mapping = {
            "all_reduce": cls.ALL_REDUCE,
            "allreduce": cls.ALL_REDUCE,
            "all_gather": cls.ALL_GATHER,
            "allgather": cls.ALL_GATHER,
            "reduce_scatter": cls.REDUCE_SCATTER,
            "reducescatter": cls.REDUCE_SCATTER,
            "broadcast": cls.BROADCAST,
            "bcast": cls.BROADCAST,
            "all_to_all": cls.ALL_TO_ALL,
            "alltoall": cls.ALL_TO_ALL,
            "reduce": cls.REDUCE,
            "gather": cls.GATHER,
            "scatter": cls.SCATTER,
            "barrier": cls.BARRIER,
            "send": cls.SEND,
            "recv": cls.RECV,
        }
        return mapping.get(normalized, cls.UNKNOWN)
class DataType(Enum):
    """Data type used in collective operation."""
    FLOAT32 = "float32"
    FLOAT16 = "float16"
    BFLOAT16 = "bfloat16"
    FLOAT64 = "float64"
    INT32 = "int32"
    INT64 = "int64"
    INT8 = "int8"
    UINT8 = "uint8"
    UNKNOWN = "unknown"
    def size_bytes(self) -> int:
        """Return size of this data type in bytes."""
        sizes = {
            DataType.FLOAT32: 4,
            DataType.FLOAT16: 2,
            DataType.BFLOAT16: 2,
            DataType.FLOAT64: 8,
            DataType.INT32: 4,
            DataType.INT64: 8,
            DataType.INT8: 1,
            DataType.UINT8: 1,
            DataType.UNKNOWN: 4,  # Assume float32 as default
        }
        return sizes[self]
    @classmethod
    def from_string(cls, s: str) -> "DataType":
        """Parse data type from string."""
        normalized = s.lower().replace(" ", "")
        mapping = {
            "float32": cls.FLOAT32,
            "float": cls.FLOAT32,
            "fp32": cls.FLOAT32,
            "float16": cls.FLOAT16,
            "half": cls.FLOAT16,
            "fp16": cls.FLOAT16,
            "bfloat16": cls.BFLOAT16,
            "bf16": cls.BFLOAT16,
            "float64": cls.FLOAT64,
            "double": cls.FLOAT64,
            "fp64": cls.FLOAT64,
            "int32": cls.INT32,
            "int": cls.INT32,
            "int64": cls.INT64,
            "long": cls.INT64,
            "int8": cls.INT8,
            "uint8": cls.UINT8,
            "byte": cls.UINT8,
        }
        return mapping.get(normalized, cls.UNKNOWN)
@dataclass(frozen=True)
class ProcessGroup:
    """Represents a process group / communicator.
    Attributes:
        comm_id: Communicator ID from NCCL/RCCL
        name: Optional human-readable name (e.g., "default", "dp", "tp")
        ranks: List of global ranks in this group
        world_size: Number of ranks in this group
    """
    comm_id: str
    name: Optional[str] = None
    ranks: tuple[int, ...] = field(default_factory=tuple)
    world_size: int = 0
    def __post_init__(self) -> None:
        """Validate process group."""
        # Ensure ranks is a tuple (not list) for hashability
        if isinstance(self.ranks, list):
            object.__setattr__(self, "ranks", tuple(self.ranks))
@dataclass(frozen=True)
class Collective:
    """Represents a single collective operation.
    This is the core data structure for representing one collective call
    observed during distributed training.
    Attributes:
        collective_type: Type of collective (AllReduce, AllGather, etc.)
        rank: Global rank that executed this collective
        start_time_ns: Start timestamp in nanoseconds
        end_time_ns: End timestamp in nanoseconds
        message_size_bytes: Size of message in bytes
        count: Number of elements
        datatype: Data type of elements
        process_group: Process group this collective belongs to
        sequence_number: Sequence number within the process group
        algorithm: Algorithm used (ring, tree, etc.)
        protocol: Protocol used (LL, LL128, Simple)
        num_channels: Number of NCCL channels used
        correlation_id: ID for correlating with GPU kernels
        pytorch_op_name: Name of PyTorch operator that triggered this collective
        pytorch_stack_trace: Stack trace from PyTorch profiler
        source_file: Source file that triggered the collective
        source_line: Line number in source file
        extra: Additional metadata from trace
    """
    collective_type: CollectiveType
    rank: int
    start_time_ns: int
    end_time_ns: int
    message_size_bytes: int = 0
    count: int = 0
    datatype: DataType = DataType.UNKNOWN
    process_group: Optional[ProcessGroup] = None
    sequence_number: Optional[int] = None
    algorithm: Optional[str] = None
    protocol: Optional[str] = None
    num_channels: Optional[int] = None
    correlation_id: Optional[int] = None
    pytorch_op_name: Optional[str] = None
    pytorch_stack_trace: Optional[str] = None
    source_file: Optional[str] = None
    source_line: Optional[int] = None
    extra: dict = field(default_factory=dict)
    @property
    def duration_ns(self) -> int:
        """Duration of this collective in nanoseconds."""
        return self.end_time_ns - self.start_time_ns
    @property
    def duration_ms(self) -> float:
        """Duration of this collective in milliseconds."""
        return self.duration_ns / 1_000_000
    @property
    def duration_us(self) -> float:
        """Duration of this collective in microseconds."""
        return self.duration_ns / 1_000
    def __post_init__(self) -> None:
        """Validate collective data."""
        assert self.rank >= 0, f"Rank must be non-negative, got {self.rank}"
        assert self.start_time_ns >= 0, f"Start time must be non-negative, got {self.start_time_ns}"
        assert (
            self.end_time_ns >= self.start_time_ns
        ), f"End time ({self.end_time_ns}) must be >= start time ({self.start_time_ns})"
        assert (
            self.message_size_bytes >= 0
        ), f"Message size must be non-negative, got {self.message_size_bytes}"
