"""

EEG conformer 

Test on the datasets 2b

"""

import argparse

import os

gpus = [1]

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'

os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, gpus))

import numpy as np

import math

import glob

import random

import itertools

import datetime

import time

import datetime

import sys

import scipy.io

import torchvision.transforms as transforms

from torchvision.utils import save_image, make_grid

from torch.utils.data import DataLoader

from torch.autograd import Variable

from torchsummary import summary

import torch.autograd as autograd

from torchvision.models import vgg19

import torch.nn as nn

import torch.nn.functional as F

import torch

import torch.nn.init as init

from torch.utils.data import Dataset

from PIL import Image

import torchvision.transforms as transforms

from sklearn.decomposition import PCA

import torch

import torch.nn.functional as F

import matplotlib.pyplot as plt

from torch import nn

from torch import Tensor

from PIL import Image

from torchvision.transforms import Compose, Resize, ToTensor

from einops import rearrange, reduce, repeat

from einops.layers.torch import Rearrange, Reduce

# from common_spatial_pattern import csp

import matplotlib.pyplot as plt

# from torch.utils.tensorboard import SummaryWriter

from torch.backends import cudnn

cudnn.benchmark = False

cudnn.deterministic = True

class PatchEmbedding(nn.Module):

    def __init__(self, emb_size=40):

        # self.patch_size = patch_size

        super().__init__()

        self.shallownet = nn.Sequential(

            nn.Conv2d(1, 40, (1, 25), (1, 1)),

            nn.Conv2d(40, 40, (3, 1), (1, 1)),

            nn.BatchNorm2d(40),

            nn.ELU(),

            nn.AvgPool2d((1, 75), (1, 15)),

            nn.Dropout(0.5),

        )

        self.projection = nn.Sequential(

            nn.Conv2d(40, emb_size, (1, 1), stride=(1, 1)),  # 5 is better than 1

            Rearrange('b e (h) (w) -> b (h w) e'),

        )

    def forward(self, x: Tensor) -> Tensor:

        b, _, _, _ = x.shape

        x = self.shallownet(x)

        x = self.projection(x)

        return x

class MultiHeadAttention(nn.Module):

    def __init__(self, emb_size, num_heads, dropout):

        super().__init__()

        self.emb_size = emb_size

        self.num_heads = num_heads

        self.keys = nn.Linear(emb_size, emb_size)

        self.queries = nn.Linear(emb_size, emb_size)

        self.values = nn.Linear(emb_size, emb_size)

        self.att_drop = nn.Dropout(dropout)

        self.projection = nn.Linear(emb_size, emb_size)

    def forward(self, x: Tensor, mask: Tensor = None) -> Tensor:

        queries = rearrange(self.queries(x), "b n (h d) -> b h n d", h=self.num_heads)

        keys = rearrange(self.keys(x), "b n (h d) -> b h n d", h=self.num_heads)

        values = rearrange(self.values(x), "b n (h d) -> b h n d", h=self.num_heads)

        energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys)  # batch, num_heads, query_len, key_len

        if mask is not None:

            fill_value = torch.finfo(torch.float32).min

            energy.mask_fill(~mask, fill_value)

        scaling = self.emb_size ** (1 / 2)

        att = F.softmax(energy / scaling, dim=-1)

        att = self.att_drop(att)

        out = torch.einsum('bhal, bhlv -> bhav ', att, values)

        out = rearrange(out, "b h n d -> b n (h d)")

        out = self.projection(out)

        return out

class ResidualAdd(nn.Module):

    def __init__(self, fn):

        super().__init__()

        self.fn = fn

    def forward(self, x, **kwargs):

        res = x

        x = self.fn(x, **kwargs)

        x += res

        return x

class FeedForwardBlock(nn.Sequential):

    def __init__(self, emb_size, expansion, drop_p):

        super().__init__(

            nn.Linear(emb_size, expansion * emb_size),

            nn.GELU(),

            nn.Dropout(drop_p),

            nn.Linear(expansion * emb_size, emb_size),

        )

class GELU(nn.Module):

    def forward(self, input: Tensor) -> Tensor:

        return input*0.5*(1.0+torch.erf(input/math.sqrt(2.0)))

class TransformerEncoderBlock(nn.Sequential):

    def __init__(self,

                 emb_size,

                 num_heads=5,

                 drop_p=0.5,

                 forward_expansion=4,

                 forward_drop_p=0.5):

        super().__init__(

            ResidualAdd(nn.Sequential(

                nn.LayerNorm(emb_size),

                MultiHeadAttention(emb_size, num_heads, drop_p),

                nn.Dropout(drop_p)

            )),

            ResidualAdd(nn.Sequential(

                nn.LayerNorm(emb_size),

                FeedForwardBlock(

                    emb_size, expansion=forward_expansion, drop_p=forward_drop_p),

                nn.Dropout(drop_p)

            )

            ))

class TransformerEncoder(nn.Sequential):

    def __init__(self, depth, emb_size):

        super().__init__(*[TransformerEncoderBlock(emb_size) for _ in range(depth)])

class ClassificationHead(nn.Sequential):

    def __init__(self, emb_size, n_classes):

        super().__init__()

        self.cov = nn.Sequential(

            nn.Conv1d(190, 1, 1, 1),

            nn.LeakyReLU(0.2),

            nn.Dropout(0.5)

        )

        self.clshead = nn.Sequential(

            Reduce('b n e -> b e', reduction='mean'),

            nn.LayerNorm(emb_size),

            nn.Linear(emb_size, n_classes)

        )

        self.clshead_fc = nn.Sequential(

            Reduce('b n e -> b e', reduction='mean'),

            nn.LayerNorm(emb_size),

            nn.Linear(emb_size, 32),

            nn.ELU(),

            nn.Dropout(0.5),

            nn.Linear(32, n_classes)

        )

        self.fc = nn.Sequential(

            nn.Linear(2440, 256),

            nn.ELU(),

            nn.Dropout(0.5),

            nn.Linear(256, 32),

            nn.ELU(),

            nn.Dropout(0.3),

            nn.Linear(32, 2)

        )

    def forward(self, x):

        x = x.contiguous().view(x.size(0), -1)

        out = self.fc(x)

        return x, out

# ! Rethink the use of Transformer for EEG signal

class ViT(nn.Sequential):

    def __init__(self, emb_size=40, depth=10, n_classes=2, **kwargs):

        super().__init__(

            PatchEmbedding(emb_size),

            TransformerEncoder(depth, emb_size),

            ClassificationHead(emb_size, n_classes)

        )

class ExGAN():

    def __init__(self, nsub):

        super(ExGAN, self).__init__()

        self.batch_size = 100

        self.n_epochs = 2000

        self.img_height = 22

        self.img_width = 600

        self.channels = 1

        self.c_dim = 4

        self.lr = 0.0002

        self.b1 = 0.5

        self.b2 = 0.999

        self.alpha = 0.0002

        self.dimension = (190, 50)

        self.nSub = nsub

        self.start_epoch = 0

        self.root = '/Data/strict_TE/2b/'

        self.pretrain = False

        self.log_write = open("/Code/CT/results/cf/2b/log_subject%d.txt" % self.nSub, "w")

        self.img_shape = (self.channels, self.img_height, self.img_width)

        self.Tensor = torch.cuda.FloatTensor

        self.LongTensor = torch.cuda.LongTensor

        self.criterion_l1 = torch.nn.L1Loss().cuda()

        self.criterion_l2 = torch.nn.MSELoss().cuda()

        self.criterion_cls = torch.nn.CrossEntropyLoss().cuda()

        self.model = ViT().cuda()

        self.model = nn.DataParallel(self.model, device_ids=[i for i in range(len(gpus))])

        self.model = self.model.cuda()

        self.centers = {}

    def interaug(self, timg, label):

        aug_data = []

        aug_label = []

        for cls4aug in range(2):

            cls_idx = np.where(label == cls4aug + 1)

            tmp_data = timg[cls_idx]

            tmp_label = label[cls_idx]

            tmp_aug_data = np.zeros((int(self.batch_size / 2), 1, 3, 1000))

            for ri in range(int(self.batch_size / 2)):

                for rj in range(8):

                    rand_idx = np.random.randint(0, tmp_data.shape[0], 8)

                    tmp_aug_data[ri, :, :, rj * 125:(rj + 1) * 125] = tmp_data[rand_idx[rj], :, :,

                                                                      rj * 125:(rj + 1) * 125]

            aug_data.append(tmp_aug_data)

            aug_label.append(tmp_label[:int(self.batch_size / 2)])

        aug_data = np.concatenate(aug_data)

        aug_label = np.concatenate(aug_label)

        aug_shuffle = np.random.permutation(len(aug_data))

        aug_data = aug_data[aug_shuffle, :, :]

        aug_label = aug_label[aug_shuffle]

        aug_data = torch.from_numpy(aug_data).cuda()

        aug_data = aug_data.float()

        aug_label = torch.from_numpy(aug_label-1).cuda()

        aug_label = aug_label.long()

        return aug_data, aug_label

    def get_source_data(self):

        # to get the data of target subject

        train_data = []

        train_label = []

        for session_index in range(3):

            target_tmp = scipy.io.loadmat(self.root + 'B0%d0%dT.mat' % (self.nSub, session_index+1))

            train_data_tmp = target_tmp['data']

            train_label_tmp = target_tmp['label']

            train_data_tmp = np.transpose(train_data_tmp, (2, 1, 0))

            train_data_tmp = np.expand_dims(train_data_tmp, axis=1)

            train_label_tmp = np.transpose(train_label_tmp)

            train_label_tmp = train_label_tmp[0]

            train_data.append(train_data_tmp)

            train_label.append(train_label_tmp)

        self.allData = np.concatenate(train_data)

        self.allLabel = np.concatenate(train_label)

        shuffle_num = np.random.permutation(len(self.allData))

        self.allData = self.allData[shuffle_num, :, :, :]

        self.allLabel = self.allLabel[shuffle_num]

        # test data

        test_data = []

        test_label = []

        for session_index in range(2):

            test_tmp = scipy.io.loadmat(self.root + 'B0%d0%dE.mat' % (self.nSub, session_index+4))

            test_data_tmp = test_tmp['data']

            test_label_tmp = test_tmp['label']

            test_data_tmp = np.transpose(test_data_tmp, (2, 1, 0))

            test_data_tmp = np.expand_dims(test_data_tmp, axis=1)

            test_label_tmp = np.transpose(test_label_tmp)

            test_label_tmp = test_label_tmp[0]

            test_data.append(test_data_tmp)

            test_label.append(test_label_tmp)

        self.testData = np.concatenate(test_data)

        self.testLabel = np.concatenate(test_label)

        # standardize

        target_mean = np.mean(self.allData)

        target_std = np.std(self.allData)

        self.allData = (self.allData - target_mean) / target_std

        self.testData = (self.testData - target_mean) / target_std

        return self.allData, self.allLabel, self.testData, self.testLabel

    def update_lr(self, optimizer, lr):

        for param_group in optimizer.param_groups:

            param_group['lr'] = lr

    def aug(self, img, label):

        aug_data = []

        aug_label = []

        for cls4aug in range(4):

            cls_idx = np.where(label == cls4aug + 1)

            tmp_data = img[cls_idx]

            tmp_label = label[cls_idx]

            tmp_aug_data = np.zeros(tmp_data.shape)

            for ri in range(tmp_data.shape[0]):

                for rj in range(8):

                    rand_idx = np.random.randint(0, tmp_data.shape[0], 8)

                    tmp_aug_data[ri, :, :, rj * 125:(rj + 1) * 125] = tmp_data[rand_idx[rj], :, :, rj * 125:(rj + 1) * 125]

            aug_data.append(tmp_aug_data)

            aug_label.append(tmp_label)

        aug_data = np.concatenate(aug_data)

        aug_label = np.concatenate(aug_label)

        aug_shuffle = np.random.permutation(len(aug_data))

        aug_data = aug_data[aug_shuffle, :, :]

        aug_label = aug_label[aug_shuffle]

        return aug_data, aug_label

    def update_centers(self, feature, label):

            deltac = {}

            count = {}

            count[0] = 0

            for i in range(len(label)):

                l = label[i]

                if l in deltac:

                    deltac[l] += self.centers[l]-feature[i]

                else:

                    deltac[l] = self.centers[l]-feature[i]

                if l in count:

                    count[l] += 1

                else:

                    count[l] = 1

            for ke in deltac.keys():

                deltac[ke] = deltac[ke]/(count[ke]+1)

            return deltac

    def train(self):

        img, label, test_data, test_label = self.get_source_data()

        img = torch.from_numpy(img)

        label = torch.from_numpy(label - 1)

        dataset = torch.utils.data.TensorDataset(img, label)

        self.dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=self.batch_size, shuffle=True)

        test_data = torch.from_numpy(test_data)

        test_label = torch.from_numpy(test_label - 1)

        test_dataset = torch.utils.data.TensorDataset(test_data, test_label)

        self.test_dataloader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=self.batch_size, shuffle=True)

        for i in range(self.c_dim):

            self.centers[i] = torch.randn(self.dimension)

            self.centers[i] = self.centers[i].cuda()

        # Optimizers

        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, betas=(self.b1, self.b2))

        test_data = Variable(test_data.type(self.Tensor))

        test_label = Variable(test_label.type(self.LongTensor))

        bestAcc = 0

        averAcc = 0

        num = 0

        Y_true = 0

        Y_pred = 0

        # Train the cnn model

        total_step = len(self.dataloader)

        curr_lr = self.lr

        for e in range(self.n_epochs):

            in_epoch = time.time()

            self.model.train()

            for i, (img, label) in enumerate(self.dataloader):

                img = Variable(img.cuda().type(self.Tensor))

                # img = self.active_function(img)

                label = Variable(label.cuda().type(self.LongTensor))

                aug_data, aug_label = self.interaug(self.allData, self.allLabel)

                img = torch.cat((img, aug_data))

                label = torch.cat((label, aug_label))

                tok, outputs = self.model(img)

                loss = self.criterion_cls(outputs, label)

                self.optimizer.zero_grad()

                loss.backward()

                self.optimizer.step()

            out_epoch = time.time()

            if (e + 1) % 1 == 0:

                self.model.eval()

                Tok, Cls = self.model(test_data)

                loss_test = self.criterion_cls(Cls, test_label)

                y_pred = torch.max(Cls, 1)[1]

                acc = float((y_pred == test_label).cpu().numpy().astype(int).sum()) / float(test_label.size(0))

                train_pred = torch.max(outputs, 1)[1]

                train_acc = float((train_pred == label).cpu().numpy().astype(int).sum()) / float(label.size(0))

                print('Epoch:', e,

                      '  Train loss: %.6f' % loss.detach().cpu().numpy(),

                      '  Test loss: %.6f' % loss_test.detach().cpu().numpy(),

                      '  Train accuracy %.6f' % train_acc,

                      '  Test accuracy is %.6f' % acc)

                self.log_write.write(str(e) + "    " + str(acc) + "\n")

                num = num + 1

                averAcc = averAcc + acc

                if acc > bestAcc:

                    bestAcc = acc

                    Y_true = test_label

                    Y_pred = y_pred

        torch.save(self.model.module.state_dict(), 'model.pth')

        averAcc = averAcc / num

        print('The average accuracy is:', averAcc)

        print('The best accuracy is:', bestAcc)

        self.log_write.write('The average accuracy is: ' + str(averAcc) + "\n")

        self.log_write.write('The best accuracy is: ' + str(bestAcc) + "\n")

        return bestAcc, averAcc, Y_true, Y_pred

def main():

    best = 0

    aver = 0

    result_write = open("/Code/CT/results/cf/2b/sub_result.txt", "w")

    for i in range(9):

        starttime = datetime.datetime.now()

        seed_n = np.random.randint(2021)

        print('seed is ' + str(seed_n))

        random.seed(seed_n)

        np.random.seed(seed_n)

        torch.manual_seed(seed_n)

        torch.cuda.manual_seed(seed_n)

        torch.cuda.manual_seed_all(seed_n)

        print('Subject %d' % (i+1))

        exgan = ExGAN(i + 1)

        bestAcc, averAcc, Y_true, Y_pred = exgan.train()

        print('THE BEST ACCURACY IS ' + str(bestAcc))

        result_write.write('Subject ' + str(i + 1) + ' : ' + 'Seed is: ' + str(seed_n) + "\n")

        result_write.write('**Subject ' + str(i + 1) + ' : ' + 'The best accuracy is: ' + str(bestAcc) + "\n")

        result_write.write('Subject ' + str(i + 1) + ' : ' + 'The average accuracy is: ' + str(averAcc) + "\n")

        endtime = datetime.datetime.now()

        print('subject %d duration: '%(i+1) + str(endtime - starttime))

        best = best + bestAcc

        aver = aver + averAcc

        if i == 0:

            yt = Y_true

            yp = Y_pred

        else:

            yt = torch.cat((yt, Y_true))

            yp = torch.cat((yp, Y_pred))

    best = best / 9

    aver = aver / 9

    result_write.write('**The average Best accuracy is: ' + str(best) + "\n")

    result_write.write('The average Aver accuracy is: ' + str(aver) + "\n")

    result_write.close()

if __name__ == "__main__":

    print(time.asctime(time.localtime(time.time())))

    main()

    print(time.asctime(time.localtime(time.time())))