"""Data models for distributed traces analysis.

This module provides frozen dataclasses representing collective operations,
rank timelines, and multi-rank trace sessions.
"""

from wafer_core.lib.distributed_traces.models.collective import (
    Collective,
    CollectiveType,
    DataType,
    ProcessGroup,
)
from wafer_core.lib.distributed_traces.models.rank_timeline import (
    ComputeEvent,
    RankTimeline,
    TimelineEvent,
    TimelineWindow,
    WindowType,
)
from wafer_core.lib.distributed_traces.models.trace_session import (
    AnalysisResult,
    BandwidthMetrics,
    CollectiveMatch,
    ComparisonResult,
    HangRiskLevel,
    HangRiskReport,
    OverlapMetrics,
    StragglerReport,
    TraceFormat,
    TraceSession,
)

__all__ = [
    # Collective types
    "Collective",
    "CollectiveType",
    "DataType",
    "ProcessGroup",
    # Timeline types
    "ComputeEvent",
    "RankTimeline",
    "TimelineEvent",
    "TimelineWindow",
    "WindowType",
    # Session and analysis types
    "AnalysisResult",
    "BandwidthMetrics",
    "CollectiveMatch",
    "ComparisonResult",
    "HangRiskLevel",
    "HangRiskReport",
    "OverlapMetrics",
    "StragglerReport",
    "TraceFormat",
    "TraceSession",
]
