from __future__ import annotations

from collections import defaultdict as defaultdict
from typing import Any, Optional

from paddle.fluid.framework import get_flags as get_flags

from .dist_attribute import OperatorDistributedAttribute as OperatorDistributedAttribute
from .dist_attribute import TensorDistributedAttribute as TensorDistributedAttribute
from .dist_op import DistributedOperator as DistributedOperator
from .dist_tensor import DistributedTensor as DistributedTensor
from .process_mesh import ProcessMesh as ProcessMesh

def get_default_distributed_context(): ...
def set_default_distributed_context(dist_context: Any) -> None: ...

class DistributedContext:
    def __init__(
        self,
        serial_main_prog: Optional[Any] = ...,
        serial_startup_prog: Optional[Any] = ...,
        dist_main_progs: Optional[Any] = ...,
        dist_startup_progs: Optional[Any] = ...,
    ) -> None: ...
    @property
    def serial_program(self): ...
    @property
    def serial_graph(self): ...
    @serial_program.setter
    def serial_program(self, program: Any) -> None: ...
    @property
    def serial_ordered_nodes(self): ...
    @property
    def process_meshes(self): ...
    @property
    def dist_op_context(self): ...
    @property
    def block_state(self): ...
    @property
    def dist_main_programs(self): ...
    @property
    def dist_startup_programs(self): ...
    def add_process_mesh(self, process_mesh: Any) -> None: ...
    def add_dist_tensor_for_program(self, dist_tensor: Any) -> None: ...
    def add_dist_op_for_program(self, dist_op: Any) -> None: ...
    def get_dist_tensor_for_program(self, serial_tensor: Any): ...
    def get_dist_tensor_for_graph(self, serial_tensor_node: Any): ...
    def get_dist_op_for_program(self, serial_op: Any): ...
    def del_dist_op_for_program(self, serial_tensor: Any) -> None: ...
    def get_dist_op_for_graph(self, serial_op_node: Any): ...
    def get_tensor_dist_attr_for_program(self, serial_tensor: Any): ...
    def get_tensor_dist_attr_for_program_with_id(self, tensor_id: Any): ...
    def set_tensor_dist_attr_for_program(self, serial_tensor: Any, dist_attr: Any) -> None: ...
    def get_tensor_dist_attr_for_graph(self, serial_tensor_node: Any): ...
    def get_op_dist_attr_for_program(self, serial_op: Any): ...
    def get_op_dist_attr_for_program_with_id(self, op_id: Any): ...
    def set_op_dist_attr_for_program(self, serial_op: Any, dist_attr: Any) -> None: ...
    def get_op_dist_attr_for_graph(self, serial_op_node: Any): ...
    def get_dist_attr_for_graph(self, serial_node: Any): ...
    def init_dist_attr_for_program(self) -> None: ...
    def order_nodes_by_program_order(self): ...
    def init_dist_attr_for_graph(self) -> None: ...
    def clear_dist_info_for_program(self) -> None: ...
    def clear_dist_info_for_graph(self) -> None: ...
    def copy_dist_attr_from_graph_to_program(self) -> None: ...
    def amend_dist_attr_for_program(self) -> None: ...
    def validate_dist_attr_for_program(self): ...
    def __deepcopy__(self, memo: Any): ...

class DistributedOperatorContext:
    grad_op_id_to_op_id: Any = ...
    already_init_sync_vars: Any = ...
    varname_mapping: Any = ...
    rank_id: Any = ...
    def __init__(self) -> None: ...
    def __deepcopy__(self, memo: Any): ...
    @property
    def dst_main_program(self): ...
    @dst_main_program.setter
    def dst_main_program(self, prog: Any) -> None: ...
    @property
    def main_block(self): ...
    @property
    def dst_startup_program(self): ...
    @dst_startup_program.setter
    def dst_startup_program(self, prog: Any) -> None: ...
    @property
    def startup_block(self): ...
    @property
    def work_block(self): ...
    @work_block.setter
    def work_block(self, block: Any) -> None: ...
    @property
    def cur_src_op(self): ...
    def prepare_context(self, src_op: Any): ...

class BlockState:
    nblock: int = ...
    forward_indices: Any = ...
    backward_indices: Any = ...
    backward_to_forward_index_map: Any = ...
    def __init__(self) -> None: ...
    def parse_forward_blocks(self, program: Any) -> None: ...
    def parse_backward_blocks(self, program: Any) -> None: ...
