import pdb

import torch
import torch.nn.functional as F
import util.gan_util as gan_util
from data.transform import diffaug
from ignite.utils import convert_tensor
from kornia.augmentation import Resize


class ClassifierUpdater:
    def __init__(self, *args, **kwargs):
        self.classifier = kwargs.pop("classifier")
        self.generator = kwargs.pop("generator")
        self.finder = kwargs.pop("finder")
        self.optimizer_c = kwargs.pop("optimizer_c")
        self.optimizer_f = kwargs.pop("optimizer_f")
        self.device = kwargs.pop("device")
        self.ema_model = kwargs.pop("ema_model")
        self.lambda_p = kwargs.pop("lambda_p")
        self.u_accum_count = kwargs.pop("ubatch_ratio")
        self.batchsize_p = kwargs.pop("batchsize_p")
        self.warmup_epoch = kwargs.pop("warmup_epoch")
        self.resolution = kwargs.pop("resolution")
        self.resizer = Resize(size=self.resolution)
        self.loss = F.cross_entropy
        self.augment = diffaug

    def get_batch(self, batch):
        x, y = batch
        return (
            convert_tensor(x, device=self.device, non_blocking=True),
            convert_tensor(y, device=self.device, non_blocking=True),
        )

    def _sample_noize_and_label(self, n_gen_samples=None):
        if n_gen_samples is None:
            n_gen_samples = self.n_gen_samples
        gen = self.generator if (torch.cuda.device_count() < 2) else self.generator.module
        z = gan_util.sample_z(gen, n_gen_samples, self.device)
        y = gan_util.sample_categorical_labels(gen.num_classes, n_gen_samples, self.device)
        return z, y

    def __call__(self, engine, batch):
        report = {}
        self.classifier.train()

        # Generate pseudo samples and all logits
        x, y = self.get_batch(batch)
        batchsize = x.shape[0]
        z_p, y_p = self._sample_noize_and_label(n_gen_samples=self.batchsize_p)
        x_p = self.resizer(self.generator(z_p, y_p).detach())
        images = torch.cat([x, x_p], dim=0)
        logit_all, _ = self.classifier(images)
        logit_real, logit_p = torch.split(logit_all, [batchsize, self.batchsize_p], dim=0)

        # Calculate supervised loss
        loss_supervised = self.loss(logit_real, y)
        entropy = (-F.softmax(logit_real.detach(), dim=1) * F.log_softmax(logit_real.detach(), dim=1)).sum(dim=1).mean()
        report.update({"y_pred": logit_real.detach()})
        report.update({"y": y.detach()})
        report.update({"loss": loss_supervised.detach().item()})
        report.update({"entropy": entropy.detach().item()})

        # Calculate pseudo supervised loss
        loss_pseudo = self.loss(logit_p, y_p)
        lambda_p = self.lambda_p if self.warmup_epoch < engine.state.epoch else 0.0
        loss_pseudo = lambda_p * loss_pseudo
        loss_log = loss_pseudo.detach().item()
        report.update({"loss_pseudo": loss_log})

        loss_target = loss_supervised + loss_pseudo
        self.optimizer_c.zero_grad()
        loss_target.backward()
        self.optimizer_c.step()

        del x_p, logit_p, logit_all

        if self.ema_model is not None:
            self.ema_model.update(self.classifier)
        return report
