from torch.nn import functional as F
import torch
from dassl.engine import TRAINER_REGISTRY, TrainerX
from dassl.metrics import compute_accuracy
from dassl.modeling.ops import random_efdmix, crossdomain_efdmix
import numpy as np
import os.path as osp
from dassl.utils import FocalLoss
import pdb
import torch.nn as nn
@TRAINER_REGISTRY.register()
class Vanilla2(TrainerX):
    """Vanilla baseline.

    Slightly modified for mixstyle.
    """

    def __init__(self, cfg):
        super().__init__(cfg)
        mix = cfg.TRAINER.VANILLA2.MIX
        self.cfg = cfg

        if mix == 'random':
            self.model.apply(random_efdmix)
            print('EFDMix: random mixing')

        elif mix == 'crossdomain':
            self.model.apply(crossdomain_efdmix)
            print('EFDMix: cross-domain mixing')

        else:
            raise NotImplementedError

        self.hloss = HLoss()

    def forward_backward(self, batch, class_weight=None):
        input, label, domain = self.parse_batch_train(batch)
        if 'cb_uniform' in self.cfg.MODEL.BACKBONE.NAME:
            output, statistic = self.model(input, domain, label)
            #loss = F.cross_entropy(output, label, weight=class_weight)
        else:
            output = self.model(input, domain)

        loss = F.cross_entropy(output, label) 
        
        
        #c1 = self.cls1(entropy_loss[0])
        #c2 = self.cls2(entropy_loss[1])
        #loss_entropy_1 = self.hloss(en[0])
        #loss_entropy_2 = self.hloss(en[1])

        #loss_ce1 = F.cross_entropy(c1, label)
        #loss_ce2 =  F.cross_entropy(c2, label)
        
        
        #loss = loss + (loss_entropy_1 + loss_entropy_2) * 0.5
        #loss_ce = (loss_ce1 + loss_ce2) * 0.3

        
        
        #focal_loss = FocalLoss()
        #loss = focal_loss(output,label)
        
        self.model_backward_and_update(loss)

        #self.optim_cls1.zero_grad()
        #self.optim_cls2.zero_grad()
        #loss_ce.backward()
        #self.optim_cls1.step()
        #self.optim_cls2.step()


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

        if (self.batch_idx + 1) == self.num_batches:
            self.update_lr()
        if 'cb_uniform' in self.cfg.MODEL.BACKBONE.NAME:
            return loss_summary, statistic
        else:
            return loss_summary

    def parse_batch_train(self, batch):
        input = batch['img']
        label = batch['label']
        domain = batch['domain']
        input = input.to(self.device)
        label = label.to(self.device)
        return input, label, domain
    
    @torch.no_grad()
    def vis(self):
        self.set_model_mode('eval')
        output_dir = self.cfg.OUTPUT_DIR
        source_domains = self.cfg.DATASET.SOURCE_DOMAINS
        print('Source domains:', source_domains)

        out_embed = []
        out_domain = []
        out_label = []

        split = 'test' #self.cfg.TEST.SPLIT
        data_loader = self.val_loader if split == 'val' else self.test_loader

        print('Extracting style features')

        for batch_idx, batch in enumerate(self.train_loader_x):
            input = batch['img'].to(self.device)
            label = batch['label']
            domain = batch['domain']
            impath = batch['impath']

            # model should directly output features or style statistics
            #raise NotImplementedError
            _, output = self.model(input, domain=domain)
            #pdb.set_trace()
            output = torch.cat(output[0:2], dim=1)
            output = torch.squeeze(output)
            output = output.cpu().numpy()
            out_embed.append(output)
            out_domain.append(domain.numpy())
            out_label.append(label.numpy()) # CLASS LABEL

            print('processed batch-{}'.format(batch_idx + 1))

        for batch_idx, batch in enumerate(self.test_loader):
            input = batch['img'].to(self.device)
            label = batch['label']
            domain = batch['domain']+3
            impath = batch['impath']

            # model should directly output features or style statistics
            #raise NotImplementedError
            _, output = self.model(input, domain=domain)
            #pdb.set_trace()
            output = torch.cat(output[0:2], dim=1)
            output = torch.squeeze(output)
            output = output.cpu().numpy()
            out_embed.append(output)
            out_domain.append(domain.numpy())
            out_label.append(label.numpy()) # CLASS LABEL#

            print('processed batch-{}'.format(batch_idx + 1))


        out_embed = np.concatenate(out_embed, axis=0)
        out_domain = np.concatenate(out_domain, axis=0)
        out_label = np.concatenate(out_label, axis=0)

        print('shape of feature matrix:', out_embed.shape)
        out = {
            'embed': out_embed,
            'domain': out_domain,
            'dnames': source_domains,
            'label': out_label
        }
        out_path = osp.join(output_dir, 'embed.pt')
        torch.save(out, out_path)
        print('File saved to "{}"'.format(out_path))



class HLoss(nn.Module):
    def __init__(self):
        super(HLoss, self).__init__()
        self.criterion = torch.nn.KLDivLoss(reduction='batchmean')

    def forward(self, x):
        #b = F.softmax(x,dim=1) * F.log_softmax(x+1e-6, dim=1)
        #b = torch.sum(b, 1)
        pred = F.log_softmax(x+1e-6, dim=1)
        uniform = torch.ones_like(pred) / pred.size(1)
        #pdb.set_trace()
        
        return self.criterion(pred, uniform)



