"""Achieved/theoretical bandwidth calculation.
Computes bandwidth metrics for collective operations including:
- Achieved bandwidth (bytes/second from actual transfer)
- Theoretical bandwidth (max based on hardware interconnect)
- Bandwidth efficiency (achieved / theoretical)
Hardware Bandwidth Reference:
- NVLink 3.0: 600 GB/s bidirectional per link
- NVLink 4.0: 900 GB/s bidirectional per link
- PCIe 4.0 x16: 32 GB/s per direction
- PCIe 5.0 x16: 64 GB/s per direction
- xGMI: 46 GB/s per link (MI250)
- InfiniBand HDR: 200 Gb/s (25 GB/s)
- InfiniBand NDR: 400 Gb/s (50 GB/s)
"""
from dataclasses import dataclass
from enum import Enum
from typing import Optional
from wafer_core.lib.distributed_traces.models.collective import Collective, CollectiveType
from wafer_core.lib.distributed_traces.models.trace_session import (
    BandwidthMetrics,
    CollectiveMatch,
    TraceSession,
)
class InterconnectType(Enum):
    """Type of interconnect for bandwidth calculation."""
    NVLINK_3 = "nvlink_3"  # A100
    NVLINK_4 = "nvlink_4"  # H100
    NVLINK_5 = "nvlink_5"  # B200
    PCIE_4 = "pcie_4"
    PCIE_5 = "pcie_5"
    XGMI = "xgmi"  # AMD MI250/MI300
    IB_HDR = "ib_hdr"  # InfiniBand HDR
    IB_NDR = "ib_ndr"  # InfiniBand NDR
    ETHERNET = "ethernet"
    UNKNOWN = "unknown"
# Theoretical bandwidth in GB/s for each interconnect type
THEORETICAL_BANDWIDTH_GBPS: dict[InterconnectType, float] = {
    InterconnectType.NVLINK_3: 600.0,
    InterconnectType.NVLINK_4: 900.0,
    InterconnectType.NVLINK_5: 1800.0,
    InterconnectType.PCIE_4: 32.0,
    InterconnectType.PCIE_5: 64.0,
    InterconnectType.XGMI: 400.0,  # MI300 has 8 links @ 50GB/s each
    InterconnectType.IB_HDR: 25.0,
    InterconnectType.IB_NDR: 50.0,
    InterconnectType.ETHERNET: 12.5,  # 100 GbE
    InterconnectType.UNKNOWN: 100.0,  # Conservative estimate
}
def get_theoretical_bandwidth(
    interconnect: InterconnectType = InterconnectType.UNKNOWN,
) -> float:
    """Get theoretical bandwidth for an interconnect type.
    Args:
        interconnect: Type of interconnect
    Returns:
        Theoretical bandwidth in GB/s
    """
    return THEORETICAL_BANDWIDTH_GBPS.get(interconnect, 100.0)
@dataclass
class BandwidthAnalysisResult:
    """Result of bandwidth analysis.
    Attributes:
        per_collective: Bandwidth metrics for each collective
        aggregate: Aggregated bandwidth across all collectives
        by_type: Aggregated bandwidth by collective type
        low_efficiency_collectives: Collectives with efficiency < 50%
    """
    per_collective: list[tuple[Collective, BandwidthMetrics]]
    aggregate: BandwidthMetrics
    by_type: dict[CollectiveType, BandwidthMetrics]
    low_efficiency_collectives: list[tuple[Collective, BandwidthMetrics]]
class BandwidthAnalyzer:
    """Analyzes bandwidth utilization of collective operations."""
    def __init__(
        self,
        theoretical_bandwidth_gbps: float = 100.0,
        efficiency_threshold: float = 50.0,
    ) -> None:
        """Initialize the analyzer.
        Args:
            theoretical_bandwidth_gbps: Theoretical max bandwidth in GB/s
            efficiency_threshold: Threshold below which to flag low efficiency
        """
        self.theoretical_bandwidth_gbps = theoretical_bandwidth_gbps
        self.efficiency_threshold = efficiency_threshold
    def analyze_session(self, session: TraceSession) -> BandwidthAnalysisResult:
        """Analyze bandwidth for all collectives in a session.
        Args:
            session: TraceSession to analyze
        Returns:
            BandwidthAnalysisResult with per-collective and aggregate metrics
        """
        per_collective = []
        by_type: dict[CollectiveType, list[BandwidthMetrics]] = {}
        low_efficiency = []
        for timeline in session.timelines.values():
            for collective in timeline.collectives:
                metrics = self.analyze_collective(collective)
                per_collective.append((collective, metrics))

                if collective.collective_type not in by_type:
                    by_type[collective.collective_type] = []
                by_type[collective.collective_type].append(metrics)
                # Flag low efficiency
                if metrics.efficiency_percent < self.efficiency_threshold:
                    low_efficiency.append((collective, metrics))
        aggregate = self._aggregate_metrics([m for _, m in per_collective])
        by_type_aggregate = {
            ctype: self._aggregate_metrics(metrics) for ctype, metrics in by_type.items()
        }
        return BandwidthAnalysisResult(
            per_collective=per_collective,
            aggregate=aggregate,
            by_type=by_type_aggregate,
            low_efficiency_collectives=low_efficiency,
        )
    def analyze_collective(self, collective: Collective) -> BandwidthMetrics:
        """Analyze bandwidth for a single collective.
        Args:
            collective: Collective to analyze
        Returns:
            BandwidthMetrics for this collective
        """
        return BandwidthMetrics.calculate(
            message_size_bytes=collective.message_size_bytes,
            duration_ns=collective.duration_ns,
            theoretical_bandwidth_gbps=self.theoretical_bandwidth_gbps,
            algorithm=collective.algorithm,
        )
    def analyze_match(self, match: CollectiveMatch) -> BandwidthMetrics:
        """Analyze bandwidth for a matched collective across ranks.
        Uses the median duration to compute bandwidth, as this represents
        typical behavior excluding straggler impact.
        Args:
            match: CollectiveMatch to analyze
        Returns:
            BandwidthMetrics using median duration
        """
        return BandwidthMetrics.calculate(
            message_size_bytes=match.message_size_bytes,
            duration_ns=match.median_duration_ns,
            theoretical_bandwidth_gbps=self.theoretical_bandwidth_gbps,
        )
    def _aggregate_metrics(self, metrics_list: list[BandwidthMetrics]) -> BandwidthMetrics:
        """Aggregate multiple bandwidth metrics.
        Args:
            metrics_list: List of BandwidthMetrics to aggregate
        Returns:
            Aggregated BandwidthMetrics
        """
        if not metrics_list:
            return BandwidthMetrics(
                achieved_bandwidth_gbps=0,
                theoretical_bandwidth_gbps=self.theoretical_bandwidth_gbps,
                efficiency_percent=0,
                message_size_bytes=0,
                duration_ns=0,
            )
        total_bytes = sum(m.message_size_bytes for m in metrics_list)
        total_duration_ns = sum(m.duration_ns for m in metrics_list)
        return BandwidthMetrics.calculate(
            message_size_bytes=total_bytes,
            duration_ns=total_duration_ns,
            theoretical_bandwidth_gbps=self.theoretical_bandwidth_gbps,
        )
def analyze_bandwidth(
    session: TraceSession,
    interconnect: InterconnectType = InterconnectType.UNKNOWN,
    efficiency_threshold: float = 50.0,
) -> BandwidthAnalysisResult:
    """Convenience function to analyze bandwidth in a session.
    Args:
        session: TraceSession to analyze
        interconnect: Type of interconnect for theoretical bandwidth
        efficiency_threshold: Threshold for flagging low efficiency
    Returns:
        BandwidthAnalysisResult
    """
    theoretical = get_theoretical_bandwidth(interconnect)
    analyzer = BandwidthAnalyzer(
        theoretical_bandwidth_gbps=theoretical,
        efficiency_threshold=efficiency_threshold,
    )
    return analyzer.analyze_session(session)
