from __future__ import annotations

from typing import Any, Optional

from .. import core as core
from .. import framework as framework
from .. import initializer as initializer
from .. import unique_name as unique_name
from ..distribute_lookup_table import (
    find_distributed_lookup_table as find_distributed_lookup_table,
)
from ..framework import Block as Block
from ..framework import Parameter as Parameter
from ..framework import Program as Program
from ..framework import default_main_program as default_main_program
from ..framework import default_startup_program as default_startup_program
from ..framework import grad_var_name as grad_var_name
from . import collective as collective
from .details import UnionFind as UnionFind
from .details import VarsDistributed as VarsDistributed
from .details import VarStruct as VarStruct
from .details import delete_ops as delete_ops
from .details import find_op_by_output_arg as find_op_by_output_arg
from .details import wait_server_ready as wait_server_ready
from .ps_dispatcher import PSDispatcher as PSDispatcher
from .ps_dispatcher import RoundRobin as RoundRobin

LOOKUP_TABLE_TYPE: Any
LOOKUP_TABLE_GRAD_TYPE: Any
OP_NAME_SCOPE: str
CLIP_OP_NAME_SCOPE: str
OP_ROLE_VAR_ATTR_NAME: Any
RPC_OP_ROLE_ATTR_NAME: Any
op_role_attr_name: Any
OPT_OP_ROLE_ATTR_VALUE: Any
RPC_OP_ROLE_ATTR_VALUE: Any
DIST_OP_ROLE_ATTR_VALUE: Any
LR_SCHED_OP_ROLE_ATTR_VALUE: Any
PRINT_LOG: bool

class DistributedMode:
    SYNC: int = ...
    ASYNC: int = ...
    HALF_ASYNC: int = ...
    GEO: int = ...

def log(*args: Any) -> None: ...

class VarBlock:
    varname: Any = ...
    offset: Any = ...
    size: Any = ...
    def __init__(self, varname: Any, offset: Any, size: Any) -> None: ...

def same_or_split_var(p_name: Any, var_name: Any): ...
def slice_variable(var_list: Any, slice_count: Any, min_block_size: Any): ...

class DistributeTranspilerConfig:
    slice_var_up: bool = ...
    split_method: Any = ...
    min_block_size: int = ...
    enable_dc_asgd: bool = ...
    mode: str = ...
    print_log: bool = ...
    wait_port: bool = ...
    half_async: bool = ...
    completely_not_async: bool = ...
    geo_sgd_mode: bool = ...
    geo_sgd_need_push_nums: int = ...
    nccl_comm_num: int = ...
    use_hierarchical_allreduce: bool = ...
    hierarchical_allreduce_inter_nranks: int = ...
    collective_mode: Any = ...
    def __init__(self) -> None: ...
    @property
    def runtime_split_send_recv(self): ...
    @runtime_split_send_recv.setter
    def runtime_split_send_recv(self, value: Any) -> None: ...
    @property
    def sync_mode(self): ...
    @sync_mode.setter
    def sync_mode(self, value: Any) -> None: ...

class ServerRuntimeConfig:
    def __init__(self) -> None: ...

class DistributeTranspiler:
    config: Any = ...
    distributed_mode: Any = ...
    counter_var: Any = ...
    def __init__(self, config: Optional[Any] = ...) -> None: ...
    origin_program: Any = ...
    startup_program: Any = ...
    origin_startup_program: Any = ...
    trainer_num: Any = ...
    sync_mode: Any = ...
    trainer_id: Any = ...
    pserver_endpoints: Any = ...
    vars_overview: Any = ...
    table_name: Any = ...
    has_distributed_lookup_table: Any = ...
    param_name_to_grad_name: Any = ...
    grad_name_to_param_name: Any = ...
    sparse_update_ops: Any = ...
    sparse_param_to_height_sections: Any = ...
    need_delete_optimize_vars: Any = ...
    grad_name_to_send_dummy_out: Any = ...
    def transpile(
        self,
        trainer_id: Any,
        program: Optional[Any] = ...,
        pservers: str = ...,
        trainers: int = ...,
        sync_mode: bool = ...,
        startup_program: Optional[Any] = ...,
        current_endpoint: str = ...,
    ) -> None: ...
    def get_trainer_program(self, wait_port: bool = ...): ...
    param_bak_list: Any = ...
    pserver_program: Any = ...
    def get_pserver_program(self, endpoint: Any): ...
    def get_pserver_programs(self, endpoint: Any): ...
    def get_startup_program(
        self, endpoint: Any, pserver_program: Optional[Any] = ..., startup_program: Optional[Any] = ...
    ): ...
