"""Distributed Traces Profiler - NCCL/RCCL Communication Analysis.

This module provides unified analysis for distributed training communication,
supporting both NVIDIA (NCCL) and AMD (RCCL) collective operations.

Submodules:
    models: Data structures for collectives, timelines, and trace sessions
    parsers: Format-specific parsing (PyTorch, NCCL Inspector, rocprofv3, NSYS, ROCm Systems)
    correlation: Cross-rank alignment and PyTorch-to-NCCL linking
    analysis: Metrics computation (bandwidth, overlap, straggler, hang risk)
    export: Output formats (JSON, CSV, timeline)
    collection: Wafer-managed trace collection setup

Example:
    from wafer_core.lib.distributed_traces import DistributedTracesAnalyzer

    analyzer = DistributedTracesAnalyzer()
    session, error = analyzer.load_traces(["trace_rank0.json", "trace_rank1.json"])
    if error:
        print(f"Error: {error}")
    else:
        result = analyzer.analyze(session)
        print(f"Overlap: {result.overlap_metrics.overlap_percent:.1f}%")
"""

from wafer_core.lib.distributed_traces.api import (
    AnalyzerConfig,
    DistributedTracesAnalyzer,
    analyze_distributed_traces,
)
from wafer_core.lib.distributed_traces.models import (
    AnalysisResult,
    Collective,
    CollectiveType,
    HangRiskLevel,
    RankTimeline,
    TraceFormat,
    TraceSession,
)

__all__ = [
    # Main API
    "AnalyzerConfig",
    "DistributedTracesAnalyzer",
    "analyze_distributed_traces",
    # Models
    "AnalysisResult",
    "Collective",
    "CollectiveType",
    "HangRiskLevel",
    "RankTimeline",
    "TraceFormat",
    "TraceSession",
]
