from __future__ import annotations

from typing import Any, Optional

COMM_OP_TYPE: Any
NON_COMP_TYPE: Any
OP_COST_FACTORY: Any

def parse_to_desc(op: Optional[Any] = ..., dist_op: Optional[Any] = ..., dist_context: Optional[Any] = ...): ...
def parse_desc_to_str(desc: Any): ...

class CommContext:
    cluster: Any = ...
    def __init__(self, cluster: Any) -> None: ...
    def __new__(cls, *args: Any, **kwargs: Any): ...
    @property
    def alpha_inter(self): ...
    @property
    def alpha_intra(self): ...
    @property
    def alpha_base_ring(self): ...
    @property
    def alpha_base_tree(self): ...
    def get_beta(self, ranks: Any): ...

class Cost:
    def __init__(self, time: int = ..., memory: int = ..., flops: int = ...) -> None: ...
    @property
    def time(self): ...
    @time.setter
    def time(self, val: Any) -> None: ...
    @property
    def memory(self): ...
    @memory.setter
    def memory(self, val: Any) -> None: ...
    @property
    def flops(self): ...
    @flops.setter
    def flops(self, val: Any) -> None: ...
    def __add__(self, rhs: Any): ...
    def __sub__(self, rhs: Any): ...

class OpCost:
    def __init__(self, op: Optional[Any] = ..., op_desc: Optional[Any] = ...) -> None: ...
    @property
    def op(self): ...
    @property
    def op_desc(self): ...
    @property
    def cost(self): ...
    def calc_time(self): ...
    def calc_memory(self): ...
    def calc_flops(self): ...
    def calc_cost(self): ...

class CommOpCost(OpCost):
    OP_TYPE: str = ...
    def __init__(
        self, op: Optional[Any] = ..., op_desc: Optional[Any] = ..., comm_context: Optional[Any] = ...
    ) -> None: ...
    @property
    def comm_context(self): ...

class CompOpCost(OpCost):
    OP_TYPE: str = ...
    cluster: Any = ...
    def __init__(self, op: Optional[Any] = ..., op_desc: Optional[Any] = ..., cluster: Optional[Any] = ...) -> None: ...

def register_op_cost(cls): ...
def calc_time_from_model(
    op: Optional[Any] = ..., desc: Optional[Any] = ..., cluster: Optional[Any] = ..., comm_context: Optional[Any] = ...
): ...
