"""Multi-rank trace session container.
Defines data structures for representing a complete trace session with
multiple ranks and the results of analysis.
"""
from dataclasses import dataclass, field
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.rank_timeline import RankTimeline
class TraceFormat(Enum):
    """Supported trace file formats."""
    PYTORCH_PROFILER = "pytorch_profiler"  # Chrome trace JSON from torch.profiler
    NCCL_INSPECTOR = "nccl_inspector"  # NCCL Inspector JSONL
    ROCPROFV3 = "rocprofv3"  # rocprofv3 CSV/JSON
    NSYS = "nsys"  # Nsight Systems SQLite
    ROCSYS = "rocsys"  # ROCm Systems Profiler
    UNKNOWN = "unknown"
    @classmethod
    def detect_format(cls, file_path: str) -> "TraceFormat":
        """Auto-detect trace format from file path and content."""
        import json
        from pathlib import Path
        path = Path(file_path)
        suffix = path.suffix.lower()

        if suffix in [".nsys-rep", ".sqlite"]:
            return cls.NSYS
        if suffix in [".rocpd", ".db"]:
            return cls.ROCSYS
        if suffix == ".jsonl":
            return cls.NCCL_INSPECTOR
        if suffix == ".csv":
            # Could be rocprofv3 or other - check header
            try:
                with open(path, encoding="utf-8") as f:
                    header = f.readline().strip().lower()
                    if "rccl" in header or "domain" in header:
                        return cls.ROCPROFV3
            except Exception:
                pass
            return cls.ROCPROFV3  # Default for CSV
        if suffix in [".json", ".gz"]:

            try:
                with open(path, encoding="utf-8") as f:
                    # Read first few KB to detect format
                    content = f.read(4096)
                    data = json.loads(content) if not content.startswith("{") else None
                    # PyTorch profiler has traceEvents array
                    if "traceEvents" in content or "trace_events" in content:
                        return cls.PYTORCH_PROFILER
                    # NCCL Inspector has specific event types
                    if '"op"' in content and '"comm_id"' in content:
                        return cls.NCCL_INSPECTOR
            except Exception:
                pass
            # Default to PyTorch profiler for JSON
            return cls.PYTORCH_PROFILER
        return cls.UNKNOWN
class HangRiskLevel(Enum):
    """Risk level for potential hang detection."""
    NORMAL = "normal"  # < 1% of timeout
    ELEVATED = "elevated"  # 1-10% of timeout
    WARNING = "warning"  # 10-50% of timeout
    CRITICAL = "critical"  # > 50% of timeout
@dataclass(frozen=True)
class CollectiveMatch:
    """Represents a matched collective operation across ranks.
    When the same logical collective is identified across multiple ranks,
    this structure holds the matching collectives for analysis.
    Attributes:
        collective_type: Type of the collective
        sequence_number: Sequence number in the process group
        process_group_id: ID of the process group
        collectives: Map of rank -> Collective for this match
        message_size_bytes: Message size (should be same across ranks)
    """
    collective_type: CollectiveType
    sequence_number: int
    process_group_id: str
    collectives: dict[int, Collective]  # rank -> Collective
    message_size_bytes: int = 0
    @property
    def num_ranks(self) -> int:
        """Number of ranks in this match."""
        return len(self.collectives)
    @property
    def ranks(self) -> list[int]:
        """List of ranks in this match."""
        return sorted(self.collectives.keys())
    @property
    def min_duration_ns(self) -> int:
        """Minimum duration across ranks."""
        return min(c.duration_ns for c in self.collectives.values())
    @property
    def max_duration_ns(self) -> int:
        """Maximum duration across ranks."""
        return max(c.duration_ns for c in self.collectives.values())
    @property
    def median_duration_ns(self) -> int:
        """Median duration across ranks."""
        durations = sorted(c.duration_ns for c in self.collectives.values())
        n = len(durations)
        if n % 2 == 0:
            return (durations[n // 2 - 1] + durations[n // 2]) // 2
        return durations[n // 2]
    @property
    def straggler_rank(self) -> int:
        """Rank with the longest duration (straggler)."""
        return max(self.collectives.items(), key=lambda x: x[1].duration_ns)[0]
    @property
    def straggler_impact_ns(self) -> int:
        """Impact of straggler: (max - median) * (num_ranks - 1)."""
        return (self.max_duration_ns - self.median_duration_ns) * (self.num_ranks - 1)
@dataclass(frozen=True)
class BandwidthMetrics:
    """Bandwidth metrics for a collective or aggregated analysis.
    Attributes:
        achieved_bandwidth_gbps: Achieved bandwidth in GB/s
        theoretical_bandwidth_gbps: Theoretical max bandwidth in GB/s
        efficiency_percent: Achieved / theoretical * 100
        message_size_bytes: Message size
        duration_ns: Duration
        algorithm: Algorithm used (if known)
    """
    achieved_bandwidth_gbps: float
    theoretical_bandwidth_gbps: float
    efficiency_percent: float
    message_size_bytes: int
    duration_ns: int
    algorithm: Optional[str] = None
    @classmethod
    def calculate(
        cls,
        message_size_bytes: int,
        duration_ns: int,
        theoretical_bandwidth_gbps: float,
        algorithm: Optional[str] = None,
    ) -> "BandwidthMetrics":
        """Calculate bandwidth metrics from message size and duration."""
        duration_s = duration_ns / 1_000_000_000
        achieved_gbps = (message_size_bytes / duration_s / 1_000_000_000) if duration_s > 0 else 0
        efficiency = (achieved_gbps / theoretical_bandwidth_gbps * 100) if theoretical_bandwidth_gbps > 0 else 0
        return cls(
            achieved_bandwidth_gbps=achieved_gbps,
            theoretical_bandwidth_gbps=theoretical_bandwidth_gbps,
            efficiency_percent=efficiency,
            message_size_bytes=message_size_bytes,
            duration_ns=duration_ns,
            algorithm=algorithm,
        )
@dataclass(frozen=True)
class OverlapMetrics:
    """Communication/compute overlap metrics.
    Attributes:
        total_compute_time_ns: Total time in compute operations
        total_communication_time_ns: Total time in communication
        overlapped_time_ns: Time when both compute and communication active
        compute_only_time_ns: Time with only compute active
        communication_only_time_ns: Time with only communication active
        idle_time_ns: Time with neither active
        overlap_ratio: overlapped_time / communication_time (0-1)
        iteration_metrics: Per-iteration breakdown (if available)
    """
    total_compute_time_ns: int
    total_communication_time_ns: int
    overlapped_time_ns: int
    compute_only_time_ns: int
    communication_only_time_ns: int
    idle_time_ns: int
    overlap_ratio: float
    iteration_metrics: list["IterationOverlap"] = field(default_factory=list)
    @property
    def overlap_percent(self) -> float:
        """Overlap ratio as percentage."""
        return self.overlap_ratio * 100
@dataclass(frozen=True)
class IterationOverlap:
    """Overlap metrics for a single iteration.
    Attributes:
        iteration: Iteration number
        start_time_ns: Start of iteration
        end_time_ns: End of iteration
        compute_time_ns: Compute time in this iteration
        communication_time_ns: Communication time
        overlapped_time_ns: Overlapped time
        overlap_ratio: Overlap ratio for this iteration
    """
    iteration: int
    start_time_ns: int
    end_time_ns: int
    compute_time_ns: int
    communication_time_ns: int
    overlapped_time_ns: int
    overlap_ratio: float
@dataclass(frozen=True)
class StragglerReport:
    """Straggler detection report.
    Attributes:
        rank_scores: Per-rank cumulative straggler scores
        flagged_ranks: Ranks exceeding the straggler threshold
        top_collectives: Most impactful collective instances
        total_straggler_impact_ns: Total time lost to stragglers
        is_consistent: True if same rank is consistently slow
    """
    rank_scores: dict[int, float]  # rank -> cumulative score
    flagged_ranks: list[int]  # ranks exceeding threshold
    top_collectives: list[tuple[CollectiveMatch, int]]  # (match, rank) for top impacts
    total_straggler_impact_ns: int
    is_consistent: bool  # Same rank always slow
@dataclass(frozen=True)
class HangRiskReport:
    """Hang risk assessment report.
    Attributes:
        timeout_threshold_ns: Configured timeout threshold
        collectives_by_risk: Collectives grouped by risk level
        highest_risk_collective: Collective with highest risk score
        risk_distribution: Count of collectives at each risk level
        recommendations: Actionable recommendations
    """
    timeout_threshold_ns: int
    collectives_by_risk: dict[HangRiskLevel, list[Collective]]
    highest_risk_collective: Optional[Collective]
    risk_distribution: dict[HangRiskLevel, int]
    recommendations: list[str]
@dataclass(frozen=True)
class AnalysisResult:
    """Complete analysis result for a trace session.
    Attributes:
        bandwidth_metrics: Per-collective bandwidth analysis
        aggregate_bandwidth: Aggregated bandwidth metrics
        overlap_metrics: Communication/compute overlap analysis
        straggler_report: Straggler detection report
        hang_risk_report: Hang risk assessment
        collective_matches: Matched collectives across ranks
        summary: Summary statistics
    """
    bandwidth_metrics: dict[int, BandwidthMetrics]  # collective index -> metrics
    aggregate_bandwidth: BandwidthMetrics
    overlap_metrics: OverlapMetrics
    straggler_report: StragglerReport
    hang_risk_report: HangRiskReport
    collective_matches: list[CollectiveMatch]
    summary: dict[str, object]
@dataclass(frozen=True)
class ComparisonResult:
    """Result of comparing two trace sessions.
    Attributes:
        baseline_summary: Summary of baseline trace
        comparison_summary: Summary of comparison trace
        bandwidth_delta: Change in bandwidth metrics
        overlap_delta: Change in overlap metrics
        straggler_delta: Change in straggler impact
        per_collective_deltas: Per-collective type changes
        regressions: Metrics that got worse
        improvements: Metrics that improved
        recommendations: Recommendations based on comparison
    """
    baseline_summary: dict[str, object]
    comparison_summary: dict[str, object]
    bandwidth_delta: float  # Percentage change
    overlap_delta: float  # Percentage change in overlap ratio
    straggler_delta: float  # Percentage change in straggler impact
    per_collective_deltas: dict[CollectiveType, dict[str, float]]
    regressions: list[str]
    improvements: list[str]
    recommendations: list[str]
@dataclass
class TraceSession:
    """Multi-rank trace session container.
    Holds all rank timelines and provides methods for cross-rank analysis.
    Note: This class is mutable to allow incremental building.
    Attributes:
        name: Session name (e.g., training run identifier)
        trace_format: Format of the source traces
        world_size: Total number of ranks
        timelines: Per-rank timelines
        metadata: Additional metadata
        aligned: Whether clock alignment has been performed
    """
    name: str = ""
    trace_format: TraceFormat = TraceFormat.UNKNOWN
    world_size: int = 0
    timelines: dict[int, RankTimeline] = field(default_factory=dict)  # rank -> timeline
    metadata: dict[str, object] = field(default_factory=dict)
    aligned: bool = False
    @property
    def ranks(self) -> list[int]:
        """List of ranks in this session."""
        return sorted(self.timelines.keys())
    @property
    def num_ranks(self) -> int:
        """Number of ranks with timelines."""
        return len(self.timelines)
    @property
    def total_collectives(self) -> int:
        """Total number of collectives across all ranks."""
        return sum(len(t.collectives) for t in self.timelines.values())
    @property
    def start_time_ns(self) -> int:
        """Earliest timestamp across all ranks."""
        if not self.timelines:
            return 0
        return min(t.start_time_ns for t in self.timelines.values())
    @property
    def end_time_ns(self) -> int:
        """Latest timestamp across all ranks."""
        if not self.timelines:
            return 0
        return max(t.end_time_ns for t in self.timelines.values())
    @property
    def duration_ns(self) -> int:
        """Total duration of the session."""
        return self.end_time_ns - self.start_time_ns
    def add_timeline(self, timeline: RankTimeline) -> None:
        """Add a timeline for a rank."""
        assert timeline is not None
        assert timeline.rank >= 0
        self.timelines[timeline.rank] = timeline
        self.world_size = max(self.world_size, len(self.timelines))
    def get_timeline(self, rank: int) -> Optional[RankTimeline]:
        """Get timeline for a specific rank."""
        return self.timelines.get(rank)
    def get_all_collectives(self) -> list[Collective]:
        """Get all collectives from all ranks."""
        collectives = []
        for timeline in self.timelines.values():
            collectives.extend(timeline.collectives)
        return collectives
    def get_collectives_by_type(self, collective_type: CollectiveType) -> list[Collective]:
        """Get all collectives of a specific type across all ranks."""
        collectives = []
        for timeline in self.timelines.values():
            collectives.extend(
                c for c in timeline.collectives if c.collective_type == collective_type
            )
        return collectives
