# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.

"""
Train a network on multiple GPUs.
"""

import torch
import ctypes

from fairseq import optim, utils
from fairseq.meters import AverageMeter
from fairseq.optim import lr_scheduler
from fairseq.trainer import Trainer

def fused_norm(input):
    return input.norm(dtype=torch.float32,p=2).item()

class DynamicLossScaler:

    def __init__(self, init_scale=2.**15, scale_factor=2., scale_window=2000):
        self.loss_scale = init_scale
        self.scale_factor = scale_factor
        self.scale_window = scale_window
        self._iter = 0
        self._last_overflow_iter = -1

    def update_scale(self, overflow):
        if overflow:
            self.loss_scale /= self.scale_factor
            self._last_overflow_iter = self._iter
        elif (self._iter - self._last_overflow_iter) % self.scale_window == 0:
            self.loss_scale *= self.scale_factor
        self._iter += 1

    @staticmethod
    def has_overflow(grad_norm):
        # detect inf and nan
        if grad_norm == float('inf') or grad_norm != grad_norm:
            return True
        return False


class FP16Trainer(Trainer):
    """Modified trainer for FP16.

    We maintain two copies of the model's parameters, both in FP16 and FP32.
    We do forward/backward with FP16 and compute the loss + optimize with FP32.
    """

    def __init__(self, args, task, model, criterion, allreduce_communicators=None):
        super().__init__(args, task, model, criterion, allreduce_communicators)

        # convert model to FP16 (but keep criterion FP32)
        self.model.half()

        # broadcast initial weights from rank=0
        # this broadcast isn't required in DistributedFP16Trainer because the
        # broadcast is done by DistributedFusedAdam
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            for p in self.model.parameters():
                torch.distributed.broadcast(p, 0)

        # dynamically scale loss to reduce overflow
        self.scaler = DynamicLossScaler(init_scale=2.**7)
        self.meters['loss_scale'] = AverageMeter()

        self.grad_denom = 1.0

        if self.args.enable_parallel_backward_allred_opt:
            import numpy as np

            self._flat_grads_parallel = torch.tensor([], dtype=torch.float16).cuda()
            self._grads_info = []
            grads_size = 0
            p_offset = 0
            for p_i, p in enumerate([p for p in self.model.parameters() if p.requires_grad]):
                p_grads_size = np.prod(list(p.size()))
                grads_size += p_grads_size
                # register hooks
                def wrapper(param, param_i, param_grads_size, param_offset):
                    def allreduce_hook(grad):
                        self._do_allreduce(param_i, param_grads_size, param_offset, grad)

                    if param.requires_grad:
                        param.register_hook(allreduce_hook)
                # print(p_i, p.size(), p_grads_size, p_offset)
                self._grads_info.append({"param_grads_size":p_grads_size, "param_offset":p_offset})
                wrapper(p, p_i, p_grads_size, p_offset)
                p_offset += p_grads_size
            self._flat_grads_parallel.resize_(grads_size)
            # print(grads_size, len(self._flat_grads_parallel), self._flat_grads_parallel.dtype, self._flat_grads_parallel.get_device())

            self._allreduce_flush_min_threshold = self.args.parallel_backward_allred_opt_threshold
            print("| parallel all-reduce ENABLED. all-reduce threshold: " + str(self._allreduce_flush_min_threshold))
            self._grads_generated = [False]*len(self._grads_info)
            self._allreduce_processed_idx = len(self._grads_info)-1

            self._num_allreduce_sent = 0
            print("| # of parallel all-reduce cuda streams: " + str(self.args.parallel_backward_allred_cuda_nstreams))
            if allreduce_communicators:
                self._allreduce_groups = allreduce_communicators[0]
                self._allreduce_streams = allreduce_communicators[1]
            else:
                raise RuntimeError('Moved communicator init before RUN_START (invalid code path)')
                self._allreduce_groups = [torch.distributed.new_group() for _ in range(self.args.parallel_backward_allred_cuda_nstreams)]
                self._allreduce_streams = [torch.cuda.Stream() for _ in range(self.args.parallel_backward_allred_cuda_nstreams)]

            if self.args.enable_parallel_backward_allred_opt_correctness_check:
                self._num_grads_generated = 0
                self._all_grads_generated = False
                self._allreduce_schedule = []

    def _get_flush_bucket(self):
        # print([1 if x else 0 for x in self._grads_generated])
        flush_bucket = []

        size = 0
        allreduce_processed_idx_list = []
        allreduce_processed_end_idx = self._allreduce_processed_idx
        remaining_grads_for_allreduce = self._grads_generated[allreduce_processed_end_idx-len(self._grads_generated)::-1]
        # print([1 if x else 0 for x in remaining_grads_for_allreduce])
        for s in remaining_grads_for_allreduce:
            # print(s,allreduce_processed_end_idx,size)
            if s:
                allreduce_processed_idx_list.append(allreduce_processed_end_idx)
                size += self._grads_info[allreduce_processed_end_idx]["param_grads_size"]
                allreduce_processed_end_idx -= 1
            else:
                break

        # print(size, allreduce_processed_idx_list)
        ignore_threshold = all(self._grads_generated)

        if size >= self._allreduce_flush_min_threshold or ignore_threshold:
            # for i in allreduce_processed_idx_list:
            #     print(i, self._grads_info[i]["param_grads_size"], self._grads_info[i]["param_offset"],size)
            if allreduce_processed_idx_list:
                start = self._grads_info[(allreduce_processed_idx_list[-1])]["param_offset"]
                end = start + size
                # print("->", start, end)
                flush_bucket = [start, end]

            self._allreduce_processed_idx = allreduce_processed_end_idx
            if self._allreduce_processed_idx < 0:
                # reset
                self._grads_generated = [False]*len(self._grads_info)
                self._allreduce_processed_idx = len(self._grads_info)-1

        return flush_bucket

    def _do_allreduce(self, param_i, param_grads_size, param_offset, grad):
        if self._last_step == False:

            # # ----------------------
            # # debugging: do all-reduce in the same stream
            # print(self._last_step, self._grads_total, len(self._backward_grads_schedule), param_i, param_offset, param_grads_size, grad.size(), grad.numel(), grad.dtype)
            # self._flat_grads_parallel[param_offset:param_offset+param_grads_size].copy_(grad.view(-1))
            # self._flat_grads_parallel[param_offset:param_offset+param_grads_size].div_(self.args.distributed_world_size)
            # torch.distributed.all_reduce(self._flat_grads_parallel[param_offset:param_offset+param_grads_size])
            # # ----------------------

            # # ----------------------
            # # option #1: send per-layer gradients
            # torch.div(grad.view(-1), self.args.distributed_world_size, out=self._flat_grads_parallel[param_offset:param_offset+param_grads_size])
            # orig_stream = torch.cuda.current_stream()
            # self._reduction_stream.wait_stream(orig_stream)
            # with torch.cuda.stream(self._reduction_stream):
            #     torch.distributed.all_reduce(self._flat_grads_parallel[param_offset:param_offset+param_grads_size])
            # # ----------------------

            # ----------------------
            # option #2: bucket all-reduce based on threshold
            self._flat_grads_parallel.record_stream(torch.cuda.current_stream())
            torch.div(grad.view(-1), self.args.distributed_world_size, out=self._flat_grads_parallel[param_offset:param_offset+param_grads_size])
            self._grads_generated[param_i]=True
            flush_bucket = self._get_flush_bucket()
            if flush_bucket:
                start = flush_bucket[0]
                end = flush_bucket[1]
                # print("->", start, end)

                if self.args.enable_parallel_backward_allred_opt_correctness_check and not self._all_grads_generated:
                    self._allreduce_schedule.append(flush_bucket)

                # orig_stream = torch.cuda.current_stream()
                # self._reduction_stream.wait_stream(orig_stream)
                # with torch.cuda.stream(self._reduction_stream):
                #     torch.distributed.all_reduce(self._flat_grads_parallel[start:end])
                orig_stream = torch.cuda.current_stream()
                allreduce_group = self._allreduce_groups[self._num_allreduce_sent%len(self._allreduce_groups)]
                allreduce_stream = self._allreduce_streams[self._num_allreduce_sent%len(self._allreduce_streams)]
                allreduce_stream.wait_stream(orig_stream)
                with torch.cuda.stream(allreduce_stream):
                    self._flat_grads_parallel.record_stream(torch.cuda.current_stream())
                    torch.distributed.all_reduce(self._flat_grads_parallel[start:end], group=allreduce_group)
                self._num_allreduce_sent += 1


            if self.args.enable_parallel_backward_allred_opt_correctness_check:
                self._num_grads_generated += 1
                if self._num_grads_generated == len(self._grads_info):
                    self._all_grads_generated = True
            # ----------------------


    def _build_optimizer(self):
        # create FP32 copy of parameters and grads
        params = [p for p in self.model.parameters() if p.requires_grad]
        total_param_size = sum(p.data.numel() for p in params)
        self.fp32_params = params[0].new(0).float().new(total_param_size)
        offset = 0
        for p in params:
            numel = p.data.numel()
            self.fp32_params[offset:offset+numel].copy_(p.data.view(-1))
            offset += numel
        self.fp32_params = torch.nn.Parameter(self.fp32_params)
        #self.fp32_params.grad = self.fp32_params.data.new(total_param_size)

        # create optimizer using the copied FP32 params
        self._optimizer = optim.build_optimizer(self.args, [self.fp32_params])
        self.lr_scheduler = lr_scheduler.build_lr_scheduler(self.args, self.optimizer)

    def save_checkpoint(self, filename, extra_state):
        """Save all training state in a checkpoint file."""
        extra_state['loss_scale'] = self.scaler.loss_scale
        super().save_checkpoint(filename, extra_state)

    def load_checkpoint(self, filename):
        """Load all training state from a checkpoint file."""
        extra_state = super().load_checkpoint(filename)
        if extra_state is not None and 'loss_scale' in extra_state:
            self.scaler.loss_scale = extra_state['loss_scale']
        return extra_state

    def zero_grad(self):
        # zero both the FP16 and FP32 grads
#        self.model.zero_grad()      # FP16
#        self.optimizer.zero_grad()  # FP32
#        r"""Clears the gradients of all optimized :class:`torch.Tensor` s."""
        for p in self.model.parameters():
            p.grad = None





    def _backward(self, loss):
        self.meters['loss_scale'].reset()
        self.meters['loss_scale'].update(self.scaler.loss_scale)
        if loss is not None:
            # dynamically rescale loss to stay in FP16 range
            loss = loss * self.scaler.loss_scale
        return super()._backward(loss)

    def _all_reduce_and_rescale(self, grad_denom, has_grad = True):
        # undo effect of dynamic loss scaling on gradients
        self.grad_denom = grad_denom * self.scaler.loss_scale

        if self.args.distributed_world_size > 1:
            self.grad_denom /= self.args.distributed_world_size

            if not self.args.enable_parallel_backward_allred_opt or self._last_step:
                # flatten grads into a single buffer
                self._flat_grads = self._get_flat_grads(out=None, has_grad = has_grad)

                # scale gradients to avoid overflow in all-reduce
                self._flat_grads.div_(self.args.distributed_world_size)

                # all-reduce flat grads
                torch.distributed.all_reduce(self._flat_grads)
            else:
                # torch.cuda.current_stream().wait_stream(self._reduction_stream)
                for allreduce_stream in self._allreduce_streams:
                    torch.cuda.current_stream().wait_stream(allreduce_stream)
                self._flat_grads_parallel.record_stream(torch.cuda.current_stream())
                self._flat_grads = self._flat_grads_parallel

                if self.args.enable_parallel_backward_allred_opt_correctness_check:
                    # # ----------------------
                    # # option #1: send per-layer gradients
                    # grads = self._get_grads()
                    # offset = 0
                    # for g in grads:
                    #     numel = g.numel()
                    #     out = grads[0].new(numel).zero_()
                    #     out.copy_(g.view(-1))
                    #     out.div_(self.args.distributed_world_size)
                    #     torch.distributed.all_reduce(out)
                    #     is_parallel_grads_finite = torch.all(torch.isfinite(self._flat_grads_parallel[offset:offset+numel]))
                    #     is_out_finite = torch.all(torch.isfinite(out))
                    #     assert(is_out_finite == is_parallel_grads_finite)
                    #     if not is_out_finite:
                    #         print("| OVERLAP-CHECK: check inf/nan detected. this batch should be skipped")
                    #     else:
                    #         if not torch.all(torch.eq(out, self._flat_grads_parallel[offset:offset+numel])):
                    #             print(out[0:10], self._flat_grads_parallel[offset:offset+10])
                    #             # for i,_ in enumerate(out):
                    #             #     if out[i] != self._flat_grads_parallel[i]:
                    #             #         print(i,out[i],self._flat_grads_parallel[i])
                    #             raise RuntimeError('w-gradients received in parallel vs. end differ')
                    #     offset += numel
                    # # ----------------------

                    # ----------------------
                    # option #2: bucket all-reduce based on threshold
                    # print(self._allreduce_schedule)
                    out = self._get_flat_grads()
                    out.div_(self.args.distributed_world_size)
                    grads_size = 0
                    for s in self._allreduce_schedule:
                        start = s[0]
                        end = s[1]
                        assert(end > start)
                        grads_size += (end - start)
                        torch.distributed.all_reduce(out[start:end])
                        is_parallel_grads_finite = torch.all(torch.isfinite(self._flat_grads_parallel[start:end]))
                        is_out_finite = torch.all(torch.isfinite(out[start:end]))
                        assert(is_out_finite == is_parallel_grads_finite)
                        if not is_out_finite:
                            print("| OVERLAP-CHECK: check inf/nan detected. this batch should be skipped")
                        else:
                            if not torch.all(torch.eq(out[start:end], self._flat_grads_parallel[start:end])):
                                print(start, end, out[start:end], self._flat_grads_parallel[start:end])
                                raise RuntimeError('w-gradients received in parallel vs. end differ')
                    assert(grads_size == len(self._flat_grads_parallel))
                    # ----------------------
        else:
            # flatten grads into a single buffer
            self._flat_grads = self._get_flat_grads(out=None, has_grad = has_grad)

        grad_norm = fused_norm(self._flat_grads)

        # detect overflow and adjust loss scale
        overflow = DynamicLossScaler.has_overflow(grad_norm)
        self.scaler.update_scale(overflow)
        if overflow:
            if self.scaler.loss_scale <= self.args.min_loss_scale:
                raise Exception((
                    'Minimum loss scale reached ({}). Your loss is probably exploding. '
                    'Try lowering the learning rate, using gradient clipping or '
                    'increasing the batch size.'
                ).format(self.args.min_loss_scale))
            raise OverflowError('setting loss scale to: ' + str(self.scaler.loss_scale))

        return grad_norm

    def _opt(self):
        # take an optimization step using the FP32 params and grads
        #super()._opt()
        new_params = self._flat_grads.new_empty(self._flat_grads.size())
        self.optimizer.optimizer.step(closure=None, grads=[self._flat_grads], output_params=[new_params], scale=self.grad_denom)
        self.zero_grad()
        self._num_updates += 1

        # update learning rate
        self.lr_scheduler.step_update(self._num_updates)

        # copy FP32 params back into FP16 model
        offset = 0
        with torch.no_grad():
            for p in self.model.parameters():
                if not p.requires_grad:
                    continue
                numel = p.data.numel()
                p.set_(new_params[offset:offset+numel].view_as(p.data))
                offset += numel


class DistributedFP16Trainer(Trainer):
    """Modified trainer for FP16.

    We maintain two copies of the model's parameters, both in FP16 and FP32.
    We do forward/backward with FP16 and compute the loss + optimize with FP32.
    """

    def __init__(self, args, task, model, criterion, allreduce_communicators=None):
        super().__init__(args, task, model, criterion, allreduce_communicators)

        # convert model to FP16 (but keep criterion FP32)
        self.model.half()

        # dynamically scale loss to reduce overflow
        self.scaler = DynamicLossScaler(init_scale=2.**7)
        self.meters['loss_scale'] = AverageMeter()
        # FIXME: Add more meters

        self.grad_denom = 1.0

        assert (not self.args.enable_parallel_backward_allred_opt), "--distributed-weight-update cannot be combined with --enable-parallel-backward-allred-opt"

    def save_checkpoint(self, filename, extra_state):
        """Save all training state in a checkpoint file."""
        # To-Do: gather optimizer buffer chunks before saving state
        extra_state['loss_scale'] = self.scaler.loss_scale
        super().save_checkpoint(filename, extra_state)

    def load_checkpoint(self, filename):
        """Load all training state from a checkpoint file."""
        # To-Do: scatter optimizer buffer chunks after restoring state
        extra_state = super().load_checkpoint(filename)
        if extra_state is not None and 'loss_scale' in extra_state:
            self.scaler.loss_scale = extra_state['loss_scale']
        return extra_state

    #def zero_grad(self):
    #    for p in self.model.parameters():
    #        p.grad = None

    def _backward(self, loss):
        self.meters['loss_scale'].reset()
        self.meters['loss_scale'].update(self.scaler.loss_scale)

        if loss is not None:
            # dynamically rescale loss to stay in FP16 range
            loss = loss * self.scaler.loss_scale
        rval = super()._backward(loss)
        self.optimizer.optimizer.complete_reductions()
        return rval

    def __process_overflow(self, overflow):
        self.scaler.update_scale(overflow)
        if overflow:
            if self.scaler.loss_scale <= self.args.min_loss_scale:
                raise Exception((
                    'Minimum loss scale reached ({}). Your loss is probably exploding. '
                    'Try lowering the learning rate, using gradient clipping or '
                    'increasing the batch size.'
                ).format(self.args.min_loss_scale))
            raise OverflowError('setting loss scale to: ' + str(self.scaler.loss_scale))

    def _all_reduce_and_rescale(self, grad_denom, has_grad = True):
        grad_norm = self.optimizer.optimizer.L2_grad_norm
        if grad_norm is not None:
            overflow = self.scaler.has_overflow(grad_norm)
            self.__process_overflow(overflow)
            return grad_norm
        else:
            return None

    def _opt(self):
        self.optimizer.optimizer.step(skip_overflow_check=self.args.dwu_compute_L2_grad_norm)
        self.zero_grad()

        self.__process_overflow(False if self.args.dwu_compute_L2_grad_norm or not self.optimizer.optimizer.has_overflow else True)

        self._num_updates += 1

        # update learning rate
        self.lr_scheduler.step_update(self._num_updates)

