import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.optim as optim
import numpy as np
import util
from termcolor import cprint

import itertools
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
import torch.nn.functional as F
from scipy.stats import entropy
from sklearn.metrics import confusion_matrix

import time


class CLASSIFIER:
    # train_Y is interger
    # CLASSIFIER(syn_feature,util.map_label(syn_label,data.unseenclasses),data,data.unseenclasses.size(0),opt.cuda,opt.classifier_lr, 0.5, 25, opt.syn_num, False)
    def __init__(self, _train_X, _train_Y, data_loader, _nclass, _cuda, _lr=0.001, _beta1=0.5, _nepoch=20,
                 _batch_size=100, generalized=True, epoch=20):
        self.train_X = _train_X
        self.train_Y = _train_Y
        self.test_seen_feature = data_loader.test_seen_feature
        self.test_seen_label = data_loader.test_seen_label
        self.test_unseen_feature = data_loader.test_unseen_feature
        self.test_unseen_label = data_loader.test_unseen_label
        self.seenclasses = data_loader.seenclasses
        self.unseenclasses = data_loader.unseenclasses
        self.batch_size = _batch_size
        self.nepoch = _nepoch
        self.nclass = _nclass
        self.input_dim = _train_X.size(1)
        self.cuda = _cuda
        self.model = LINEAR_LOGSOFTMAX(self.input_dim, self.nclass)
        self.model.apply(util.weights_init)
        self.criterion = nn.NLLLoss()

        self.data = data_loader

        self.input = torch.FloatTensor(_batch_size, self.input_dim)
        self.label = torch.LongTensor(_batch_size)

        self.lr = _lr
        self.beta1 = _beta1
        # setup optimizer
        self.optimizer = optim.Adam(self.model.parameters(), lr=_lr, betas=(_beta1, 0.999))
        self.epoch = epoch

        if self.cuda:
            self.model.cuda()
            self.criterion.cuda()
            self.input = self.input.cuda()
            self.label = self.label.cuda()

        self.index_in_epoch = 0
        self.epochs_completed = 0
        self.ntrain = self.train_X.size()[0]
        self.backup_X = _train_X
        self.backup_Y = _train_Y

        if generalized:
            self.acc_seen, self.acc_unseen, self.H = self.fit()
        else:
            self.acc = self.fit_zsl()

    def pairwise_distances(self, x, y=None):
        '''
        Input: x is a Nxd matrix
               y is an optional Mxd matirx
        Output: dist is a NxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:]
                if y is not given then use 'y=x'.
        i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2
        '''
        x_norm = (x ** 2).sum(1).view(-1, 1)
        if y is not None:
            y_t = torch.transpose(y, 0, 1)
            y_norm = (y ** 2).sum(1).view(1, -1)
        else:
            y_t = torch.transpose(x, 0, 1)
            y_norm = x_norm.view(1, -1)

        dist = x_norm + y_norm - 2.0 * torch.mm(x, y_t)
        # Ensure diagonal is zero if x=y
        if y is None:
            dist = dist - torch.diag(dist.diag)
        return torch.clamp(dist, 0.0, np.inf)

    def fit_zsl(self):
        first_acc = 0
        first_all_pred = None
        first_all_output = None
        first_all_acc = []

        trun = lambda x: int(x * 100)

        all_length = self.test_unseen_feature.size(0)
        mapped_test_label = util.map_label(self.test_unseen_label, self.unseenclasses)
        for epoch in range(self.nepoch):
            for i in range(0, self.ntrain, self.batch_size):
                self.model.zero_grad()
                batch_input, batch_label = self.next_batch(self.batch_size)
                self.input.copy_(batch_input)
                self.label.copy_(batch_label)
                inputv = Variable(self.input)  # fake_feature
                labelv = Variable(self.label)  # fake_labels
                output = self.model(inputv)
                loss = self.criterion(output, labelv)  # 使用fake_unseen_feature和labels来训练分类器
                loss.backward()
                self.optimizer.step()

            acc, pred, output, all_acc = self.val(self.test_unseen_feature, self.test_unseen_label, self.unseenclasses)
            if acc > first_acc:
                first_acc = acc
                first_all_pred = pred
                first_all_output = output
                first_all_acc = all_acc

        print('First Acc: {:.2f}%'.format(first_acc * 100))
        # print([trun(x) for  x in  list(first_all_acc)])
        # cprint('First Acc: {:.2f}%'.format(first_acc * 100),'red')
        sys.stdout.flush()
        return first_acc

    def split_pred(self, all_pred, real_label):
        seen_pred = None
        seen_label = None
        unseen_pred = None
        unseen_label = None
        for i in self.seenclasses:
            idx = (real_label == i)
            if seen_pred is None:
                seen_pred = all_pred[idx]
                seen_label = real_label[idx]
            else:
                seen_pred = torch.cat((seen_pred, all_pred[idx]), 0)
                seen_label = torch.cat((seen_label, real_label[idx]))

        for i in self.unseenclasses:
            idx = (real_label == i)
            if unseen_pred is None:
                unseen_pred = all_pred[idx]
                unseen_label = real_label[idx]
            else:
                unseen_pred = torch.cat((unseen_pred, all_pred[idx]), 0)
                unseen_label = torch.cat((unseen_label, real_label[idx]), 0)

        return seen_pred, seen_label, unseen_pred, unseen_label

    # for gzsl
    def fit(self):
        # 3个length
        test_seen_length = self.test_seen_feature.shape[0]  # 1764
        test_unseen_length = self.test_unseen_feature.shape[0]  # 2967
        all_length = test_seen_length + test_unseen_length
        all_test_feature = torch.cat((self.test_seen_feature, self.test_unseen_feature), 0)
        all_test_label = torch.cat((self.test_seen_label, self.test_unseen_label), 0)#[4958+5685]
        all_classes = torch.sort(torch.cat((self.seenclasses, self.unseenclasses), 0))[0]
        first_acc = 0
        first_all_pred = None
        first_all_output = None

        best_H = 0
        for epoch in range(self.nepoch):
            for i in range(0, self.ntrain, self.batch_size):  # self.ntrain=22057, self.batch_size=300
                self.model.zero_grad()
                batch_input, batch_label = self.next_batch(self.batch_size)
                self.input.copy_(batch_input)
                self.label.copy_(batch_label)

                inputv = Variable(self.input)
                labelv = Variable(self.label)
                output = self.model(inputv)
                loss = self.criterion(output, labelv)
                loss.backward()
                self.optimizer.step()

            acc_seen, pred_seen, output_seen = self.val_gzsl(self.test_seen_feature, self.test_seen_label,self.seenclasses)
            # print(acc_seen)#float
            # print(pred_seen.shape)#[4958]
            # print(output_seen.shape)#[4958,50]
            acc_unseen, pred_unseen, output_unseen = self.val_gzsl(self.test_unseen_feature, self.test_unseen_label,self.unseenclasses)
            # print(acc_unseen)#float
            # print(pred_unseen.shape)#[5685]
            # print(output_unseen.shape)#[5685,50]
            H = 2 * acc_seen * acc_unseen / (acc_seen + acc_unseen)
            if H >= best_H:
                best_H = H
                first_all_pred = torch.cat((pred_seen, pred_unseen), 0)
                first_all_output = torch.cat((output_seen, output_unseen), 0)
        first_seen_pred, first_seen_label, first_unseen_pred, first_unseen_label = self.split_pred(first_all_pred,all_test_label)
        # print(first_seen_pred.shape)#[4958],sort labels
        # print(first_seen_label.shape)#[4958]
        # print(first_unseen_pred.shape)#[5685]
        # print(first_unseen_label.shape)#[5685]
        # def compute_per_class_acc_gzsl(self, test_label, predicted_label, target_classes):
        acc_first_seen = self.compute_per_class_acc_gzsl(first_seen_label, first_seen_pred, self.seenclasses)
        acc_first_unseen = self.compute_per_class_acc_gzsl(first_unseen_label, first_unseen_pred, self.unseenclasses)
        acc_first_H = 2 * acc_first_seen * acc_first_unseen / (acc_first_seen + acc_first_unseen)
        print('First Seen: {:.2f}%, Unseen: {:.2f}%, First H: {:.2f}%'.format(acc_first_seen * 100,acc_first_unseen * 100,acc_first_H * 100))

        sys.stdout.flush()
        return acc_first_seen, acc_first_unseen, acc_first_H

    def val(self, test_X, test_label, target_classes, second=False):
        start = 0
        ntest = test_X.size()[0]
        predicted_label = torch.LongTensor(test_label.size())
        all_output = None
        for i in range(0, ntest, self.batch_size):
            end = min(ntest, start + self.batch_size)
            if self.cuda:
                output = self.model(Variable(test_X[start:end].cuda(), volatile=True))
            else:
                output = self.model(Variable(test_X[start:end], volatile=True))
            if all_output is None:
                all_output = output
            else:
                all_output = torch.cat((all_output, output), 0)
            _, predicted_label[start:end] = torch.max(output.data, 1)
            start = end
        acc = self.compute_per_class_acc(util.map_label(test_label, target_classes), predicted_label,
                                         target_classes.size(0))
        acc_all = self.compute_every_class_acc(util.map_label(test_label, target_classes), predicted_label,
                                               target_classes.size(0))
        return acc, predicted_label, all_output, acc_all

    def val_gzsl(self, test_X, test_label, target_classes):
        start = 0
        ntest = test_X.size()[0]
        predicted_label = torch.LongTensor(test_label.size())
        all_output = None
        for i in range(0, ntest, self.batch_size):
            end = min(ntest, start + self.batch_size)
            if self.cuda:
                output = self.model(Variable(test_X[start:end].cuda(), volatile=True))
            else:
                output = self.model(Variable(test_X[start:end], volatile=True))

            if all_output is None:
                all_output = output
            else:
                all_output = torch.cat((all_output, output), 0)
            _, predicted_label[start:end] = torch.max(output.data, 1)
            start = end
        # acc = self.compute_per_class_acc(util.map_label(test_label, target_classes), predicted_label, target_classes.size(0))
        acc = self.compute_per_class_acc_gzsl(test_label, predicted_label, target_classes)
        return acc, predicted_label, all_output

    def next_batch(self, batch_size):
        start = self.index_in_epoch
        # shuffle the data at the first epoch
        if self.epochs_completed == 0 and start == 0:
            perm = torch.randperm(self.ntrain)
            self.train_X = self.train_X[perm]
            self.train_Y = self.train_Y[perm]
        # the last batch
        if start + batch_size > self.ntrain:
            self.epochs_completed += 1
            rest_num_examples = self.ntrain - start
            if rest_num_examples > 0:
                X_rest_part = self.train_X[start:self.ntrain]
                Y_rest_part = self.train_Y[start:self.ntrain]
            # shuffle the data
            perm = torch.randperm(self.ntrain)
            self.train_X = self.train_X[perm]
            self.train_Y = self.train_Y[perm]
            # start next epoch
            start = 0
            self.index_in_epoch = batch_size - rest_num_examples
            end = self.index_in_epoch
            X_new_part = self.train_X[start:end]
            Y_new_part = self.train_Y[start:end]
            if rest_num_examples > 0:
                return torch.cat((X_rest_part, X_new_part), 0), torch.cat((Y_rest_part, Y_new_part), 0)
            else:
                return X_new_part, Y_new_part
        else:
            self.index_in_epoch += batch_size
            end = self.index_in_epoch
            return self.train_X[start:end], self.train_Y[start:end]

    def compute_per_class_acc_gzsl(self, test_label, predicted_label, target_classes):
        acc_per_class = 0
        for i in target_classes:
            idx = (test_label == i)
            if torch.sum(idx).float() == 0:
                continue
            else:
                acc_per_class += torch.sum(test_label[idx] == predicted_label[idx]).float() / torch.sum(idx).float()
        acc_per_class /= target_classes.size(0)
        return acc_per_class

    def compute_per_class_acc(self, test_label, predicted_label, nclass):
        acc_per_class = torch.FloatTensor(nclass).fill_(0)
        for i in range(nclass):
            idx = (test_label == i)
            if torch.sum(idx).float() != 0:
                acc_per_class[i] = torch.sum(test_label[idx] == predicted_label[idx]).float() / torch.sum(idx).float()
        return acc_per_class.mean()

    def compute_every_class_acc(self, test_label, predicted_label, nclass):
        acc_per_class = torch.FloatTensor(nclass).fill_(0)
        for i in range(nclass):
            idx = (test_label == i)
            if torch.sum(idx).float() != 0:
                acc_per_class[i] = torch.sum(test_label[idx] == predicted_label[idx]).float() / torch.sum(idx).float()
        return acc_per_class


class LINEAR_LOGSOFTMAX(nn.Module):
    def __init__(self, input_dim, nclass):
        super(LINEAR_LOGSOFTMAX, self).__init__()
        self.fc = nn.Linear(input_dim, nclass)
        self.logic = nn.LogSoftmax(dim=1)

    def forward(self, x):
        o = self.logic(self.fc(x))
        return o