import os
import time
from collections import namedtuple
import torch

try:
    from apex import amp
    from apex.amp import _amp_state
except ImportError:
    pass
    #raise ImportError("Please install apex from https://www.github.com/nvidia/apex if you want to use fp16.")

# Parameter to pass to batch_end_callback
BatchEndParam = namedtuple('BatchEndParams',
                           ['epoch',
                            'nbatch',
                            'rank',
                            'add_step',
                            'data_in_time',
                            'data_transfer_time',
                            'forward_time',
                            'backward_time',
                            'optimizer_time',
                            'metric_time',
                            'eval_metric',
                            'locals'])


def _multiple_callbacks(callbacks, *args, **kwargs):
    """Sends args and kwargs to any configured callbacks.
    This handles the cases where the 'callbacks' variable
    is ``None``, a single function, or a list.
    """
    if isinstance(callbacks, list):
        for cb in callbacks:
            cb(*args, **kwargs)
        return
    if callbacks:
        callbacks(*args, **kwargs)


def to_cuda(batch):
    batch = list(batch)

    for i in range(len(batch)):
        if isinstance(batch[i], torch.Tensor):
            batch[i] = batch[i].cuda(non_blocking=True)
        elif isinstance(batch[i], list):
            for j, o in enumerate(batch[i]):
                if isinstance(batch[i], torch.Tensor):
                    batch[i][j] = o.cuda(non_blocking=True)

    return batch


def train(net,
          optimizer,
          lr_scheduler,
          train_loader,
          train_sampler,
          metrics,
          begin_epoch,
          end_epoch,
          logger,
          rank=None,
          batch_end_callbacks=None,
          epoch_end_callbacks=None,
          writer=None,
          validation_monitor=None,
          fp16=False,
          clip_grad_norm=-1,
          gradient_accumulate_steps=1):

    assert isinstance(gradient_accumulate_steps, int) and gradient_accumulate_steps >= 1
    # torch.autograd.set_detect_anomaly(True)
    

    for epoch in range(begin_epoch, end_epoch):
        print('PROGRESS: %.2f%%' % (100.0 * epoch / end_epoch))

        # set epoch as random seed of sampler while distributed training
        if train_sampler is not None and hasattr(train_sampler, 'set_epoch'):
            train_sampler.set_epoch(epoch)

        # reset metrics
        metrics.reset()

        # set net to train mode
        net.train()

        # clear the paramter gradients
        # optimizer.zero_grad()

        # init end time
        end_time = time.time()

        if isinstance(lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
            name, value = validation_monitor.metrics.get()
            val = value[name.index(validation_monitor.host_metric_name)]
            lr_scheduler.step(val, epoch)

        # training
        for nbatch, batch in enumerate(train_loader):
            global_steps = len(train_loader) * epoch + nbatch
            os.environ['global_steps'] = str(global_steps)

            # record time
            data_in_time = time.time() - end_time

            # transfer data to GPU
            data_transfer_time = time.time()
            batch = to_cuda(batch)
            data_transfer_time = time.time() - data_transfer_time

            # forward
            forward_time = time.time()
            outputs, loss = net(*batch)
            loss = loss.mean()
            if gradient_accumulate_steps > 1:
                loss = loss / gradient_accumulate_steps
            forward_time = time.time() - forward_time

            # backward
            backward_time = time.time()
            if fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            backward_time = time.time() - backward_time

            optimizer_time = time.time()
            if (global_steps + 1) % gradient_accumulate_steps == 0:
                # step LR scheduler
                if lr_scheduler is not None and not isinstance(lr_scheduler,
                                                               torch.optim.lr_scheduler.ReduceLROnPlateau):
                    lr_scheduler.step()

                # clip gradient
                if clip_grad_norm > 0:
                    if fp16:
                        total_norm = torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
                                                                    clip_grad_norm)
                    else:
                        total_norm = torch.nn.utils.clip_grad_norm_(net.parameters(),
                                                                    clip_grad_norm)
                    if writer is not None:
                        writer.add_scalar(tag='grad-para/Total-Norm',
                                          scalar_value=float(total_norm),
                                          global_step=global_steps)

                optimizer.step()
                # clear the parameter gradients
                optimizer.zero_grad()
            optimizer_time = time.time() - optimizer_time

            # update metric
            metric_time = time.time()
            metrics.update(outputs)
            if writer is not None:
                with torch.no_grad():
                    for group_i, param_group in enumerate(optimizer.param_groups):
                        writer.add_scalar(tag='Initial-LR/Group_{}'.format(group_i),
                                          scalar_value=param_group['initial_lr'],
                                          global_step=global_steps)
                        writer.add_scalar(tag='LR/Group_{}'.format(group_i),
                                          scalar_value=param_group['lr'],
                                          global_step=global_steps)
                    writer.add_scalar(tag='Train-Loss',
                                      scalar_value=float(loss.item()),
                                      global_step=global_steps)
                    name, value = metrics.get()
                    for n, v in zip(name, value):
                        writer.add_scalar(tag='Train-' + n,
                                          scalar_value=v,
                                          global_step=global_steps)

            metric_time = time.time() - metric_time

            # execute batch_end_callbacks
            if batch_end_callbacks is not None:
                batch_end_params = BatchEndParam(epoch=epoch, nbatch=nbatch, add_step=True, rank=rank,
                                                 data_in_time=data_in_time, data_transfer_time=data_transfer_time,
                                                 forward_time=forward_time, backward_time=backward_time,
                                                 optimizer_time=optimizer_time, metric_time=metric_time,
                                                 eval_metric=metrics, locals=locals())
                _multiple_callbacks(batch_end_callbacks, batch_end_params)

            # update end time
            end_time = time.time()

        # excute epoch_end_callbacks
        if validation_monitor is not None:
            validation_monitor(epoch, net, optimizer, writer)
        if epoch_end_callbacks is not None:
            _multiple_callbacks(epoch_end_callbacks, epoch, net, optimizer, writer, validation_monitor=validation_monitor)


