from __future__ import annotations

from enum import Enum
from typing import Any, Callable, Iterable, Optional, Union

from .profiler_statistic import SortedKeys as SortedKeys
from .profiler_statistic import StatisticData as StatisticData
from .timer import benchmark as benchmark
from .utils import RecordEvent as RecordEvent
from .utils import wrap_optimizers as wrap_optimizers

class ProfilerState(Enum):
    CLOSED = ...
    READY = ...
    RECORD = ...
    RECORD_AND_RETURN = ...

class ProfilerTarget(Enum):
    CPU = ...
    GPU = ...
    MLU = ...

def make_scheduler(closed: int, ready: int, record: int, *, repeat: int = ..., skip_first: int = ...) -> Callable: ...
def export_chrome_tracing(dir_name: str, worker_name: Optional[str] = ...) -> Callable: ...
def export_protobuf(dir_name: str, worker_name: Optional[str] = ...) -> Callable: ...

class Profiler:
    targets: Any = ...
    profiler: Any = ...
    scheduler: Any = ...
    on_trace_ready: Any = ...
    step_num: int = ...
    previous_state: Any = ...
    current_state: Any = ...
    record_event: Any = ...
    profiler_result: Any = ...
    timer_only: Any = ...
    def __init__(
        self,
        *,
        targets: Optional[Iterable[ProfilerTarget]] = ...,
        scheduler: Union[Callable[[int], ProfilerState], tuple, None] = ...,
        on_trace_ready: Optional[Callable[..., Any]] = ...,
        timer_only: Optional[bool] = ...,
    ) -> None: ...
    def __enter__(self): ...
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: ...
    def start(self) -> None: ...
    def stop(self) -> None: ...
    def step(self, num_samples: Optional[int] = ...) -> Any: ...
    def step_info(self, unit: Optional[Any] = ...): ...
    def export(self, path: str = ..., format: str = ...) -> None: ...
    def summary(
        self, sorted_by: Any = ..., op_detail: bool = ..., thread_sep: bool = ..., time_unit: str = ...
    ) -> None: ...

def get_profiler(config_path: Any): ...
