"""High-level API for distributed traces analysis.
This module provides the main entry points for:
- Parsing traces from various formats
- Analyzing distributed training communication
- Generating reports and exports
Example usage:
    from wafer_core.lib.distributed_traces.api import DistributedTracesAnalyzer
    analyzer = DistributedTracesAnalyzer()
    session = analyzer.load_traces(["trace_rank0.json", "trace_rank1.json"])
    result = analyzer.analyze(session)
    print(result.summary)
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from wafer_core.lib.distributed_traces.analysis import (
    BandwidthAnalyzer,
    HangRiskAnalyzer,
    OverlapAnalyzer,
    StragglerAnalyzer,
    analyze_bandwidth,
    analyze_hang_risk,
    analyze_overlap,
    detect_stragglers,
)
from wafer_core.lib.distributed_traces.analysis.bandwidth import InterconnectType
from wafer_core.lib.distributed_traces.correlation import (
    CollectiveMatcher,
    RankAligner,
    align_ranks,
    match_collectives_across_ranks,
)
from wafer_core.lib.distributed_traces.export import (
    export_analysis_json,
    export_collectives_csv,
    export_session_json,
    export_timeline_json,
)
from wafer_core.lib.distributed_traces.models import (
    AnalysisResult,
    BandwidthMetrics,
    CollectiveMatch,
    HangRiskReport,
    OverlapMetrics,
    RankTimeline,
    StragglerReport,
    TraceFormat,
    TraceSession,
)
from wafer_core.lib.distributed_traces.parsers import (
    parse_nccl_inspector_file,
    parse_nsys_file,
    parse_pytorch_trace_file,
    parse_rocprofv3_file,
    parse_rocsys_file,
)
@dataclass
class AnalyzerConfig:
    """Configuration for the distributed traces analyzer.
    Attributes:
        interconnect: Type of interconnect for bandwidth calculation
        timeout_sec: NCCL timeout threshold in seconds
        align_ranks: Whether to align timestamps across ranks
        bandwidth_efficiency_threshold: Threshold for flagging low bandwidth
        straggler_std_devs: Standard deviations for straggler detection
    """
    interconnect: InterconnectType = InterconnectType.UNKNOWN
    timeout_sec: float = 600.0
    align_ranks: bool = True
    bandwidth_efficiency_threshold: float = 50.0
    straggler_std_devs: float = 2.0
class DistributedTracesAnalyzer:
    """Main analyzer for distributed traces.
    Provides a unified interface for loading, analyzing, and exporting
    distributed training traces from various formats.
    """
    def __init__(self, config: Optional[AnalyzerConfig] = None) -> None:
        """Initialize the analyzer.
        Args:
            config: Analyzer configuration
        """
        self.config = config or AnalyzerConfig()
    def load_traces(
        self,
        file_paths: list[str | Path],
        session_name: str = "",
    ) -> tuple[TraceSession | None, str | None]:
        """Load traces from multiple files.
        Auto-detects format for each file and parses into a unified session.
        Args:
            file_paths: List of trace file paths (one per rank)
            session_name: Optional name for the session
        Returns:
            Tuple of (TraceSession, error). One will be None.
        """
        if not file_paths:
            return None, "No trace files provided"
        session = TraceSession(name=session_name)
        errors = []
        for i, path in enumerate(file_paths):
            path = Path(path)
            if not path.exists():
                errors.append(f"File not found: {path}")
                continue
            # Detect format
            trace_format = TraceFormat.detect_format(str(path))
            timeline, error = self._parse_file(path, trace_format, rank=i)
            if error:
                errors.append(f"{path}: {error}")
                continue
            if timeline:
                session.add_timeline(timeline)
                if session.trace_format == TraceFormat.UNKNOWN:
                    session.trace_format = trace_format
        if not session.timelines:
            return None, "; ".join(errors) if errors else "No traces could be parsed"
        # Align ranks if configured
        if self.config.align_ranks and session.num_ranks > 1:
            align_ranks(session, apply=True)
        return session, None
    def _parse_file(
        self,
        path: Path,
        trace_format: TraceFormat,
        rank: Optional[int] = None,
    ) -> tuple[RankTimeline | None, str | None]:
        """Parse a single trace file.
        Args:
            path: Path to trace file
            trace_format: Detected format
            rank: Optional rank number
        Returns:
            Tuple of (RankTimeline, error)
        """
        parsers = {
            TraceFormat.PYTORCH_PROFILER: parse_pytorch_trace_file,
            TraceFormat.NCCL_INSPECTOR: parse_nccl_inspector_file,
            TraceFormat.ROCPROFV3: parse_rocprofv3_file,
            TraceFormat.NSYS: parse_nsys_file,
            TraceFormat.ROCSYS: parse_rocsys_file,
        }
        parser = parsers.get(trace_format)
        if parser is None:
            return None, f"Unknown trace format: {trace_format}"
        return parser(str(path), rank=rank)
    def analyze(self, session: TraceSession) -> AnalysisResult:
        """Perform full analysis on a trace session.
        Args:
            session: TraceSession to analyze
        Returns:
            AnalysisResult with all metrics
        """
        assert session is not None
        assert session.num_ranks > 0
        matches = match_collectives_across_ranks(session)
        # Bandwidth analysis
        bandwidth_result = analyze_bandwidth(
            session,
            interconnect=self.config.interconnect,
            efficiency_threshold=self.config.bandwidth_efficiency_threshold,
        )
        # Overlap analysis
        overlap_result = analyze_overlap(session)
        # Straggler detection
        straggler_report = detect_stragglers(
            session,
            matches=matches,
            threshold_std_devs=self.config.straggler_std_devs,
        )
        # Hang risk assessment
        hang_risk_report = analyze_hang_risk(
            session,
            timeout_sec=self.config.timeout_sec,
        )
        per_collective_bandwidth = {
            i: metrics
            for i, (_, metrics) in enumerate(bandwidth_result.per_collective)
        }
        # Generate summary
        summary = self._generate_summary(
            session,
            matches,
            bandwidth_result,
            overlap_result,
            straggler_report,
            hang_risk_report,
        )
        return AnalysisResult(
            bandwidth_metrics=per_collective_bandwidth,
            aggregate_bandwidth=bandwidth_result.aggregate,
            overlap_metrics=overlap_result.overall,
            straggler_report=straggler_report,
            hang_risk_report=hang_risk_report,
            collective_matches=matches,
            summary=summary,
        )
    def _generate_summary(
        self,
        session: TraceSession,
        matches: list[CollectiveMatch],
        bandwidth_result: "BandwidthAnalysisResult",  # noqa: F821
        overlap_result: "OverlapAnalysisResult",  # noqa: F821
        straggler_report: StragglerReport,
        hang_risk_report: HangRiskReport,
    ) -> dict:
        """Generate summary statistics."""
        from wafer_core.lib.distributed_traces.models.collective import CollectiveType
        # Count collectives by type
        type_counts: dict[str, int] = {}
        for timeline in session.timelines.values():
            for c in timeline.collectives:
                type_name = c.collective_type.value
                type_counts[type_name] = type_counts.get(type_name, 0) + 1
        return {
            "session_name": session.name,
            "trace_format": session.trace_format.value,
            "num_ranks": session.num_ranks,
            "world_size": session.world_size,
            "total_collectives": session.total_collectives,
            "collectives_by_type": type_counts,
            "duration_ms": session.duration_ns / 1_000_000,
            "num_matched_collectives": len(matches),
            "bandwidth_efficiency_percent": bandwidth_result.aggregate.efficiency_percent,
            "achieved_bandwidth_gbps": bandwidth_result.aggregate.achieved_bandwidth_gbps,
            "overlap_ratio": overlap_result.overall.overlap_ratio,
            "overlap_percent": overlap_result.overall.overlap_percent,
            "num_flagged_stragglers": len(straggler_report.flagged_ranks),
            "flagged_straggler_ranks": straggler_report.flagged_ranks,
            "total_straggler_impact_ms": straggler_report.total_straggler_impact_ns / 1_000_000,
            "hang_risk_distribution": {
                level.value: count
                for level, count in hang_risk_report.risk_distribution.items()
            },
            "has_critical_hang_risk": hang_risk_report.risk_distribution.get(
                HangRiskLevel.CRITICAL, 0
            ) > 0,
        }
    def export_json(
        self,
        result: AnalysisResult,
        session: TraceSession,
        output_path: Optional[str] = None,
    ) -> str:
        """Export analysis results to JSON.
        Args:
            result: AnalysisResult from analyze()
            session: TraceSession that was analyzed
            output_path: Optional path to write to
        Returns:
            JSON string
        """
        return export_analysis_json(result, session, output_path)
    def export_csv(
        self,
        session: TraceSession,
        output_path: Optional[str] = None,
    ) -> str:
        """Export collectives to CSV.
        Args:
            session: TraceSession to export
            output_path: Optional path to write to
        Returns:
            CSV string
        """
        return export_collectives_csv(session, output_path)
    def export_timeline(
        self,
        session: TraceSession,
        matches: Optional[list[CollectiveMatch]] = None,
        output_path: Optional[str] = None,
    ) -> str:
        """Export timeline data for visualization.
        Args:
            session: TraceSession to export
            matches: Optional collective matches
            output_path: Optional path to write to
        Returns:
            JSON string for timeline view
        """
        return export_timeline_json(session, matches, output_path)
    def compare_sessions(
        self,
        baseline: TraceSession,
        comparison: TraceSession,
    ) -> "ComparisonResult":  # noqa: F821
        """Compare two trace sessions.
        Args:
            baseline: Baseline session
            comparison: Session to compare against baseline
        Returns:
            ComparisonResult with deltas
        """
        from wafer_core.lib.distributed_traces.models.trace_session import ComparisonResult
        baseline_result = self.analyze(baseline)
        comparison_result = self.analyze(comparison)
        bw_baseline = baseline_result.aggregate_bandwidth.achieved_bandwidth_gbps
        bw_comparison = comparison_result.aggregate_bandwidth.achieved_bandwidth_gbps
        bw_delta = ((bw_comparison - bw_baseline) / bw_baseline * 100) if bw_baseline > 0 else 0
        overlap_baseline = baseline_result.overlap_metrics.overlap_ratio
        overlap_comparison = comparison_result.overlap_metrics.overlap_ratio
        overlap_delta = (
            (overlap_comparison - overlap_baseline) / overlap_baseline * 100
        ) if overlap_baseline > 0 else 0
        straggler_baseline = baseline_result.straggler_report.total_straggler_impact_ns
        straggler_comparison = comparison_result.straggler_report.total_straggler_impact_ns
        straggler_delta = (
            (straggler_comparison - straggler_baseline) / straggler_baseline * 100
        ) if straggler_baseline > 0 else 0

        regressions = []
        improvements = []
        if bw_delta < -5:
            regressions.append(f"Bandwidth decreased by {abs(bw_delta):.1f}%")
        elif bw_delta > 5:
            improvements.append(f"Bandwidth increased by {bw_delta:.1f}%")
        if overlap_delta < -5:
            regressions.append(f"Overlap ratio decreased by {abs(overlap_delta):.1f}%")
        elif overlap_delta > 5:
            improvements.append(f"Overlap ratio increased by {overlap_delta:.1f}%")
        if straggler_delta > 10:
            regressions.append(f"Straggler impact increased by {straggler_delta:.1f}%")
        elif straggler_delta < -10:
            improvements.append(f"Straggler impact decreased by {abs(straggler_delta):.1f}%")
        # Generate recommendations
        recommendations = []
        if regressions:
            recommendations.append("Performance regressions detected:")
            recommendations.extend([f"  - {r}" for r in regressions])
        if improvements:
            recommendations.append("Performance improvements:")
            recommendations.extend([f"  - {i}" for i in improvements])
        return ComparisonResult(
            baseline_summary=baseline_result.summary,
            comparison_summary=comparison_result.summary,
            bandwidth_delta=bw_delta,
            overlap_delta=overlap_delta,
            straggler_delta=straggler_delta,
            per_collective_deltas={},  # Could be extended for per-type analysis
            regressions=regressions,
            improvements=improvements,
            recommendations=recommendations,
        )
# Import HangRiskLevel for type hints
from wafer_core.lib.distributed_traces.models.trace_session import HangRiskLevel
def analyze_distributed_traces(
    file_paths: list[str | Path],
    config: Optional[AnalyzerConfig] = None,
) -> tuple[AnalysisResult | None, str | None]:
    """Convenience function to analyze distributed traces.
    Args:
        file_paths: List of trace file paths
        config: Optional analyzer configuration
    Returns:
        Tuple of (AnalysisResult, error). One will be None.
    """
    analyzer = DistributedTracesAnalyzer(config)
    session, error = analyzer.load_traces(file_paths)
    if error:
        return None, error
    assert session is not None
    result = analyzer.analyze(session)
    return result, None
