from bagua.torch_api.distributed import BaguaModule
from bagua.torch_api.bucket import BaguaBucket
from bagua.torch_api.tensor import BaguaTensor
from typing import List
import torch


class Algorithm:
    """
    This is the base class that all Bagua algorithms inherit.

    It provides methods that can be override to implement different kinds of
    distributed algorithms.
    """

    def need_reset(self) -> bool:
        """
        Returns:
             ``True`` if all initialization methods of the current algorithms should be called again. \
             This is useful for algorithms that have multiple stages where each stage needs different initializations.
        """
        return False

    def init_tensors(self, bagua_module: BaguaModule) -> List[BaguaTensor]:
        """
        Given a :class:`~bagua.torch_api.distributed.BaguaModule`, return Bagua tensors to be used in Bagua for later
        operations.

        Args:
            bagua_module: A PyTorch module initialized by
                :meth:`~bagua.torch_api.distributed.BaguaModule.with_bagua` method.

        Returns:
            A list of Bagua tensors for communication.
        """

        parameters = bagua_module.bagua_build_params()
        tensors = []
        for name, param in parameters.__reversed__():
            grad = param.bagua_ensure_grad().ensure_bagua_tensor(
                name, bagua_module.bagua_module_name
            )
            param._bagua_grad = grad
            tensors.append(grad)
        self._communication_tensor_names = set(name for name, _ in parameters)
        assert len(self._communication_tensor_names) == len(
            tensors
        ), "tensor names should be unique"
        return tensors

    def tensors_to_buckets(self, tensors: List[List[BaguaTensor]]) -> List[BaguaBucket]:
        """
        Given the bucketing suggestion from Bagua, return the actual Bagua buckets.
        The default implementation follows the suggestion to do the bucketing.

        Args:
            tensors: Bagua tensors grouped in different
                lists, representing Bagua's suggestion on how to bucketing the
                tensors.

        Returns:
            A list of Bagua buckets.
        """
        bagua_buckets = []
        for idx, bucket in enumerate(tensors):
            bagua_bucket = BaguaBucket(
                bucket, flatten=True, name=str(idx)
            )  # TODO: check duplicated names
            bagua_buckets.append(bagua_bucket)
        return bagua_buckets

    def init_forward_pre_hook(self, bagua_module: BaguaModule):
        """Given a :class:`~bagua.torch_api.distributed.BaguaModule`, return a hook function that will be executed before the
        forward process.

        Args:
            bagua_module: A PyTorch module initialized by
                :meth:`~bagua.torch_api.distributed.BaguaModule.with_bagua` method.

        Returns:
            A function that takes the model's input.
        """

        def hook(input):
            pass

        return hook

    def init_backward_hook(self, bagua_module: BaguaModule):
        """Given a :class:`~bagua.torch_api.distributed.BaguaModule`, return a hook function that will be executed on every
        parameter's gradient computation completion.

        Args:
            bagua_module: A PyTorch module initialized by
                :meth:`~bagua.torch_api.distributed.BaguaModule.with_bagua` method.

        Returns:
            A function that takes the name of a parameter (as in ``torch.nn.Module.named_parameters``) and the parameter itself.
        """

        def hook(parameter_name, parameter):
            if parameter_name in self._communication_tensor_names:
                assert (
                    parameter._bagua_grad.data_ptr() == parameter.grad.data_ptr()
                ), "bagua grad data_ptr should match parameter grad"
                parameter._bagua_grad.bagua_mark_communication_ready()

        return hook

    def init_post_backward_hook(self, bagua_module: BaguaModule):
        """Given a :class:`~bagua.torch_api.distributed.BaguaModule`, return a hook function that will be executed when the
        backward pass is done.

        Args:
            bagua_module: A PyTorch module initialized by
                :meth:`~bagua.torch_api.distributed.BaguaModule.with_bagua` method.

        Returns:
            A function that takes no argument.
        """

        def hook():
            bagua_module._bagua_backend.wait_pending_comm_ops()

        return hook

    def init_post_optimizer_step_hook(self, bagua_module: BaguaModule):
        """Given a :class:`~bagua.torch_api.distributed.BaguaModule`, return a hook function that will be executed when the
        ``optimizer.step()`` is done.

        Args:
            bagua_module: A PyTorch module initialized by
                :meth:`~bagua.torch_api.distributed.BaguaModule.with_bagua` method.

        Returns:
            A function that gets called after an optimizer's ``step()`` method is called. The function takes the optimizer as its argument.
        """

        def hook(optimizer: torch.optim.Optimizer):
            pass

        return hook

    def init_operations(
        self,
        bagua_module: BaguaModule,
        bucket: BaguaBucket,
    ):
        """Given a :class:`~bagua.torch_api.distributed.BaguaModule`, and a :class:`~bagua.torch_api.bucket.BaguaBucket`,
        register operations to be executed on the bucket.

        Args:
            bagua_module: A PyTorch module initialized by
                :meth:`~bagua.torch_api.distributed.BaguaModule.with_bagua` method.
            bucket: A single bucket to register operations.
        """
