import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch import optim
import torch.utils.data as data
import numpy as np
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class MyDataset(data.Dataset):
    def __init__(self, specs, labels):
        self.specs = specs
        self.labels = labels

    def __getitem__(self, index):
        spec, target = self.specs[index], self.labels[index]
        return spec, target

    def __len__(self):
        return len(self.specs)


class AutoEncoder(nn.Module):
    def __init__(self, inputDim, hiddenDim):
        super().__init__()
        self.inputDim = inputDim
        self.hiddenDim = hiddenDim
        self.encoder = nn.Linear(inputDim, hiddenDim, bias=True)
        self.decoder = nn.Linear(hiddenDim, inputDim, bias=True)
        self.act = F.relu

    def forward(self, x, rep=False):
        hidden = self.encoder(x)
        hidden = self.act(hidden)
        if rep:
            return hidden
        else:
            out = self.decoder(hidden)
            return out


class SAE(nn.Module):
    def __init__(self, encoderList, output_dim):
        super().__init__()
        self.encoderList = encoderList
        self.en1 = encoderList[0]
        self.en2 = encoderList[1]
        self.fc = nn.Linear(128, output_dim, bias=True)  # 分类层输出维度为 num_classes

    def forward(self, x):
        out = x
        out = self.en1(out, rep=True)
        out = self.en2(out, rep=True)
        out = self.fc(out)
        return out


class SAE_net(object):
    def __init__(self, AE_epoch=200, SAE_epoch=200,
                 input_dim=404, hidden1_dim=512,
                 hidden2_dim=128, output_dim=4,  # 默认4类，可在调用时传入 num_classes
                 batch_size=128):
        self.AE_epoch = AE_epoch
        self.SAE_epoch = SAE_epoch
        self.input_dim = input_dim
        self.hidden1_dim = hidden1_dim
        self.hidden2_dim = hidden2_dim
        self.output_dim = output_dim
        self.batch_size = batch_size
        self.train_loader = None

        encoder1 = AutoEncoder(self.input_dim, self.hidden1_dim)
        encoder2 = AutoEncoder(self.hidden1_dim, self.hidden2_dim)
        self.encoder_list = [encoder1, encoder2]

    def trainAE(self, x_train, y_train, encoderList, trainLayer, batchSize, epoch, useCuda=False):
        if useCuda:
            for encoder in encoderList:
                encoder.to(device)

        optimizer = optim.Adam(encoderList[trainLayer].parameters())
        criterion = nn.MSELoss()

        data_train = MyDataset(x_train, y_train)
        self.train_loader = torch.utils.data.DataLoader(data_train, batch_size=batchSize, shuffle=True)

        for _ in range(epoch):
            for batch_idx, (x, target) in enumerate(self.train_loader):
                optimizer.zero_grad()
                if useCuda:
                    x, target = x.to(device), target.to(device)
                x = Variable(x).type(torch.FloatTensor)
                x = x.view(x.size(0), -1)

                out = x
                if trainLayer != 0:
                    for i in range(trainLayer):
                        out = encoderList[i](out, rep=True)

                pred = encoderList[trainLayer](out, rep=False).cpu()
                loss = criterion(pred, out)
                loss.backward()
                optimizer.step()

    def trainClassifier(self, model, epoch, useCuda=False):
        if useCuda:
            model = model.to(device)

        for param in model.parameters():
            param.requires_grad = True

        optimizer = optim.Adam(model.parameters())
        criterion = nn.CrossEntropyLoss()

        for _ in range(epoch):
            for batch_idx, (x, target) in enumerate(self.train_loader):
                optimizer.zero_grad()
                if useCuda:
                    x, target = x.to(device), target.to(device)
                x = Variable(x).type(torch.FloatTensor)
                x = x.view(-1, self.input_dim)

                out = model(x)
                loss = criterion(out, target)
                loss.backward()
                optimizer.step()
        self.model = model

    def fit(self, x_train=None, y_train=None, X_test=None, y_test=None):
        x_train = x_train[:, np.newaxis, :]
        x_train = torch.from_numpy(x_train).float()

        for i in range(2):
            self.trainAE(x_train=x_train, y_train=y_train,
                         encoderList=self.encoder_list, trainLayer=i,
                         batchSize=self.batch_size, epoch=self.AE_epoch)

        model = SAE(encoderList=self.encoder_list, output_dim=self.output_dim)

        # 训练分类器并获取训练集的评估指标
        train_accuracy, train_precision, train_recall, train_f1, train_cm = self.trainClassifier(model=model, epoch=self.SAE_epoch, X_train=x_train, y_train=y_train)

        # 计算测试集的评估指标
        test_accuracy, test_precision, test_recall, test_f1, test_cm = self.evaluate(model, X_test, y_test)

        # 返回训练集和测试集的评估结果
        train_metrics = {
            "accuracy": train_accuracy,
            "precision": train_precision,
            "recall": train_recall,
            "f1_score": train_f1,
            "confusion_matrix": train_cm
        }

        test_metrics = {
            "accuracy": test_accuracy,
            "precision": test_precision,
            "recall": test_recall,
            "f1_score": test_f1,
            "confusion_matrix": test_cm
        }

        return train_metrics, test_metrics

    def evaluate(self, model, X_test, y_test):
        X_test = torch.from_numpy(X_test).float()
        X_test = X_test[:, np.newaxis, :]
        X_test = Variable(X_test).view(-1, self.input_dim)

        out = model(X_test)
        _, y_pred = torch.max(out, 1)

        # 计算准确率、精确率、召回率、F1分数和混淆矩阵
        accuracy = accuracy_score(y_test, y_pred.numpy())
        precision = precision_score(y_test, y_pred.numpy(), average='weighted')
        recall = recall_score(y_test, y_pred.numpy(), average='weighted')
        f1 = f1_score(y_test, y_pred.numpy(), average='weighted')
        cm = confusion_matrix(y_test, y_pred.numpy())

        return accuracy, precision, recall, f1, cm


def SAE(X_train, y_train, X_test, y_test, num_classes=4):
    clf = SAE_net(output_dim=num_classes)
    train_metrics, test_metrics = clf.fit(X_train, y_train, X_test, y_test)

    return train_metrics, test_metrics