import warnings
from collections import OrderedDict

import torch
from torch.nn import DataParallel as DataParallel_
from torch.nn.parallel import parallel_apply
from torch.nn.parallel.replicate import _broadcast_coalesced_reshape
from torch.nn.parallel.scatter_gather import scatter_kwargs

from torchmetal.modules.module import MetaModule


class DataParallel(DataParallel_, MetaModule):
    __doc__ = DataParallel_.__doc__

    def scatter(self, inputs, kwargs, device_ids):
        if not isinstance(self.module, MetaModule):
            return super(DataParallel, self).scatter(inputs, kwargs, device_ids)

        params = kwargs.pop("params", None)
        inputs_, kwargs_ = scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
        # Add params argument unchanged back in kwargs
        replicas = self._replicate_params(
            params, inputs_, device_ids, detach=not torch.is_grad_enabled()
        )
        kwargs_ = tuple(
            dict(params=replica, **kwarg) for (kwarg, replica) in zip(kwargs_, replicas)
        )
        return inputs_, kwargs_

    def _replicate_params(self, params, inputs, device_ids, detach=False):
        if params is None:
            module_params = OrderedDict(self.module.named_parameters())
        else:
            # Temporarily disable the warning if no parameter with key prefix
            # `module` was found. In that case, the original params dictionary
            # is used.
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                module_params = self.get_subdict(params, key="module")
            if module_params is None:
                module_params = params

        replicas = _broadcast_coalesced_reshape(
            list(module_params.values()), device_ids[: len(inputs)], detach
        )
        replicas = tuple(
            OrderedDict(zip(module_params.keys(), replica)) for replica in replicas
        )
        return replicas
