"""Define distributed tensor."""

import os

from .. import backend as F, utils

from .dist_context import is_initialized
from .kvstore import get_kvstore
from .role import get_role
from .rpc import get_group_id


def _default_init_data(shape, dtype):
    return F.zeros(shape, dtype, F.cpu())


# These IDs can identify the anonymous distributed tensors.
DIST_TENSOR_ID = 0


class DistTensor:
    """Distributed tensor.

    ``DistTensor`` references to a distributed tensor sharded and stored in a cluster of machines.
    It has the same interface as Pytorch Tensor to access its metadata (e.g., shape and data type).
    To access data in a distributed tensor, it supports slicing rows and writing data to rows.
    It does not support any operators of a deep learning framework, such as addition and
    multiplication.

    Currently, distributed tensors are designed to store node data and edge data of a distributed
    graph. Therefore, their first dimensions have to be the number of nodes or edges in the graph.
    The tensors are sharded in the first dimension based on the partition policy of nodes
    or edges. When a distributed tensor is created, the partition policy is automatically
    determined based on the first dimension if the partition policy is not provided. If the first
    dimension matches the number of nodes of a node type, ``DistTensor`` will use the partition
    policy for this particular node type; if the first dimension matches the number of edges of
    an edge type, ``DistTensor`` will use the partition policy for this particular edge type.
    If DGL cannot determine the partition policy automatically (e.g., multiple node types or
    edge types have the same number of nodes or edges), users have to explicity provide
    the partition policy.

    A distributed tensor can be ether named or anonymous.
    When a distributed tensor has a name, the tensor can be persistent if ``persistent=True``.
    Normally, DGL destroys the distributed tensor in the system when the ``DistTensor`` object
    goes away. However, a persistent tensor lives in the system even if
    the ``DistTenor`` object disappears in the trainer process. The persistent tensor has
    the same life span as the DGL servers. DGL does not allow an anonymous tensor to be persistent.

    When a ``DistTensor`` object is created, it may reference to an existing distributed tensor or
    create a new one. A distributed tensor is identified by the name passed to the constructor.
    If the name exists, ``DistTensor`` will reference the existing one.
    In this case, the shape and the data type must match the existing tensor.
    If the name doesn't exist, a new tensor will be created in the kvstore.

    When a distributed tensor is created, its values are initialized to zero. Users
    can define an initialization function to control how the values are initialized.
    The init function has two input arguments: shape and data type and returns a tensor.
    Below shows an example of an init function:

    .. highlight:: python
    .. code-block:: python

        def init_func(shape, dtype):
            return torch.ones(shape=shape, dtype=dtype)

    Parameters
    ----------
    shape : tuple
        The shape of the tensor. The first dimension has to be the number of nodes or
        the number of edges of a distributed graph.
    dtype : dtype
        The dtype of the tensor. The data type has to be the one in the deep learning framework.
    name : string, optional
        The name of the embeddings. The name can uniquely identify embeddings in a system
        so that another ``DistTensor`` object can referent to the distributed tensor.
    init_func : callable, optional
        The function to initialize data in the tensor. If the init function is not provided,
        the values of the embeddings are initialized to zero.
    part_policy : PartitionPolicy, optional
        The partition policy of the rows of the tensor to different machines in the cluster.
        Currently, it only supports node partition policy or edge partition policy.
        The system determines the right partition policy automatically.
    persistent : bool
        Whether the created tensor lives after the ``DistTensor`` object is destroyed.
    is_gdata : bool
        Whether the created tensor is a ndata/edata or not.
    attach : bool
        Whether to attach group ID into name to be globally unique.

    Examples
    --------
    >>> init = lambda shape, dtype: th.ones(shape, dtype=dtype)
    >>> arr = dgl.distributed.DistTensor((g.num_nodes(), 2), th.int32, init_func=init)
    >>> print(arr[0:3])
    tensor([[1, 1],
            [1, 1],
            [1, 1]], dtype=torch.int32)
    >>> arr[0:3] = th.ones((3, 2), dtype=th.int32) * 2
    >>> print(arr[0:3])
    tensor([[2, 2],
            [2, 2],
            [2, 2]], dtype=torch.int32)

    Note
    ----
    The creation of ``DistTensor`` is a synchronized operation. When a trainer process tries to
    create a ``DistTensor`` object, the creation succeeds only when all trainer processes
    do the same.
    """

    def __init__(
        self,
        shape,
        dtype,
        name=None,
        init_func=None,
        part_policy=None,
        persistent=False,
        is_gdata=True,
        attach=True,
    ):
        self.kvstore = get_kvstore()
        assert (
            self.kvstore is not None
        ), "Distributed module is not initialized. Please call dgl.distributed.initialize."
        self._shape = shape
        self._dtype = dtype
        self._attach = attach
        self._is_gdata = is_gdata

        part_policies = self.kvstore.all_possible_part_policy
        # If a user doesn't provide a partition policy, we should find one based on
        # the input shape.
        if part_policy is None:
            for policy_name in part_policies:
                policy = part_policies[policy_name]
                if policy.get_size() == shape[0]:
                    # If multiple partition policies match the input shape, we cannot
                    # decide which is the right one automatically. We should ask users
                    # to provide one.
                    assert part_policy is None, (
                        "Multiple partition policies match the input shape. "
                        + "Please provide a partition policy explicitly."
                    )
                    part_policy = policy
            assert part_policy is not None, (
                "Cannot find a right partition policy. It is either because "
                + "its first dimension does not match the number of nodes or edges "
                + "of a distributed graph or there does not exist a distributed graph."
            )

        self._part_policy = part_policy
        assert (
            part_policy.get_size() == shape[0]
        ), "The partition policy does not match the input shape."

        if init_func is None:
            init_func = _default_init_data
        exist_names = self.kvstore.data_name_list()
        # If a user doesn't provide a name, we generate a name ourselves.
        # We need to generate the name in a deterministic way.
        if name is None:
            assert (
                not persistent
            ), "We cannot generate anonymous persistent distributed tensors"
            global DIST_TENSOR_ID
            # All processes of the same role should create DistTensor synchronously.
            # Thus, all of them should have the same IDs.
            name = "anonymous-" + get_role() + "-" + str(DIST_TENSOR_ID)
            DIST_TENSOR_ID += 1
        assert isinstance(name, str), "name {} is type {}".format(
            name, type(name)
        )
        name = self._attach_group_id(name)
        self._tensor_name = name
        data_name = part_policy.get_data_name(name)
        self._name = str(data_name)
        self._persistent = persistent
        if self._name not in exist_names:
            self._owner = True
            self.kvstore.init_data(
                self._name, shape, dtype, part_policy, init_func, is_gdata
            )
        else:
            self._owner = False
            dtype1, shape1, _ = self.kvstore.get_data_meta(self._name)
            assert (
                dtype == dtype1
            ), "The dtype does not match with the existing tensor"
            assert (
                shape == shape1
            ), "The shape does not match with the existing tensor"

    def __del__(self):
        initialized = (
            os.environ.get("DGL_DIST_MODE", "standalone") == "standalone"
            or is_initialized()
        )
        if not self._persistent and self._owner and initialized:
            self.kvstore.delete_data(self._name)

    def __getitem__(self, idx):
        idx = utils.toindex(idx)
        idx = idx.tousertensor()
        return self.kvstore.pull(name=self._name, id_tensor=idx)

    def __setitem__(self, idx, val):
        idx = utils.toindex(idx)
        idx = idx.tousertensor()
        # TODO(zhengda) how do we want to support broadcast (e.g., G.ndata['h'][idx] = 1).
        self.kvstore.push(name=self._name, id_tensor=idx, data_tensor=val)

    @property
    def kvstore_key(self):
        """Return the key string of this DistTensor in the associated KVStore."""
        return self._name

    @property
    def local_partition(self):
        """Return the local partition of this DistTensor."""
        return self.kvstore.data_store[self._name]

    def __or__(self, other):
        new_dist_tensor = DistTensor(
            self._shape,
            self._dtype,
            part_policy=self._part_policy,
            persistent=self._persistent,
            is_gdata=self._is_gdata,
            attach=self._attach,
        )
        kvstore = self.kvstore
        kvstore.union(self._name, other._name, new_dist_tensor._name)
        return new_dist_tensor

    def __len__(self):
        return self._shape[0]

    @property
    def part_policy(self):
        """Return the partition policy

        Returns
        -------
        PartitionPolicy
            The partition policy of the distributed tensor.
        """
        return self._part_policy

    @property
    def shape(self):
        """Return the shape of the distributed tensor.

        Returns
        -------
        tuple
            The shape of the distributed tensor.
        """
        return self._shape

    @property
    def dtype(self):
        """Return the data type of the distributed tensor.

        Returns
        ------
        dtype
            The data type of the tensor.
        """
        return self._dtype

    @property
    def name(self):
        """Return the name of the distributed tensor

        Returns
        -------
        str
            The name of the tensor.
        """
        return self._detach_group_id(self._name)

    @property
    def tensor_name(self):
        """Return the tensor name

        Returns
        -------
        str
            The name of the tensor.
        """
        return self._detach_group_id(self._tensor_name)

    def count_nonzero(self):
        """Count and return the number of nonzero value

        Returns
        -------
        int
            the number of nonzero value
        """
        return self.kvstore.count_nonzero(name=self._name)

    def _attach_group_id(self, name):
        """Attach group ID if needed

        Returns
        -------
        str
            new name with group ID attached
        """
        if not self._attach:
            return name
        return "{}_{}".format(name, get_group_id())

    def _detach_group_id(self, name):
        """Detach group ID if needed

        Returns
        -------
        str
            original name without group ID
        """
        if not self._attach:
            return name
        suffix = "_{}".format(get_group_id())
        return name[: -len(suffix)]
