import itertools
import random
from functools import partial

import fire
import pandas as pd
import seaborn as sns
import torchvision.transforms as transforms
from matplotlib import ticker
from scipy.io import savemat
from scipy.ndimage import zoom
from scipy.stats import sem
from sklearn import manifold
from sklearn.manifold import MDS
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchsummary import summary
from torchvision import datasets
from torchvision.utils import save_image, make_grid

from models import *
from utils import *


class Main(object):
    def __init__(self, img_size=225, cuda=torch.cuda.is_available(), images_folder='./data',
                 models_folder='./models', model='prior', test_image='156imagesC', latent_dim=100,
                 samples_folder='./samples', logs_folder='./logs', n_epochs=2000, batch_size=1024, n_model=None,
                 intermediate_size=512,
                 sample_interval=50, n_cpu=32, n_classes=4, train_data='train_Dataset.lmdb', dist_prior=False,
                 decoder_prior=False, deconv=True, init=True, m=7, random=False):

        super(Main, self).__init__()
        self.init = init
        self.deconv = deconv
        self.dist_prior = dist_prior
        self.decoder_prior = decoder_prior
        self.n_cpu = n_cpu
        self.intermediate_size = intermediate_size
        self.sample_interval = sample_interval
        self.img_transform = transforms.Compose([
            transforms.Resize(img_size),
            transforms.CenterCrop(img_size),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ])
        self.n_classes = n_classes
        self.cuda = cuda
        self.m = m
        self.Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
        self.batch_size = batch_size * torch.cuda.device_count()
        self.n_epochs = n_epochs
        self.main_image_folder = images_folder
        self.train_images = os.path.join(images_folder, train_data)
        self.images_folder = os.path.join(images_folder, 'test', test_image)
        self.meg_folder = os.path.join(images_folder, 'MEG_RDMs')
        self.main_model_folder = models_folder
        self.models_folder = os.path.join(models_folder, model)
        self.logs_folder = os.path.join(logs_folder, model)
        os.makedirs(self.logs_folder, exist_ok=True)
        self.main_samples_folder = samples_folder
        self.samples_folder = os.path.join(samples_folder, model)
        self.model = model
        self.latent_dim = latent_dim
        self.define_models()
        if cuda:
            self.encoder = nn.DataParallel(self.encoder)
            self.decoder = nn.DataParallel(self.decoder)
            self.discriminator = nn.DataParallel(self.discriminator)
            self.encoder.cuda()
            self.decoder.cuda()
            self.discriminator.cuda()
            summary(self.encoder, input_size=(3, img_size, img_size))
            if decoder_prior:
                summary(self.decoder, input_size=(1, latent_dim + n_classes))
            else:
                summary(self.decoder, input_size=(1, latent_dim))
            if dist_prior:
                summary(self.discriminator, input_size=(1, latent_dim + n_classes))
            else:
                summary(self.discriminator, input_size=(1, latent_dim))
        if not random:
            self.load_model(n_model=n_model, train=False)
        self.writer = SummaryWriter(log_dir=self.logs_folder)

        print(f"tensorboard --logdir={self.logs_folder}")

    def define_models(self):
        if 'aae' in self.model:
            self.encoder = Encoder(self.Tensor, latent_dim=self.latent_dim,
                                   intermediate_size=self.intermediate_size, init=self.init, reparm=True)
        elif 'vae' in self.model:
            self.encoder = Encoder(self.Tensor, latent_dim=self.latent_dim,
                                   intermediate_size=self.intermediate_size, init=self.init, reparm=True,
                                   return_mu=True)
        else:
            self.encoder = Encoder(self.Tensor, latent_dim=self.latent_dim,
                                   intermediate_size=self.intermediate_size, init=self.init, reparm=False)
        self.decoder = Decoder(latent_dim=self.latent_dim, prior=self.decoder_prior, deconv=self.deconv,
                               init=self.init,
                               m=self.m)
        self.discriminator = Discriminator(latent_dim=self.latent_dim, n_classes=self.n_classes,
                                           prior=self.dist_prior,
                                           init=self.init)

    def read_data(self):
        assert os.path.exists(self.logs_folder)
        event_acc = EventAccumulator(self.logs_folder)
        event_acc.Reload()
        print(event_acc.Tags())
        print(event_acc.Scalars('Loss/g_loss'))

    def gen_rdm_plots(self, scatter_size=200, extension='svg'):
        assert os.path.exists(os.path.join(self.samples_folder, 'rdms'))
        save_path = os.path.join(self.samples_folder, 'rdm_plots')
        os.makedirs(save_path, exist_ok=True)
        self.delete_dir_content(save_path)
        for i in os.listdir(os.path.join(self.samples_folder, 'rdms')):
            rdm = normalize(get_RDMs(os.path.join(self.samples_folder, 'rdms', i)))
            plt.figure()
            plt.imshow(rdm)
            plt.axis('off')
            plt.tight_layout()
            plt.savefig(os.path.join(save_path, f'RDM_{i.replace(".mat", "")}.{extension}'), bbox_inches='tight',
                        pad_inches=0)
            plt.close()
            plt.figure()
            Y = MDS(n_components=2, n_init=100, max_iter=1000, dissimilarity='precomputed', eps=1e-6)
            Y = Y.fit_transform(rdm)
            plt.figure()
            plt.axis('off')
            plt.scatter(Y[0:32, 0], Y[0:32, 1], c='red', s=scatter_size)
            plt.scatter(Y[32:84, 0], Y[32:84, 1], c='orange', s=scatter_size)
            plt.scatter(Y[84:120, 0], Y[84:120, 1], c='blue', s=scatter_size)
            plt.scatter(Y[120:157, 0], Y[120:157, 1], c='green', s=scatter_size)
            plt.tight_layout()
            plt.savefig(os.path.join(save_path, f'MDS_{i.replace(".mat", "")}.{extension}'), bbox_inches='tight',
                        pad_inches=0)
            plt.close()

    def gen_rdm(self):
        assert os.path.exists(self.images_folder)
        save_path = os.path.join(self.samples_folder, 'rdms')
        os.makedirs(save_path, exist_ok=True)
        print(save_path)
        self.delete_dir_content(save_path)
        dataloader = self.load_data(self.images_folder, 156)

        def hook_fn(m, i, o, e=True):
            o = o.flatten(1).cpu().detach().numpy()
            if e:
                if m in etem_act:
                    etem_act[m] = np.vstack((etem_act[m], o))
                else:
                    etem_act[m] = o
            else:
                if m in dtem_act:
                    dtem_act[m] = np.vstack((dtem_act[m], o))
                else:
                    dtem_act[m] = o

        def get_all_layers(net, par=True):
            for layer in net.modules():
                if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear) or isinstance(layer, nn.ConvTranspose2d):
                    print(layer)
                    layer.register_forward_hook(partial(hook_fn, e=par))

        get_all_layers(self.encoder)
        get_all_layers(self.decoder, False)

        for idx, (img, labels) in enumerate(dataloader):
            etem_act = {}
            dtem_act = {}
            img = img.cuda()
            de_img = self.encoder(self.Tensor(img))
            # print(len(de_img))
            if len(de_img) == 3:
                _ = self.decoder(de_img[0])
                # print(de_img[0].flatten(1).cpu().detach().numpy().shape)
                etem_act['z'] = de_img[0].flatten(1).cpu().detach().numpy()

            else:
                _ = self.decoder(de_img)
                # print(de_img.flatten(1).cpu().detach().numpy().shape)
                etem_act['z'] = de_img.flatten(1).cpu().detach().numpy()

        for i, activations in tqdm(enumerate(etem_act.values())):
            savemat(os.path.join(save_path, f'E_L{i + 1}.mat'),
                    {'RDM': computeRDM(activations.squeeze())})
        for i, activations in tqdm(enumerate(dtem_act.values())):
            savemat(os.path.join(save_path, f'D_L{i + 1}.mat'),
                    {'RDM': computeRDM(activations.squeeze())})

    def gen_compare_rdms(self):
        assert os.path.exists(os.path.join(self.samples_folder, 'rdms'))
        save_path = os.path.join(self.samples_folder, 'compare_rdms')
        os.makedirs(save_path, exist_ok=True)
        self.delete_dir_content(save_path)
        files = os.listdir(os.path.join(self.samples_folder, 'rdms'))
        files.sort()
        files = [os.path.join(self.samples_folder, 'rdms', x) for x in files]
        RDMs = list(map(get_RDMs, files))
        x_labels = [file.split('/')[-1].replace('.mat', '').replace('_L', '') for file in files]
        res = np.zeros((len(files), len(files)))
        for i in range(len(files)):
            for j in range(i + 1, len(files)):
                tres = 1 - np.corrcoef(lower_tra(RDMs[i]), lower_tra(RDMs[j]))[0, 1]
                res[i, j] = tres
                res[j, i] = tres

        fig = plt.figure(figsize=[10, 10])
        ax = fig.add_subplot(111)
        cax = ax.matshow(res / np.max(res), interpolation='nearest')
        fig.colorbar(cax)

        files = list(map(lambda x: x.split("\\")[-1].replace('.mat', '').split("/")[-1], files))
        ax.set_xticklabels([''] + x_labels)
        ax.set_yticklabels([''] + x_labels)
        ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
        ax.yaxis.set_major_locator(ticker.MultipleLocator(1))
        plt.setp(ax.get_xticklabels(), rotation=-45, ha="right",
                 rotation_mode="anchor")
        plt.savefig(os.path.join(save_path, f'{self.model}.png'))
        savemat(os.path.join(save_path, f'{self.model}.mat'), {'Matrix': res})

    def plot_meg_compare(self):
        assert os.path.exists(self.meg_folder)
        assert os.path.exists(os.path.join(self.samples_folder, 'rdms'))
        self.encoder.eval()
        self.decoder.eval()
        save_path = os.path.join(self.samples_folder, 'rdm_meg')
        os.makedirs(save_path, exist_ok=True)
        # self.delete_dir_content(save_path)
        rdms_folder = os.path.join(self.samples_folder, 'rdms')
        E_correlations = []
        D_correlations = []
        files = [os.path.join(rdms_folder, i) for i in os.listdir(rdms_folder)]
        E_files = [file for file in files if 'E_' in file]
        D_files = [file for file in files if 'D_' in file]
        if os.path.exists(os.path.join(save_path, 'E.pickle')):
            E_correlations = pickle.load(open(os.path.join(save_path, 'E.pickle'), 'rb'))
        else:
            for idx, layer_rdm in tqdm(enumerate(E_files)):
                E_correlations.append(compare_meg_rdms(self.meg_folder, get_RDMs(layer_rdm)))
            pickle.dump(E_correlations, open(os.path.join(save_path, 'E.pickle'), "wb"))
        plot_correlations_time(E_correlations, save_path)
        if os.path.exists(os.path.join(save_path, 'D.pickle')):
            D_correlations = pickle.load(open(os.path.join(save_path, 'D.pickle'), 'rb'))
        else:
            for idx, layer_rdm in tqdm(enumerate(D_files)):
                D_correlations.append(compare_meg_rdms(self.meg_folder, get_RDMs(layer_rdm)))
            pickle.dump(D_correlations, open(os.path.join(save_path, 'D.pickle'), "wb"))
        plot_correlations_time(D_correlations, save_path, 'D')

    def gen_images(self, n_row=8):
        assert os.path.exists(self.images_folder)
        save_path = os.path.join(self.samples_folder, 'image_reconstruction')
        os.makedirs(save_path, exist_ok=True)
        self.delete_dir_content(save_path)
        dataloader = self.load_data(self.images_folder, batch_size=1)
        for idx, (img, _) in enumerate(dataloader):
            decode = self.encoder(img)
            gen_imgs = self.decoder(decode)
            break
        save_image(gen_imgs, os.path.join(save_path, 'f.png'), nrow=n_row, normalize=True)
        save_image(img, os.path.join(save_path, 'r.png'), nrow=n_row, normalize=True)

    def load_data(self, dir, batch_size=None, shuffle=False, lmdb=False):
        assert os.path.exists(dir)
        if not lmdb:
            imageset = datasets.ImageFolder(root=dir, transform=self.img_transform)
        else:
            imageset = ImageFolderLMDB(dir)
        return torch.utils.data.DataLoader(imageset,
                                           batch_size=batch_size if batch_size else self.batch_size,
                                           shuffle=shuffle,
                                           num_workers=self.n_cpu)

    def load_model(self, train=False, n_model=None):
        if train:
            self.discriminator.train()
            self.decoder.train()
            self.encoder.train()
        else:
            self.discriminator.eval()
            self.decoder.eval()
            self.encoder.eval()
        if os.path.exists(self.models_folder) and len(os.listdir(self.models_folder)) > 0:
            files = os.listdir(self.models_folder)
            files.sort(key=natural_keys)
            print(f'Loading Model {files[-1]}....')
            if n_model:
                checkpoint = torch.load(os.path.join(self.models_folder, files[n_model]),
                                        map_location='cuda' if self.cuda else 'cpu')
            else:
                checkpoint = torch.load(os.path.join(self.models_folder, files[-1]),
                                        map_location='cuda' if self.cuda else 'cpu')

            self.discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
            self.decoder.load_state_dict(checkpoint['decoder_state_dict'], strict=True)
            self.encoder.load_state_dict(checkpoint['encoder_state_dict'], strict=True)

            return int(files[-1])
        print('Not loading a model')
        return 0

    def change_range(self, x):
        return (x + 1) / 2

    def get_latent_disc(self, n_images=15, p=0.00009):
        results = [[] for _ in range(self.n_classes)]
        for j in tqdm(range(self.n_classes)):
            while len(results[j]) < n_images:
                z = Variable(self.Tensor(np.random.normal(0, 1, (n_images, self.latent_dim))))
                labels = nn.functional.one_hot(
                    Variable(self.Tensor(np.array([np.full(n_images, j)]).flatten())).to(torch.int64), self.n_classes)
                zn = torch.cat((labels.float(), z), 1)
                r = self.discriminator(zn)
                for i in range(r.shape[0]):
                    if r[i] <= p:
                        results[j].append(self.decoder(z[i, :].view(1, -1))[0].squeeze().cpu().detach().numpy())
        return results

    def gen_raw_images(self, n_row=5):
        save_path = os.path.join(self.samples_folder, 'gen_images')
        self.delete_dir_content(save_path)
        os.makedirs(save_path, exist_ok=True)
        images = self.get_latent_disc()
        for i, imgs in tqdm(enumerate(images)):
            save_image(torch.tensor(imgs).data, os.path.join(save_path, f'{i}.png'), nrow=n_row, normalize=True)

    def sample_images(self, step, n_row=5, recon=None):
        save_path = os.path.join(self.samples_folder, 'sampled_images')
        os.makedirs(save_path, exist_ok=True)
        z = Variable(self.Tensor(np.random.normal(0, 1, (n_row ** 2, self.latent_dim))))
        if self.decoder_prior:
            abels = Variable(
                self.Tensor(np.random.randint(0, self.n_classes, (n_row ** 2))).to(torch.int64))
            z = torch.cat((nn.functional.one_hot(abels, self.n_classes).float().cuda(), z), 1)
        gen_imgs = self.decoder(z)
        # self.save_images(gen_imgs.data, os.path.join(save_path, image_name))
        self.writer.add_image('generated', make_grid(self.change_range(gen_imgs)), step)
        if recon:
            self.writer.add_image('real_images', make_grid(self.change_range(recon[0][:20])), step)
            self.writer.add_image('reconstructed_images', make_grid(self.change_range(recon[1][:20])), step)

    def delete_dir_content(self, dir):
        if os.path.exists(dir):
            for f in os.listdir(dir):
                os.remove(os.path.join(dir, f))

    def reset_grad(self):
        self.encoder.zero_grad()
        self.discriminator.zero_grad()
        self.decoder.zero_grad()

    def train_vae(self, lr=0.0002, save_model=10):
        optimizer_G = torch.optim.Adam(
            itertools.chain(self.encoder.parameters(), self.decoder.parameters()), lr=lr
        )
        Tensor = self.Tensor
        loss_log = tqdm(total=0, position=0, bar_format='{desc}')
        print("Loading data ... ")
        dataloader = self.load_data(self.train_images, shuffle=True, lmdb=True)
        offset = self.load_model(train=True)
        print("Starting the training process ... ")
        for epoch in range(self.n_epochs):
            for i, (imgs, labels) in enumerate(dataloader):
                real_imgs = Variable(imgs.type(Tensor).squeeze())

                # -----------------
                #  Train Generator
                # -----------------
                optimizer_G.zero_grad()

                encoded_imgs, mu, log_var = self.encoder(real_imgs)
                decoded_imgs = self.decoder(encoded_imgs)

                # Loss measures generator's ability to fool the discriminator
                g_loss = F.mse_loss(decoded_imgs, real_imgs) + torch.mean(
                    -0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim=1), dim=0)

                g_loss.backward()
                optimizer_G.step()

                loss_log.set_description_str(
                    '[Epoch %d/%d] [Batch %d/%d] [G loss: %f]' % (
                        epoch + offset, self.n_epochs, i, len(dataloader), g_loss.item())
                )
                batches_done = (epoch + offset) * len(dataloader) + i
                self.writer.add_scalar('Loss/g_loss', g_loss.item(), batches_done)
                if batches_done % self.sample_interval == 0:
                    self.sample_images(step=batches_done, recon=(real_imgs, decoded_imgs))
            if (epoch + offset + 1) % save_model == 0:
                self.save_model(epoch + offset + 1)

    def train_aae(self, lr=0.0002, save_model=10):
        adversarial_loss = torch.nn.BCELoss()
        pixelwise_loss = torch.nn.MSELoss()
        optimizer_G = torch.optim.Adam(
            itertools.chain(self.encoder.parameters(), self.decoder.parameters()), lr=lr
        )
        optimizer_D = torch.optim.Adam(self.discriminator.parameters(), lr=lr)
        Tensor = self.Tensor
        loss_log = tqdm(total=0, position=0, bar_format='{desc}')
        print("Loading data ... ")
        dataloader = self.load_data(self.train_images, shuffle=True, lmdb=True)
        offset = self.load_model(train=True)
        print("Starting the training process ... ")
        for epoch in range(self.n_epochs):
            for i, (imgs, labels) in enumerate(dataloader):
                valid = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)
                fake = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)
                real_imgs = Variable(imgs.type(Tensor).squeeze(), requires_grad=False)
                real_labels = labels.type(torch.cuda.LongTensor).squeeze()

                # -----------------
                #  Train Generator
                # -----------------
                optimizer_G.zero_grad()

                encoded_imgs = self.encoder(real_imgs)
                decoded_imgs = self.decoder(encoded_imgs)

                if self.dist_prior:
                    encoded_imgs = torch.cat(
                        (nn.functional.one_hot(real_labels, self.n_classes).float(), encoded_imgs), 1)

                encoded_imgs = Variable(encoded_imgs, requires_grad=True)
                # print(torch.argmax(encoded_imgs), encoded_imgs.max())
                # Loss measures generator's ability to fool the discriminator
                g_loss = 0.001 * adversarial_loss(self.discriminator(encoded_imgs), valid) + 0.999 * pixelwise_loss(
                    decoded_imgs, real_imgs
                )

                g_loss.backward()
                optimizer_G.step()

                # ---------------------
                #  Train Discriminator
                # ---------------------

                optimizer_D.zero_grad()

                # Sample noise as discriminator ground truth
                z = self.Tensor(imgs.shape[0], self.latent_dim).normal_()

                if self.dist_prior:
                    z = torch.cat(
                        (nn.functional.one_hot(torch.randint(0, self.n_classes, size=[imgs.shape[0]], device='cuda'),
                                               self.n_classes).float(), z), 1)
                z = Variable(z, requires_grad=True)
                # Measure discriminator's ability to classify real from generated samples
                real_loss = adversarial_loss(self.discriminator(z), valid)
                fake_loss = adversarial_loss(self.discriminator(encoded_imgs.detach()), fake)
                d_loss = 0.5 * (real_loss + fake_loss)

                d_loss.backward()
                optimizer_D.step()

                loss_log.set_description_str(
                    '[Epoch %d/%d] [Batch %d/%d] [G loss: %f] [D loss: %f]' % (
                        epoch + offset, self.n_epochs, i, len(dataloader), g_loss.item(), d_loss.item())
                )
                batches_done = (epoch + offset) * len(dataloader) + i
                self.writer.add_scalar('Loss/g_loss', g_loss.item(), batches_done)
                self.writer.add_scalar('Loss/d_loss', d_loss.item(), batches_done)
                if batches_done % self.sample_interval == 0:
                    self.sample_images(step=batches_done, recon=(real_imgs, decoded_imgs))
            if (epoch + offset + 1) % save_model == 0:
                self.save_model(epoch + offset + 1)

    def save_images(self, data, image_name, n_row=8):
        save_image(data, image_name, nrow=n_row,
                   normalize=True)

    def save_model(self, epoch):
        os.makedirs(self.models_folder, exist_ok=True)
        torch.save({
            'epoch': epoch,
            'discriminator_state_dict': self.discriminator.state_dict(),
            'decoder_state_dict': self.decoder.state_dict(),
            'encoder_state_dict': self.encoder.state_dict(),
        }, os.path.join(self.models_folder, f"{epoch}"))


if __name__ == '__main__':
    fire.Fire(Main)
