import time
import datetime

import torch
from torch.nn import functional as F

from dassl.engine import TRAINER_REGISTRY, TrainerX
from dassl.metrics import compute_accuracy
from dassl.utils import (
    MetricMeter, AverageMeter, count_num_param, load_pretrained_weights
)
from dassl.engine import SimpleNet
from dassl.optim import build_optimizer
from dassl.modeling.ops import InvariancePenaltyLoss, InvariancePenaltyLoss2


@TRAINER_REGISTRY.register()
class IRM(TrainerX):
    """Invariant Risk Minimization."""

    def __init__(self, cfg):
        super().__init__(cfg)
        self.anneal_epochs = self.cfg.TRAINER.IRM.ANNEAL_EPOCHS
        self.num_domains = self.dm.num_source_domains
        if cfg.TRAINER.IRM.LOSS_TYPE == 'multiply':
            self.invariance_penalty_loss = InvariancePenaltyLoss().to(
                self.device
            )
        elif cfg.TRAINER.IRM.LOSS_TYPE == 'square':
            self.invariance_penalty_loss = InvariancePenaltyLoss2().to(
                self.device
            )
        else:
            raise ValueError('wrong IRM loss type')

        # self.global_iter = 0
        self.trade_off = 1.0

    def check_cfg(self, cfg):
        assert cfg.DATALOADER.TRAIN_X.SAMPLER == 'RandomDomainSampler'

    def before_epoch(self):
        if self.epoch == self.anneal_epochs:
            self.optim = build_optimizer(self.model, self.cfg.OPTIM)
            self.sched = torch.optim.lr_scheduler.CosineAnnealingLR(
                self.optim, self.max_epoch - self.anneal_epochs
            )
            self._optims['model'] = self.optim
            self._scheds['model'] = self.sched

        if self.epoch >= self.anneal_epochs:
            self.trade_off = self.cfg.TRAINER.IRM.TRADE_OFF
        else:
            self.trade_off = 1.0

    def forward_backward(self, batch):
        input, label = self.parse_batch_train(batch)
        output = self.model(input)
        # classification loss
        loss_ce = F.cross_entropy(output, label)
        # penalty loss
        loss_penalty = 0
        for output_per_domain, label_per_domain in zip(
            output.chunk(self.num_domains, dim=0),
            label.chunk(self.num_domains, dim=0)
        ):
            # normalize loss by domain num
            loss_penalty += self.invariance_penalty_loss(
                output_per_domain, label_per_domain
            ) / self.num_domains

        # if self.global_iter >= self.anneal_iters:
        #     trade_off = self.cfg.TRAINER.IRM.TRADE_OFF
        # else:
        #     trade_off = 1.0

        loss = loss_ce + loss_penalty * self.trade_off

        self.model_backward_and_update(loss)

        loss_summary = {
            'loss': loss.item(),
            'loss_ce': loss_ce.item(),
            'loss_penalty': loss_penalty.item(),
            'acc': compute_accuracy(output, label)[0].item()
        }

        # # update learning rate at every iteration
        # self.update_lr()

        if (self.batch_idx + 1) == self.num_batches:
            self.update_lr()

        return loss_summary

    def parse_batch_train(self, batch):
        input = batch['img']
        label = batch['label']
        input = input.to(self.device)
        label = label.to(self.device)
        return input, label

    # def run_epoch(self):
    #     self.set_model_mode('train')
    #     losses = MetricMeter()
    #     batch_time = AverageMeter()
    #     data_time = AverageMeter()
    #     self.num_batches = len(self.train_loader_x)

    #     end = time.time()
    #     for self.batch_idx, batch in enumerate(self.train_loader_x):

    #         self.global_iter = self.epoch * self.num_batches + self.batch_idx
    #         # reset optimizer to avoid sharp jump in gradient magnitudes
    #         if self.global_iter == self.anneal_iters:
    #             self.optim = build_optimizer(self.model, self.cfg.OPTIM)
    #             self.sched = torch.optim.lr_scheduler.CosineAnnealingLR(
    #                 self.optim,
    #                 self.max_epoch * len(self.train_loader_x) -
    #                 self.anneal_iters
    #             )
    #             self._optims['model'] = self.optim
    #             self._scheds['model'] = self.sched

    #         data_time.update(time.time() - end)
    #         loss_summary = self.forward_backward(batch)
    #         batch_time.update(time.time() - end)
    #         losses.update(loss_summary)

    #         if (self.batch_idx + 1) % self.cfg.TRAIN.PRINT_FREQ == 0:
    #             nb_this_epoch = self.num_batches - (self.batch_idx + 1)
    #             nb_future_epochs = (
    #                 self.max_epoch - (self.epoch + 1)
    #             ) * self.num_batches
    #             eta_seconds = batch_time.avg * (nb_this_epoch+nb_future_epochs)
    #             eta = str(datetime.timedelta(seconds=int(eta_seconds)))
    #             print(
    #                 'epoch [{0}/{1}][{2}/{3}]\t'
    #                 'time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
    #                 'data {data_time.val:.3f} ({data_time.avg:.3f})\t'
    #                 'eta {eta}\t'
    #                 '{losses}\t'
    #                 'lr {lr}'.format(
    #                     self.epoch + 1,
    #                     self.max_epoch,
    #                     self.batch_idx + 1,
    #                     self.num_batches,
    #                     batch_time=batch_time,
    #                     data_time=data_time,
    #                     eta=eta,
    #                     losses=losses,
    #                     lr=self.get_current_lr()
    #                 )
    #             )

    #         n_iter = self.epoch * self.num_batches + self.batch_idx
    #         for name, meter in losses.meters.items():
    #             self.write_scalar('train/' + name, meter.avg, n_iter)
    #         self.write_scalar('train/lr', self.get_current_lr(), n_iter)

    #         end = time.time()

    # def build_model(self):
    #     """Build and register model.

    #     The default builds a classification model along with its
    #     optimizer and scheduler.

    #     Custom trainers can re-implement this method if necessary.
    #     """
    #     cfg = self.cfg

    #     print('Building model')
    #     self.model = SimpleNet(cfg, cfg.MODEL, self.num_classes)
    #     if cfg.MODEL.INIT_WEIGHTS:
    #         load_pretrained_weights(self.model, cfg.MODEL.INIT_WEIGHTS)
    #     self.model.to(self.device)
    #     print('# params: {:,}'.format(count_num_param(self.model)))
    #     self.optim = build_optimizer(self.model, cfg.OPTIM)
    #     # adjust learning reate at every iteration
    #     self.sched = torch.optim.lr_scheduler.CosineAnnealingLR(
    #         self.optim, self.max_epoch * len(self.train_loader_x)
    #     )
    #     self.register_model('model', self.model, self.optim, self.sched)
