"""
Here we encapsulate some common Horovod functions.

Note that you are supposed to be able to import this module even if Horovod is not installed.

The usage of this module / global context is also considered optional at this point.
Horovod is enabled <==> ``use_horovod`` is enabled in the config.

For relevant further config options, see the code of :class:`HorovodContext` below.
Most importantly:

* ``horovod_dataset_distribution``, recommended value ``"random_seed_offset"``, default value ``"shard"``
* ``horovod_reduce_type``, recommended value ``"param"``, default value ``"grad"``
* ``horovod_param_sync_step``, recommended value ``100``, default value ``1``
* ``horovod_param_sync_time_diff``, alternative to ``horovod_param_sync_step``, e.g. ``100.`` (secs), default ``None``

Also see :ref:`multi_gpu`.
Also see :mod:`TFDistributed`.
"""

import os
import socket
import typing
from returnn.config import Config


class HorovodContext:
    """
    This setups some helper functions.
    """

    def __init__(self, config):
        """
        :param Config config:
        """
        # noinspection PyUnresolvedReferences,PyPackageRequirements
        import horovod

        print("Horovod:", horovod.__version__, horovod.__file__)
        # noinspection PyUnresolvedReferences,PyPackageRequirements
        import horovod.tensorflow as hvd

        hvd.init()
        print(
            "Horovod initialized. Hostname %s, pid %i, rank %i / size %i, local rank %i / local size %i."
            % (socket.gethostname(), os.getpid(), hvd.rank(), hvd.size(), hvd.local_rank(), hvd.local_size())
        )
        self._config = config
        self._hvd_mod = hvd
        self._local_rank = hvd.local_rank()
        self._local_size = hvd.local_size()
        self._rank = hvd.rank()
        self._size = hvd.size()
        # Note: Rather than reading the config here, we will instead read it in the functions below on-the-fly.
        # This allows e.g. the pretraining or custom get_network to temporarily overwrite the behavior for some epoch.

    def should_sync_every_step(self):
        """
        :return: whether we should sync every step.
          This is both for the signal for more data, and also loss/error/score reduction.
        :rtype: bool
        """
        if self.is_dataset_distribution_random_seed_offset() and self.is_reduce_type_param():
            # In this setting, we don't need a sync every step.
            # Thus avoid it to have faster more async training.
            return False
        # By default, we should sync every step.
        return True

    def get_reduce_type(self):
        """
        :rtype: str
        """
        reduce_type = self._config.value("horovod_reduce_type", "grad")
        assert reduce_type in {"grad", "param"}
        return reduce_type

    def is_reduce_type_grad(self):
        """
        :rtype: bool
        """
        return self.get_reduce_type() == "grad"

    def is_reduce_type_param(self):
        """
        :rtype: bool
        """
        return self.get_reduce_type() == "param"

    def get_param_sync_time_diff(self):
        """
        :rtype: float|None
        """
        assert self.is_reduce_type_param()
        return self._config.float("horovod_param_sync_time_diff", None)

    def get_param_sync_step(self):
        """
        :rtype: int
        """
        assert self.is_reduce_type_param()
        return self._config.int("horovod_param_sync_step", 1)

    def get_dataset_distribution_type(self):
        """
        :rtype: str
        """
        dataset_distribution = self._config.value("horovod_dataset_distribution", "shard")
        assert dataset_distribution in {"shard", "random_seed_offset"}
        return dataset_distribution

    def is_dataset_distribution_shard(self):
        """
        :rtype: bool
        """
        return self.get_dataset_distribution_type() == "shard"

    def get_dataset_shard_batch_slice(self):
        """
        :rtype: slice
        """
        assert self.is_dataset_distribution_shard()
        return slice(self.rank(), None, self.size())

    def is_dataset_distribution_random_seed_offset(self):
        """
        :rtype: bool
        """
        return self.get_dataset_distribution_type() == "random_seed_offset"

    def rank(self):
        """
        :rtype: int
        """
        return self._rank

    def size(self):
        """
        :rtype: int
        """
        return self._size

    def local_rank(self):
        """
        :rtype: int
        """
        return self._local_rank

    def local_size(self):
        """
        :rtype: int
        """
        return self._local_size


_is_set_up = False
_ctx = None  # type: typing.Optional[HorovodContext]


def get_ctx(config=None):
    """
    :param Config|None config:
    :returns: the global context if Horovod is enabled, or None otherwise.
      If we did not setup the context yet, it will automatically create it.
    :rtype: HorovodContext|None
    """
    global _is_set_up, _ctx
    if _is_set_up:
        return _ctx
    if not config:
        from returnn.config import get_global_config

        config = get_global_config(raise_exception=False)
        if not config:
            return None
    _is_set_up = True
    if not config.is_true("use_horovod"):
        return None
    _ctx = HorovodContext(config=config)
    return _ctx
