from __future__ import annotations

from typing import Any, Optional

from .cost_model import estimate_cost as estimate_cost
from .dist_attribute import OperatorDistributedAttribute as OperatorDistributedAttribute
from .dist_attribute import TensorDistributedAttribute as TensorDistributedAttribute
from .dist_context import DistributedContext as DistributedContext
from .dist_context import DistributedOperatorContext as DistributedOperatorContext
from .dist_op import DistributedOperator as DistributedOperator
from .operators.common import (
    get_distributed_operator_impl_container as get_distributed_operator_impl_container,
)
from .operators.common import is_elementwise_op as is_elementwise_op
from .process_group import ProcessGroup as ProcessGroup
from .process_group import get_process_group as get_process_group
from .utils import get_all_distributed_main_program as get_all_distributed_main_program
from .utils import (
    update_op_dims_mapping_by_default_dist_impl as update_op_dims_mapping_by_default_dist_impl,
)
from .utils import (
    update_op_dims_mapping_by_elementwise_like_dist_impl as update_op_dims_mapping_by_elementwise_like_dist_impl,
)

class PlanFilter:
    @staticmethod
    def check_dims_mapping_for_tensor(process_mesh_topology: Any, tensor_shape: Any, dims_mapping: Any): ...
    @staticmethod
    def check_dims_mapping_for_op(op: Any, op_dist_attr: Any, vars: Any): ...
    @staticmethod
    def check_dims_mapping_for_special_op(op: Any, op_dist_attr: Any, vars: Any): ...

class PlanSpace:
    not_enum_ops: Any = ...
    special_vars: Any = ...
    @staticmethod
    def enum_process_mesh_topology(processes: Any): ...
    @staticmethod
    def enum_valid_dist_attr_for_program(program: Any, process_mesh_topology: Any, is_pipeline: bool = ...): ...

class SearchAlgorithm:
    def __init__(self, name: Any) -> None: ...
    @property
    def name(self) -> None: ...
    def search(self) -> None: ...

class MCMC(SearchAlgorithm):
    def __init__(self, serial_program_info: Any, parallelizer: Any, max_search_times: int = ...) -> None: ...
    @property
    def serial_program_info(self): ...
    @property
    def parallelizer(self): ...
    @property
    def max_search_times(self): ...
    def make_special_op_unshard(
        self, op: Any, ops: Any, vars: Any, dist_context: Any, valid_dist_attr_dict: Any
    ) -> None: ...
    def init_program(
        self, valid_dist_attr_dict: Any, program: Any, pipeline_process_meshes: Any, global_process_mesh: Any
    ): ...
    def estimate_searched_strategy_cost(self, dist_context: Any, pipeline_process_meshes: Optional[Any] = ...): ...
    def set_tensor_dist_attr(self, op: Any, op_dist_attr: Any, vars: Any, dist_context: Any) -> None: ...
    def change_process_mesh(self, op: Any, changed_process_mesh: Any, vars: Any, dist_context: Any) -> None: ...
    def search_once(
        self, program: Any, valid_dist_attr_dict: Any, dist_context: Any, pipeline_process_meshes: Optional[Any] = ...
    ): ...
    def search(self): ...

class Planner:
    def __init__(self, serial_program_info: Any, parallelizer: Any, algorithm_config: Optional[Any] = ...) -> None: ...
    @property
    def serial_program_info(self): ...
    @property
    def algorithm_config(self): ...
    @property
    def algorithm_searcher(self): ...
    @property
    def parallelizer(self): ...
    def create_algorithm_searcher(self, algorithm_config: Any): ...
    def search(self): ...
