import functools
import time
import warnings
from abc import ABC
from typing import Optional

import numpy as np
import torch
import torch.cuda
import torch.distributed
from torch.autograd import Variable
from torch.nn.parallel import DistributedDataParallel

from kubega import checkpoint, env, utils
from kubega.torch._metrics import update_sync_time, update_grad_params, update_progress
from kubega.torch.data import current_dataloader
from kubega.torch.gradient_noise_scale import GradientNoiseScale, AdamGradientNoiseScale
from kubega.torch.scaling_rules import AdaScale, AdamScale, ScalingRuleBase


class _AdaptiveDataParallelState(checkpoint.State):
    def __init__(self, model, optimizer, lr_scheduler, mp_scaler, name="kubega-data-parallel"):
        super().__init__(name)
        self.model = model
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.mp_scaler = mp_scaler
        # TODO: Gain/Goodput should be tracked in the metrics module instead, not here.
        self.gain = 1.0
        # lr_factor summary
        self.lr_factor = 1.0

    def save(self, fileobj):
        state_dicts = [self.model.state_dict(), self.optimizer.state_dict()]

        if self.lr_scheduler is not None:
            state_dicts.append(self.lr_scheduler.state_dict())
        else:
            state_dicts.append(None)
        if self.mp_scaler is not None:
            state_dicts.append(self.mp_scaler.state_dict())
        else:
            state_dicts.append(None)

        torch.save((state_dicts, self.gain, self.lr_factor), fileobj)

    def load(self, fileobj):
        state_dicts, self.gain, self.lr_factor = torch.load(fileobj)
        self.model.load_state_dict(state_dicts[0])
        self.optimizer.load_state_dict(state_dicts[1])
        if state_dicts[2] is not None:
            self.lr_scheduler.load_state_dict(state_dicts[2])
        if state_dicts[3] is not None:
            self.mp_scaler.load_state_dict(state_dicts[3])


class AdaptiveDataParallel(DistributedDataParallel, ABC):
    """
    This class extends PyTorch DistributedDataParallel with support for
    adaptive batch sizes and checkpoint-restart elasticity. It automatically
    saves the given model, optimizer, and (optionally) LR scheduler whenever a
    checkpoint is triggered, and restores their states after restart. The
    optimizer is automatically patched with the chosen scaling rule.

    Arguments:
        model (torch.nn.Module): Model to be distributed.
        optimizer (torch.optim.Optimizer): Optimizer used to update the given
            model's parameters, will be patched using subclass of
        :class:`kubega.torch.scaling_rules.ScalingRuleBase`.
        scaling_rule (ScalingRuleBase): Scaling rule used to
            patch the given optimizer, default to AdaScale.
        lr_scheduler (torch.optim.lr_scheduler._LRScheduler): LR scheduler used
            to anneal the learning rate for the given optimizer.
        name (string): Unique name for each instance of this class, needed only
            if multiple instances exist.
    """

    def __init__(self, model, optimizer, lr_scheduler=None, mp_scaler=None,
                 scaling_rule: Optional[ScalingRuleBase] = None,
                 name="kubega-data-parallel", **kwargs):
        super().__init__(model, **kwargs)
        self._key = id(self)

        # Register backward hooks on model parameters. Depends on these hooks
        # being invoked before gradients are averaged. This is technically an
        # internal behavior of DistributedDataParallel, but seems to be abused
        # pretty widely so there should be little chance of it changing.
        # https://discuss.pytorch.org/t/59291
        for param in model.parameters():
            param.register_hook(functools.partial(self._backward_hook, param))

        # Setup for the scaling_rule, must be after registering backward hooks
        # because some of them need to register their own backward hooks.
        if not scaling_rule and (isinstance(optimizer, torch.optim.Adam) or
                                 isinstance(optimizer, torch.optim.AdamW)):
            self.scaling_rule = AdamScale()
        else:
            self.scaling_rule = scaling_rule or AdaScale()

        if isinstance(scaling_rule, AdamScale):
            self.gns = AdamGradientNoiseScale(self, optimizer, mp_scaler=mp_scaler)
        else:
            self.gns = GradientNoiseScale(self, optimizer, mp_scaler=mp_scaler)

        self.scaling_rule.initialize(self, optimizer, patch_optimizer=True)

        self._state = _AdaptiveDataParallelState(model, optimizer, lr_scheduler, mp_scaler, name)
        checkpoint.load_state(self._state)

        self._sync_start = None

    def forward(self, *args, **kwargs):
        # Do not do gradient synchronization during gradient accumulation.
        dataloader = current_dataloader()
        if dataloader is not None and dataloader.training:
            self.require_backward_grad_sync = dataloader.is_optim_step()
            accum_scale = (dataloader.current_local_bsz * env.num_replicas() / dataloader.batch_size)
            self.gns.set_accum_scale(accum_scale)
        return super().forward(*args, **kwargs)

    @utils.print_exc
    def _backward_hook(self, param, grad):
        # This method should be invoked once for each parameter during the
        # backward pass, before gradients are synchronized between replicas.
        if grad.device.type.startswith("cuda"):
            self._sync_start = torch.cuda.Event(enable_timing=True)
            self._sync_start.record()
        else:
            self._sync_start = time.time()
        self._final_callback_queued = False
        Variable._execution_engine.queue_callback(self._queue_callback)

    @utils.print_exc
    def _queue_callback(self):
        # This method should be invoked after the entire backward pass. We want
        # to make sure self._final_callback is invoked once, only after all
        # gradients have been synchronized between each replica. However, the
        # synchronization code in DistributedDataParallel is also done in a
        # callback, which might not yet be executed. Therefore, we enqueue
        # self._final_callback from this method, which should ensure it is
        # invoked after the gradient synchronization callback.
        if self._final_callback_queued:
            return
        self._final_callback_queued = True
        Variable._execution_engine.queue_callback(self._final_callback)

    @utils.print_exc
    def _final_callback(self):
        # This method should be invoked once for each backward pass, after
        # gradients have been synchronized between each replica.
        self._final_callback_queued = False
        # self._sync_start should mark the last time any local gradient
        # from this module was produced. We assume the duration until now was
        # primarily spent waiting for gradient synchronization.
        # TODO: Depends on the internal behavior of DistributedDataParallel,
        #       which might break with future versions of PyTorch. Any better
        #       and well-supported way to measure the synchronization time?
        if isinstance(self._sync_start, torch.cuda.Event):
            sync_end = torch.cuda.Event(enable_timing=True)
            sync_end.record()
            sync_end.synchronize()
            update_sync_time(self._sync_start.elapsed_time(sync_end) / 1e3)
        else:
            update_sync_time(time.time() - self._sync_start)

        dataloader = current_dataloader()
        if dataloader is None:
            # Don't allow backpropagation outside a dataloader loop, because
            # the batch size would be unknown.
            raise RuntimeError("backpropagation outside AdaptiveDataLoader")
        dataloader.train()

        scale = dataloader.current_batch_size / dataloader.batch_size
        self._state.gain = self.gns.gain(scale)
        self._state.lr_factor = \
            np.average(self.scaling_rule.scale_lr(scale))
        update_progress(self.gns.get_progress())
        if dataloader.max_batch_size and \
                dataloader.max_batch_size > dataloader.batch_size:
            update_grad_params(self._key, self.gns.sqr_avg(),
                               self.gns.var_avg())
        self._sync_start = None

    def zero_grad(self, *args, **kwargs):
        warnings.warn("zero_grad has no effect with AdaptiveDataParallel")

    @property
    def gain(self):  # TODO: should be tracked in the metrics module instead.
        """
        Current estimate of the AdaScale gain (r_t) value.
        """
        return self._state.gain

    def to_tensorboard(self, writer, global_step, tag_prefix=""):
        """
        Output some useful metrics to TensorBoard.

        Arguments:
            writer (torch.utils.tensorboard.SummaryWriter): ``SummaryWriter``
                object to output metrics to.
            global_step (int): Global step value to record.
            tag_prefix (str): Prefix added to each metric's tag.
        """
        if tag_prefix and not tag_prefix.endswith("/"):
            tag_prefix += "/"
        writer.add_scalar(tag_prefix + "Gradient_Norm_Sqr",
                          self.gns.sqr_avg(), global_step)
        writer.add_scalar(tag_prefix + "Gradient_Variance",
                          self.gns.var_avg(), global_step)
        writer.add_scalar(tag_prefix + "Gain",
                          self._state.gain, global_step)
        writer.add_scalar(tag_prefix + "Learning_Rate_Factor",
                          self._state.lr_factor, global_step)
        writer.add_scalar(tag_prefix + "Accum_Scale",
                          self.gns.accum_scale, global_step)
        if self.gns.accum_count > 0:
            writer.add_scalar(tag_prefix + "Accum_Count",
                              self.gns.accum_count, global_step)
        writer.add_scalar(tag_prefix + "Progress",
                          self.gns.get_progress(), global_step)
