from __future__ import annotations

from enum import Enum
from typing import Any

from .statistic_helper import *

class SortedKeys(Enum):
    CPUTotal = ...
    CPUAvg = ...
    CPUMax = ...
    CPUMin = ...
    GPUTotal = ...
    GPUAvg = ...
    GPUMax = ...
    GPUMin = ...

class HostStatisticNode:
    hostnode: Any = ...
    children_node: Any = ...
    runtime_node: Any = ...
    cpu_time: int = ...
    self_cpu_time: int = ...
    gpu_time: int = ...
    self_gpu_time: int = ...
    general_gpu_time: int = ...
    self_general_gpu_time: int = ...
    is_terminal_operator_node: bool = ...
    def __init__(self, hostnode: Any) -> None: ...
    def cal_statistic(self) -> None: ...
    @property
    def end_ns(self): ...
    @property
    def start_ns(self): ...
    def __getattr__(self, name: Any): ...

def traverse_tree(nodetrees: Any): ...
def get_device_nodes(hostnode: Any): ...
def wrap_tree(nodetrees: Any): ...

class TimeRangeSummary:
    CPUTimeRange: Any = ...
    GPUTimeRange: Any = ...
    CPUTimeRangeSum: Any = ...
    GPUTimeRangeSum: Any = ...
    call_times: Any = ...
    def __init__(self): ...
    def parse(self, nodetrees: Any): ...
    def get_gpu_devices(self): ...
    def get_gpu_range_sum(self, device_id: Any, event_type: Any): ...
    def get_cpu_range_sum(self, event_type: Any): ...

class DistributedSummary:
    cpu_communication_range: Any = ...
    gpu_communication_range: Any = ...
    communication_range: Any = ...
    computation_range: Any = ...
    overlap_range: Any = ...
    cpu_calls: int = ...
    gpu_calls: int = ...
    def __init__(self) -> None: ...
    def parse(self, nodetrees: Any) -> None: ...

class EventSummary:
    class DeviceItem:
        name: Any = ...
        call: int = ...
        gpu_time: int = ...
        max_gpu_time: int = ...
        min_gpu_time: Any = ...
        def __init__(self, name: Any) -> None: ...
        @property
        def avg_gpu_time(self): ...
        def add_gpu_time(self, time: Any) -> None: ...
        def add_item(self, node: Any) -> None: ...

    class OperatorItem:
        name: Any = ...
        call: int = ...
        cpu_time: int = ...
        gpu_time: int = ...
        max_cpu_time: int = ...
        min_cpu_time: Any = ...
        max_gpu_time: int = ...
        min_gpu_time: Any = ...
        devices: Any = ...
        operator_inners: Any = ...
        general_gpu_time: int = ...
        min_general_gpu_time: Any = ...
        max_general_gpu_time: int = ...
        def __init__(self, name: Any) -> None: ...
        @property
        def avg_cpu_time(self): ...
        @property
        def avg_gpu_time(self): ...
        @property
        def avg_general_gpu_time(self): ...
        def add_cpu_time(self, time: Any) -> None: ...
        def add_gpu_time(self, time: Any) -> None: ...
        def add_general_gpu_time(self, time: Any) -> None: ...
        def add_call(self) -> None: ...
        def add_item(self, node: Any) -> None: ...

    class GeneralItem:
        name: Any = ...
        call: int = ...
        cpu_time: int = ...
        max_cpu_time: int = ...
        min_cpu_time: Any = ...
        gpu_time: int = ...
        max_gpu_time: int = ...
        min_gpu_time: Any = ...
        general_gpu_time: int = ...
        min_general_gpu_time: Any = ...
        max_general_gpu_time: int = ...
        def __init__(self, name: Any) -> None: ...
        @property
        def avg_cpu_time(self): ...
        @property
        def avg_gpu_time(self): ...
        @property
        def avg_general_gpu_time(self): ...
        def add_cpu_time(self, time: Any) -> None: ...
        def add_gpu_time(self, time: Any) -> None: ...
        def add_general_gpu_time(self, time: Any) -> None: ...
        def add_call(self) -> None: ...
        def add_item(self, node: Any) -> None: ...
    items: Any = ...
    thread_items: Any = ...
    userdefined_items: Any = ...
    userdefined_thread_items: Any = ...
    model_perspective_items: Any = ...
    memory_manipulation_items: Any = ...
    kernel_items: Any = ...
    def __init__(self) -> None: ...
    def parse(self, nodetrees: Any) -> None: ...
    def add_operator_item(self, operator_node: Any) -> None: ...
    def add_userdefined_item(self, userdefined_node: Any) -> None: ...
    def add_memory_manipulation_item(self, memory_manipulation_node: Any) -> None: ...
    def add_model_perspective_item(self, model_perspective_node: Any) -> None: ...
    def add_kernel_item(self, root_node: Any) -> None: ...

class StatisticData:
    node_trees: Any = ...
    extra_info: Any = ...
    time_range_summary: Any = ...
    event_summary: Any = ...
    distributed_summary: Any = ...
    def __init__(self, node_trees: Any, extra_info: Any) -> None: ...
