# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# Defined in torch/csrc/monitor/python_init.cpp

import datetime
from enum import Enum
from typing import Callable, Dict, List, Union

class Aggregation(Enum):
    VALUE = "value"
    MEAN = "mean"
    COUNT = "count"
    SUM = "sum"
    MAX = "max"
    MIN = "min"

class Stat:
    name: str
    count: int
    def __init__(
        self,
        name: str,
        aggregations: List[Aggregation],
        window_size: int,
        max_samples: int = -1,
    ) -> None: ...
    def add(self, v: float) -> None: ...
    def get(self) -> Dict[Aggregation, float]: ...

class Event:
    name: str
    timestamp: datetime.datetime
    data: Dict[str, Union[int, float, bool, str]]
    def __init__(
        self,
        name: str,
        timestamp: datetime.datetime,
        data: Dict[str, Union[int, float, bool, str]],
    ) -> None: ...

def log_event(e: Event) -> None: ...

class EventHandlerHandle: ...

def register_event_handler(handler: Callable[[Event], None]) -> EventHandlerHandle: ...
def unregister_event_handler(handle: EventHandlerHandle) -> None: ...
