from tqdm import tqdm
import abc, os
from models.utils import build_model
import torch
import torch.nn.functional as F
import torch.nn as nn
from common.loss_function import TaylorCE, NCELoss, MAELoss, AUELoss, AExpLoss, FocalLoss, GCE, SCE, PHuberCE, CoresLoss, NCEandMAE, NGCEandMAE, NCEandAGCE
class SingleModel:
    __metaclass__ = abc.ABCMeta
    def __init__(self, args, device, num_classes, train_loader):
        self.device = device
        self.args = args
        self.num_classes = num_classes
        # Create model
        self.net = build_model(args.model, num_classes, device,args)
        self.iterations = 0

        self.optimizer_model = torch.optim.SGD(self.net.parameters(), args.learning_rate,
                                            momentum=args.momentum, weight_decay=args.decay)
        if len(self.args.gpu) > 1:
            gpu_list = [int(s) for s in args.gpu.split(',')]
            self.optimizer_model = nn.DataParallel(self.optimizer_model, device_ids=gpu_list)
            self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer_model.module, [80,140], gamma=0.1)
        else:
            self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer_model, [80,140], gamma=0.1)

        self.loss_function = self.get_loss(args.loss, args)


    def train(self, train_loader, epoch):
        self.net.train()
        loss_avg = 0.0
        for data, target, index in tqdm(train_loader):
            loss = self.train_batch(index, data, target, epoch)
            # backward

            if len(self.args.gpu) > 1:
                self.optimizer_model.module.zero_grad()
                loss.backward()
                self.optimizer_model.module.step()
            else:
                self.optimizer_model.zero_grad()
                loss.backward()
                self.optimizer_model.step()
            # exponential moving average
            loss_avg = loss_avg * 0.8 + float(loss) * 0.2

        self.scheduler.step()
        return loss_avg


    @abc.abstractmethod
    def train_batch(self, batch_idx, inputs, targets, epoch):
        pass

    def test(self, test_loader):
        self.net.eval()
        loss_avg = 0.0
        correct = 0
        with torch.no_grad():
            for dict in test_loader:
                data, target = dict[0].to(self.device), dict[1].to(self.device)

                # forward
                output = self.net(data)
                loss = F.cross_entropy(output, target)

                # accuracy
                pred = output.data.max(1)[1]
                correct += pred.eq(target.data).sum().item()

                # test loss average
                loss_avg += float(loss.data)

        return loss_avg / len(test_loader), correct / len(test_loader.dataset)

    def get_loss(self, loss_type, args):

        CIFAR10_CONFIG = {
            "ce": nn.CrossEntropyLoss(),
            "focal": FocalLoss(gamma=0.5),
            "mae": MAELoss(num_classes=self.num_classes),
            "gce": GCE(self.device, k=self.num_classes),
            "sce": SCE(alpha=0.5, beta=1.0, num_classes=self.num_classes),
            "nce": NCELoss(num_classes=self.num_classes),
            "ael": AExpLoss(num_classes=10, a=2.5),
            "aul": AUELoss(num_classes=10, a=5.5, q=3),
            "phuber": PHuberCE(tau=10),
            "taylor": TaylorCE(device=self.device, series=args.series),
            "cores": CoresLoss(device=self.device),
            "ncemae": NCEandMAE(alpha=1, beta=1, num_classes=10),
            "nceagce": NCEandAGCE(alpha=1, beta=4, a=6, q=1.5, num_classes=10),
        }

        CIFAR100_CONFIG = {
            "ce": nn.CrossEntropyLoss(),
            "focal": FocalLoss(gamma=0.5),
            "mae": MAELoss(num_classes=self.num_classes),
            "gce": GCE(self.device, k=self.num_classes),
            "sce": SCE(alpha=0.5, beta=1.0, num_classes=self.num_classes),
            "nce": NCELoss(num_classes=self.num_classes),
            "ael": AExpLoss(num_classes=100, a=2.5),
            "aul": AUELoss(num_classes=100, a=5.5, q=3),
            "phuber": PHuberCE(tau=30),
            "taylor": TaylorCE(device=self.device, series=args.series),
            "cores": CoresLoss(device=self.device),
            "ncemae": NCEandMAE(alpha=50, beta=1, num_classes=100),
            "nceagce": NCEandAGCE(alpha=50, beta=0.1, a=1.8, q=3.0, num_classes=100),
        }

        if args.dataset == "cifar10":
            return CIFAR10_CONFIG[loss_type]
        elif args.dataset == "cifar100":
            return CIFAR100_CONFIG[loss_type]

