# pip install pytorch-lightning
# pip install neptune-client
#%%
from __future__ import print_function
from utils.hessian_penalty.hessian_penalty_pytorch import hessian_penalty
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks.progress import ProgressBar
from sklearn.model_selection import train_test_split
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from collections import defaultdict
import torch
import torch.utils.data
from torch import nn, optim
from torch.nn import functional as F
import math
import pytorch_lightning as pl
# import utils.plot_utils as utils
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pickle
import wandb
import time
import os
from utils.dsprites.datasets import get_dataloaders
from utils import run_utils, data_utils, utils, metric_utils, settings, disentangle_utils, latent_space_utils  # ,morpho_utils
import utils.morphomnist.io as morpho_io
#ToDo EDA:
# - Long Tail graphics
# - Remove user who had less than a threshold of seen items
# - Create Markdown with EDA results

#ToDo input_params:
# Parameter that should be tweakable by invoking the routine:
# - epochs
# - learning_rate
# - batch_size
# - simplified_rating

#ToDo metrics:
# Add https://towardsdatascience.com/evaluation-metrics-for-recommender-systems-df56c6611093


from torchvision import datasets, transforms

seed = 42
torch.manual_seed(seed)

class VAE(pl.LightningModule):
    def __init__(self, conf:dict, *args, **kwargs):
        super().__init__()

        # self.kwargs = kwargs
        self.save_hyperparameters(conf)
        self.ls_predicted_movies = []
        self.is_hessian_penalty_activated = self.hparams["is_hessian_penalty_activated"]
        self.used_data = self.hparams["used_data"]
        self.expanded_user_item = self.hparams["expanded_user_item"]
        self.generative_factors = self.hparams["generative_factors"]
        self.mixup = self.hparams["mixup"]
        self.np_synthetic_data = self.hparams["synthetic_data"]
        self.ls_syn_y = self.hparams["syn_y"]
        self.experiment_path_train = conf["experiment_path"]
        self.experiment_path_test = self.experiment_path_train
        self.beta = self.hparams["beta"]
        self.avg_mce = 0.0
        self.train_dataset = None
        self.test_dataset = None
        self.test_size = self.hparams["test_size"]
        # self.test_size = 0.005
        self.no_latent_factors = self.hparams["latent_dim"]
        self.max_unique_movies = 0
        self.unique_movies =0
        self.np_user_item = None
        self.small_dataset = self.hparams["small_dataset"]
        self.simplified_rating = self.hparams["simplified_rating"]
        self.max_epochs = self.hparams["max_epochs"]
        self.dct_index2itemId = None
        self.test_y_bin = None
        self.df_movies_z_combined =None

        # if(self.np_synthetic_data is None):
        #     self.load_dataset() #additionaly assigns self.unique_movies and self.np_user_item
        #     self.df_movies = pd.read_csv('../data/generated/df_movies_cleaned3.csv')
        #     self.dct_attribute_distribution = utils.load_json_as_dict(
        #         'attribute_distribution.json')  # load relative frequency distributioon from dictionary (pickle it)
        #
        # else:
        #     self.train_dataset, self.test_dataset = train_test_split(self.np_synthetic_data, test_size=self.test_size, random_state=42)
        #     self.train_y, self.test_y = train_test_split(self.ls_syn_y, test_size=self.test_size, random_state=42)
        #     self.test_y_bin = np.asarray(pd.get_dummies(pd.DataFrame(data=self.test_y)))
        #     self.unique_movies = self.np_synthetic_data.shape[1]
        #     self.df_movies = pd.read_csv('../data/generated/syn.csv')
        #     self.dct_attribute_distribution = utils.load_json_as_dict(
        #         'syn_attribute_distribution.json')  # load relative frequency distributioon from dictionary (pickle it)

        self.bs = 100

        #nn.Linear layer creates a linear function (¸x + b), with its parameters initialized
        self.input_dimension = int(self.unique_movies *math.pow(4, self.generative_factors)) if self.expanded_user_item == True else self.unique_movies

        if(self.used_data=='morpho'):
            self.input_dimension = 28*28
        elif(self.used_data=='dsprites'):
            self.input_dimension = 64 * 64

        self.fc11 = nn.Linear(in_features=self.input_dimension, out_features=1200) #input
        self.fc12 = nn.Linear(in_features=1200, out_features=1200) #input
        self.fc13 = nn.Linear(in_features=1200, out_features=1200) #input
        self.encoder = nn.Sequential(self.fc11, nn.ReLU(),
                                     self.fc12, nn.ReLU(),
                                    self.fc13#, nn.ReLU(),
                                     )

        self.fc21 = nn.Linear(in_features=1200, out_features=self.no_latent_factors) #encoder mean
        self.fc22 = nn.Linear(in_features=1200, out_features=self.no_latent_factors) #encoder variance

        self.fc31 = nn.Linear(in_features=self.no_latent_factors, out_features=1200)
        self.fc32 = nn.Linear(in_features=1200, out_features=1200)
        self.fc33 = nn.Linear(in_features=1200, out_features=1200)
        self.fc331 = nn.Linear(in_features=1200, out_features=1200)

        self.fc34 = nn.Linear(in_features=1200, out_features=self.input_dimension)
        self.decoder = nn.Sequential(self.fc31, nn.Tanh(),
                                     self.fc32, nn.Tanh(),
                                     self.fc33, nn.Tanh(),
                                     self.fc331, nn.Tanh(),
                                     self.fc34)

        self.KLD = None
        self.ls_kld = []
        self.dis_KLD = None

        self.z = None
        self.kld_matrix = np.empty((0, self.no_latent_factors))
        self.np_z_test = np.empty((0, self.no_latent_factors))#self.test_dataset.shape[0]
        self.np_mu_test = np.empty((0, self.no_latent_factors))
        self.np_logvar_test = np.empty((0, self.no_latent_factors))

        self.np_z_train = np.empty((0, self.no_latent_factors))  # self.test_dataset.shape[0]
        self.np_mu_train = np.empty((0, self.no_latent_factors))
        self.np_logvar_train = np.empty((0, self.no_latent_factors))

        self.sigmoid_annealing_threshold = self.hparams['sigmoid_annealing_threshold']
        self.mce_batch_train = None
        self.mce_batch_test = None

        self.z_mean_train = []
        self.z_min_train = []
        self.z_max_train = []

        # Initialize weights
        self.encoder.apply(run_utils.weight_init)
        self.decoder.apply(run_utils.weight_init)

        self.batch_size =512
        # np_user_item, ls_y = sklearn.utils.shuffle(np_user_item, ls_y)
        if (self.used_data == 'dsprites'):
            self.batch_size =4096
            loader = get_dataloaders('dsprites', batch_size=512, shuffle=False)
            dsprites_data = loader.dataset.imgs#[:5000]
            self.dsprites_lat_names = loader.dataset.lat_names
            dsprites_gen_fac_values = loader.dataset.lat_values#[:5000]


            # self.train_dataset, self.test_dataset = train_test_split(dsprites_data, test_size=self.test_size, shuffle=False,random_state=42)
            #dsprites_data[:int(dsprites_data.shape[0]*0.1)]
            self.train_dataset, self.test_dataset = train_test_split(dsprites_data,
                                                                     test_size=self.test_size, shuffle=True,random_state=42)
            self.train_y, self.test_y = train_test_split(dsprites_gen_fac_values,
                                                         test_size=self.test_size, shuffle=True,random_state=42)

    def encode(self, x):
        h1 = F.relu(self.encoder(x))
        return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z):
        # return torch.bernoulli(self.decoder(z))
        return torch.bernoulli(torch.sigmoid(self.decoder(z)))

    def compute_z(self, mu, logvar):
        z = self.reparameterize(mu, logvar)
        return z

    def sample(self, mu, log_var):
        std = torch.exp(log_var / 2)
        p = torch.distributions.Normal(torch.zeros_like(mu), torch.ones_like(std))
        q = torch.distributions.Normal(mu, std)
        z = q.rsample()
        return p, q, z

    def forward(self, x, **kwargs):
        if(kwargs):
            z = kwargs['z']
            mu = kwargs['mu']
            logvar = kwargs['logvar']
            p = None
            q = None
        else:
            # print(x.view(-1, self.unique_movies)[0])
            # print(x[0])
            mu, logvar = self.encode(x) #40960/512 (Batchsize) results in 512,80
            # z = self.compute_z(mu, logvar)
            p, q, z = self.sample(mu, logvar)
            self.z = z

        return self.decode(z), mu, logvar, p, q

    def train_dataloader(self):

        # MNIST Dataset


        if(self.used_data == 'morpho'):
            # MORPHO_MNIST_FILE_TRAIN_Y = "/Users/d069735/workspace/Study/decoding-latent-space-rs/data/morpho-mnist/global/train-pert-idx1-ubyte.gz"
            MORPHO_MNIST_FILE_TRAIN_Y = "/models/_output_dir/pm-pert-idx1-ubyte.gz"
            MORPHO_MNIST_FILE_TRAIN_X = "/models/_output_dir/pm-images-idx3-ubyte.gz"
            self.train_dataset = morpho_io.load_idx(MORPHO_MNIST_FILE_TRAIN_X)[:59000]
            self.train_y = morpho_io.load_idx(MORPHO_MNIST_FILE_TRAIN_Y)[:59000]
        elif(self.used_data == 'dsprites'):
            # train_loader = get_dataloaders('dsprites', batch_size=512, shuffle=False)
            train_loader = torch.utils.data.DataLoader(dataset=self.train_dataset, num_workers=0,batch_size=self.batch_size, shuffle=True)
            return train_loader
        #regular mnist
        else:
            self.train_dataset = datasets.MNIST(root='../data/mnist_data/', train=True, transform=transforms.ToTensor(), download=True)
            self.train_y = self.train_dataset.targets.tolist()

        train_loader = torch.utils.data.DataLoader(dataset=self.train_dataset, batch_size=self.batch_size, shuffle=False)

        return train_loader

    def test_dataloader(self):
        #bar = datasets.MNIST(root='../data/mnist_data/', train=False, transform=transforms.ToTensor(),download=False)

        if(self.used_data == 'morpho'):
            # MORPHO_MNIST_FILE_TEST_Y = "/Users/d069735/workspace/Study/decoding-latent-space-rs/data/morpho-mnist/global/t10k-pert-idx1-ubyte.gz"
            MORPHO_MNIST_FILE_TEST_Y = "/models/_output_dir/pm-pert-idx1-ubyte.gz"
            # MORPHO_MNIST_FILE_TEST_X = "/Users/d069735/workspace/Study/decoding-latent-space-rs/data/morpho-mnist/global/t10k-images-idx3-ubyte.gz"
            MORPHO_MNIST_FILE_TEST_X = "/models/_output_dir/pm-images-idx3-ubyte.gz"
            self.test_dataset = morpho_io.load_idx(MORPHO_MNIST_FILE_TEST_X)[59000:60000]
            self.test_y = morpho_io.load_idx(MORPHO_MNIST_FILE_TEST_Y)[59000:60000]
        elif(self.used_data == 'dsprites'):
            # self.test_dataset =
            test_loader = torch.utils.data.DataLoader(dataset=self.test_dataset, batch_size=int(self.batch_size/2), shuffle=False)
            return test_loader
        #regular mnist

        else:
            self.test_y = self.test_dataset.targets.tolist()

        test_loader = torch.utils.data.DataLoader(dataset=self.test_dataset, batch_size=self.bs, shuffle=False)

        return test_loader

    def configure_optimizers(self):
        optimizer = optim.Adagrad(self.parameters(), lr=1e-2) #Adam
        # criterion = nn.MSELoss()  # mean-squared error loss
        # scheduler = StepLR(optimizer, step_size=1)
        return optimizer#, scheduler

    def collect_z_values(self, ts_mu_chunk, ts_logvar_chunk):#, ls_y
        start = time.time()
        ls_grad_z = self.compute_z(ts_mu_chunk, ts_logvar_chunk)
        self.np_z_train = np.append(self.np_z_train, np.asarray(ls_grad_z.tolist()),
                                    axis=0)  # TODO Describe in thesis that I get back a grad object instead of a pure tensor as it is in the test method since we are in the training method.
        self.np_mu_train = np.append(self.np_mu_train, np.asarray(ts_mu_chunk.tolist()), axis=0)
        self.np_logvar_train = np.append(self.np_logvar_train, np.asarray(ts_logvar_chunk.tolist()), axis=0)

        # print('Shape np_z_train: {}'.format(self.np_z_train.shape))


        z_mean = self.np_z_train.mean(axis=0)
        z_min = self.np_z_train.min(axis=0)
        z_max = self.np_z_train.max(axis=0)

        if(len(self.z_mean_train) == 0):
            self.z_mean_train = z_mean
            self.z_min_train = z_min
            self.z_max_train = z_max

        else:
            self.z_mean_train = (z_mean + self.z_mean_train) / 2
            self.z_max_train = np.amax(np.vstack((self.z_max_train, z_max)), axis=0) #Stack old and new together and find the max
            self.z_min_train = np.amin(np.vstack((self.z_min_train, z_min)), axis=0)


        # print('collect_z_values in seconds: {}'.format(time.time() - start))

    def average_mce_batch(self, mce_batch, mce_mini_batch):
        if (mce_batch == None):
            mce_batch = mce_mini_batch
        else:
            for key_lf, mce_lf in mce_batch.items():
                for key, val in mce_lf.items():
                    new_val = mce_mini_batch[key_lf].get(key)
                    if(new_val):
                        mce_batch[key_lf][key] = (new_val + val)/2
        return mce_batch

    #taken from https://github.com/facebookresearch/mixup-cifar10
    def mixup_data(self, x, y, alpha=1.0, use_cuda=True):
        '''Returns mixed inputs, pairs of targets, and lambda'''
        if alpha > 0:
            lam = np.random.beta(alpha, alpha)
        else:
            lam = 1

        batch_size = x.size()[0]
        if use_cuda:
            index = torch.randperm(batch_size).cuda()
        else:
            index = torch.randperm(batch_size)

        mixed_x = lam * x + (1 - lam) * x[index, :]
        y_a, y_b = y, y[index]
        return mixed_x, y_a, y_b, lam

    def training_step(self, batch, batch_idx):
        mce_minibatch=None

        # print('train step')
        # (data, _) = batch
        if(self.used_data=='morpho'):
            ts_batch_user_features = batch.view(-1, 784).float() / 255. #.unsqueeze(1)
        elif(self.used_data=='dsprites'):
            ts_batch_user_features = batch.view(-1, 64*64).float() / 255. #.unsqueeze(1)
        recon_batch, ts_mu_chunk, ts_logvar_chunk, p, q = self.forward(ts_batch_user_features)  # sample data
        if(self.current_epoch == self.max_epochs-1):
            # print("Last round..")
            self.collect_z_values(ts_mu_chunk, ts_logvar_chunk)#, ls_preference

        if (self.current_epoch == self.sigmoid_annealing_threshold ):
            self.collect_z_values(ts_mu_chunk, ts_logvar_chunk)
            # mce_minibatch = metric_utils.mce_batch(self, ts_batch_user_features, self.dct_index2itemId, k=3)
            # self.mce_batch_train = self.average_mce_batch(self.mce_batch_train, mce_minibatch)

        batch_mse, batch_kld = self.loss_function(recon_batch.view(-1, 64*64),
                                                  ts_batch_user_features, #ts_batch_user_features,
                                                  ts_mu_chunk,
                                                  ts_logvar_chunk,
                                                  self.beta,
                                                  self.unique_movies,
                                                  p,
                                                  q,
                                                  new_kld_function = False)
        hp_loss =0
        batch_mse = batch_mse / len(ts_batch_user_features)
        if(self.is_hessian_penalty_activated and self.current_epoch > int(1/3*self.max_epochs-1)):
            print('<---- Applying Hessian Penalty ---->')
            np_z = self.compute_z(ts_mu_chunk, ts_logvar_chunk)
            hp_loss = hessian_penalty(G=self.decode, z=np_z)
            print('Hessian Penalty:{}'.format(hp_loss))
            batch_loss = hp_loss
        else:

            batch_loss = batch_mse + batch_kld
        self.ls_kld.append(self.KLD.tolist())

        #Additional logs go into tensorboard_logs
        tensorboard_logs = {'train_loss': batch_loss,
                'KLD-Train': batch_kld,
                'MSE-Train': batch_mse} #
        return {'loss': batch_loss,
                'log': tensorboard_logs,
                'logvar': np.asarray(ts_mu_chunk.tolist()).mean(axis=0),
                'mu': np.asarray(ts_mu_chunk.tolist()).mean(axis=0)}



    def training_epoch_end(self, outputs):
        print("Saving MCE before KLD is applied...")
        avg_logvar = np.array([x['logvar'] for x in outputs]).mean(axis=0)
        # avg_var = np.array([x['var'] for x in outputs]).mean(axis=0)
        avg_mu = np.array([x['mu'] for x in outputs]).mean(axis=0)

        self.np_logvar_train = np.append(self.np_logvar_train, [avg_logvar], axis=0)
        self.np_mu_train = np.append(self.np_mu_train, [avg_mu], axis=0)



        if(self.current_epoch == self.sigmoid_annealing_threshold ):
            utils.save_dict_as_json(self.mce_batch_train, 'mce_results_wo_kld.json', self.experiment_path_train)
        return {}

    # def validation_step(self, batch, batch_idx):
    #     return 0

    def test_step(self, batch, batch_idx):
        # print('test step')

        batch_mce =0
        test_loss = 0

        # self.eval()
        # ts_batch_user_features = batch.view(-1, self.input_dimension)
        # (data, _) = batch
        if (self.used_data == 'morpho'):
            ts_batch_user_features = batch.view(-1, 784).float() / 255.
        elif (self.used_data == 'dsprites'):
            ts_batch_user_features = batch.view(-1, 64*64).float() / 255.

        if (self.mixup):
            ts_batch_user_features, y_a, y_b, lam = self.mixup_data(ts_batch_user_features, self.test_y_bin,
                                                                    alpha=1.0, use_cuda=False)

        recon_batch, ts_mu_chunk, ts_logvar_chunk, p, q = self.forward(ts_batch_user_features)
        ls_z = self.compute_z(ts_mu_chunk, ts_logvar_chunk)

        self.ls_predicted_movies.extend((-recon_batch).argsort()[:,0].tolist())

        self.np_z_test = np.append(self.np_z_test, np.asarray(ls_z), axis=0) #TODO get rid of np_z_chunk and use np.asarray(mu_chunk)

        self.np_mu_test = np.append(self.np_mu_test, np.asarray(ts_mu_chunk), axis =0)
        self.np_logvar_test = np.append(self.np_logvar_test, np.asarray(ts_logvar_chunk), axis =0)
        # self.np_z = np.vstack((self.np_z, np_z_chunk))

        # batch_rmse_w_zeros, batch_mse_w_zeros, batch_rmse, batch_mse = self.calculate_batch_metrics(recon_batch=recon_batch, ts_batch_user_features =ts_batch_user_features)
        batch_mse, kld = self.loss_function(recon_batch,
                                            ts_batch_user_features, #ts_batch_user_features,
                                            ts_mu_chunk,
                                            ts_logvar_chunk,
                                            self.beta,
                                            self.unique_movies,
                                            p,
                                            q,
                                            new_kld_function=False)
        batch_loss = batch_mse + kld

        # mce_minibatch = metric_utils.mce_batch(self, ts_batch_user_features, self.dct_index2itemId, k=3)
        # self.mce_batch_test = self.average_mce_batch(self.mce_batch_test, mce_minibatch)

        #to be rermoved mean_mce = { for single_mce in batch_mce}
        loss = batch_loss.item() / len(ts_batch_user_features)

        # bce = batch_bce/len(ts_batch_user_features)
        tensorboard_logs = {'KLD-Test': kld,
                            'MSE-test': batch_mse}

        return {'test_loss': loss,
                # 'rmse': batch_rmse,
                'mse': batch_mse,
                # 'rmse_w_zeros': batch_rmse_w_zeros,
                # 'mse_w_zeros': batch_mse_w_zeros,
                'log':tensorboard_logs,
                'KLD-Test': kld,
                            'MSE-Test': batch_mse
                }

        # test_loss /= len(test_loader.dataset)
        # print('====> Test set loss: {:.4f}'.format(test_loss))



    def test_epoch_end(self, outputs):
        avg_loss = np.array([x['test_loss'] for x in outputs]).mean()
        mse_test = np.array([x['MSE-Test'] for x in outputs])
        kld_test =np.array([x['KLD-Test'] for x in outputs])
        # ls_mce = {x['mce'] for x in outputs}
        # utils.save_dict_as_json(self.mce_batch_test, 'mce_results.json', self.experiment_path_test)
        # avg_mce = dict(utils.calculate_mean_of_ls_dict(ls_mce))

        # avg_rmse = np.array([x['rmse'] for x in outputs]).mean()
        # avg_rmse_w_zeros = np.array([x['rmse_w_zeros'] for x in outputs]).mean()
        # avg_mse = np.array([x['mse'] for x in outputs]).mean()
        # avg_mse_w_zeros = np.array([x['mse_w_zeros'] for x in outputs]).mean()

        tensorboard_logs = {'test_loss': avg_loss, 'MSE-Test':mse_test,'KLD-Test': kld_test }
        assert len(mse_test)==len(kld_test)
        for i in range(0, len(mse_test)):
            wandb_logger.log_metrics({'MSE-Test': mse_test[i],'KLD-Test': kld_test[i]} )

        # wandb_logger.log_metrics({'rmse': avg_rmse,
        #                           'rmse_w_zeros':avg_rmse_w_zeros,
        #                           'mse': avg_mse,
        #                           'mse_w_zeros': avg_mse_w_zeros})#, 'kld_matrix':self.kld_matrix

        return {'test_loss': avg_loss,
                'log': tensorboard_logs,
                # 'rmse': avg_rmse,
                'MSE-Test':mse_test,
                'KLD-test': kld_test }#, , 'mce':avg_mce

    def sigmoid_annealing(self, beta, epoch):

        stretch_factor = 0.5
        if(epoch < self.sigmoid_annealing_threshold):
            return 0
        else:
            kld_weight = beta/(1+ math.exp(-epoch * stretch_factor + self.sigmoid_annealing_threshold)) #epoch_threshold moves e function along the x-axis
            return kld_weight

    def kl_divergence(self,p, q):
        return np.sum(np.where(p != 0, p * np.log(p / q), 0))

    def new_kld_func(self, p, q):
        log_qz = q.log_prob(self.z)
        log_pz = p.log_prob(self.z)

        kl = log_qz - log_pz

        return kl

    # Reconstruction + KL divergence losses summed over all elements and batch
    def loss_function(self, recon_x, x, mu, logvar, beta, unique_movies, p, q, new_kld_function=False):
        # MSE = F.binary_cross_entropy(recon_x, x,reduction='sum')
        MSE = F.binary_cross_entropy(recon_x, torch.bernoulli(x),reduction='sum')

        kld_latent_factors = torch.exp(logvar) + mu ** 2 - 1. - logvar
        kld_mean = -0.5 * torch.mean(
            torch.sum(-kld_latent_factors, dim=1))
        self.kld_matrix = np.append(self.kld_matrix, np.asarray(kld_latent_factors.tolist()), axis=0)


        if (self.training):
            kld_weight = self.sigmoid_annealing(beta, self.current_epoch)
        else:
            kld_weight = beta
        self.KLD = kld_mean * kld_weight
        return MSE, self.KLD

    def calculate_batch_metrics(self, recon_batch, ts_batch_user_features):
        # Compute MSE
        # TODO MOre generic ...

        # mask = generate_mask(ts_batch_user_features, tsls_yhat_user, user_based_items_filter=loss_user_items_only)
        # tsls_yhat_user_filtered = tsls_yhat_user[~mask]  # Predicted: Filter out unseen+unrecommended items
        # ts_user_features_seen = ts_batch_user_features[~mask]  # Ground Truth: Filter out unseen+unrecommended items

        # TODO ...than this approach

        batch_rmse = 0
        batch_mse = 0
        batch_rmse_wo_zeros = 0
        batch_mse_wo_zeros = 0
        ls_yhat_user = recon_batch * ts_batch_user_features  # Set all items to zero that are of no interest and haven't been seen
        for idx, tensor in enumerate(ls_yhat_user):
            np_y = ts_batch_user_features[idx].data.numpy()
            np_y_wo_zeros = np_y[np.nonzero(np_y)]  # inner returns the index

            np_yhat = tensor.data.numpy()
            np_yhat_wo_zeros = np_yhat[np.nonzero(np_y)] #This must be np_y

            rmse, mse = metric_utils.calculate_metrics(np_y, np_yhat)
            batch_mse += mse
            batch_rmse += rmse

            if(len(np_yhat_wo_zeros)>0):
                rmse_wo_zeros, mse_wo_zeros = metric_utils.calculate_metrics(np_y_wo_zeros, np_yhat_wo_zeros)
                batch_rmse_wo_zeros += rmse_wo_zeros
                batch_mse_wo_zeros += mse_wo_zeros

        # batch_rmse, batch_mse = utils.calculate_metrics(ts_batch_user_features,ls_yhat_user)
        avg_rmse = batch_rmse / ls_yhat_user.shape[0]
        avg_rmse_wo_zeros = batch_rmse_wo_zeros / ls_yhat_user.shape[0]

        avg_mse = batch_mse / ls_yhat_user.shape[0]
        avg_mse_wo_zeros = batch_mse_wo_zeros / ls_yhat_user.shape[0]
        return avg_rmse, avg_mse, avg_rmse_wo_zeros, avg_mse_wo_zeros

    def load_attributes_and_files(self, path): #'filename.pickle'
        with open(path, 'rb') as handle:
            dct_attributes = pickle.load(handle)
        self.np_z_train = dct_attributes['np_z_train']
        self.np_logvar_train = dct_attributes['np_logvar_train']
        self.np_mu_train = dct_attributes['np_mu_train']
        self.train_y = dct_attributes['train_y']
        self.ls_kld = dct_attributes['ls_kld']

        # self.dct_attribute_distribution = utils.load_json_as_dict(
        #     'attribute_distribution.json')  # load relative frequency distributioon from dictionary (pickle it)

        self.z_max_train = dct_attributes['z_max_train']
        self.z_min_train = dct_attributes['z_min_train']
        self.z_mean_train = dct_attributes['z_mean_train']
        print('Attributes loaded')

    def save_attributes(self, path):
        dct_attributes = {'np_z_train':self.np_z_train,
                          'np_logvar_train': self.np_logvar_train,
                          'np_mu_train': self.np_mu_train,
                          'train_y': self.train_y,
                          'ls_kld':self.ls_kld,
                          'z_max_train': self.z_max_train,
                          'z_min_train': self.z_min_train,
                          'z_mean_train': self.z_mean_train}
        with open(path, 'wb') as handle:
            pickle.dump(dct_attributes, handle)
        print('Attributes saved')


def generate_distribution_df():
    dct_attribute_distribution = utils.compute_relative_frequency(
        pd.read_csv('../data/generated/syn.csv'))
    utils.save_dict_as_json(dct_attribute_distribution, 'syn_attribute_distribution.json')

if __name__ == '__main__':
    torch.manual_seed(100)
    args = run_utils.create_training_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #  use gpu if available
    settings.init()


    #%%
    train = False
    synthetic_data = True
    expanded_user_item = False
    mixup = False
    is_hessian_penalty_activated = False
    continous_data=False
    normalvariate = False
    # morpho = True
    # used_data = "morpho"
    used_data = "dsprites"
    base_path = 'results/models/vae/'

    ls_epochs = [70]#5 with new data, 70 was trained w/ old mnist
    #Note: Mit steigender Epoche wird das disentanglement verstärkt
    #
    ls_latent_factors = [10]
    ls_betas = [4] #disentangle_factors .0003
    no_generative_factors = 3

    for epoch in ls_epochs:
        for lf in ls_latent_factors:
            if(len(ls_betas)==0):
                if(expanded_user_item):
                    beta_normalized = lf/(800)
                else:
                    beta_normalized = lf / (20 * no_generative_factors) #lf/input_size, e.g. 2/10000 = 0.0002
                ls_betas.append(beta_normalized)
            for beta in ls_betas:
                train_tag = "train"
                if(not train):
                    train_tag = "test"

                print("Processing model with: {} epochs, {} latent factors, {} beta".format(epoch, lf, beta))
                exp_name =  "dsprites_mse_{}_data_{}_beta_{}_epochs_{}_lf_synt_{}_normal_{}_continous_{}_hessian_{}".format(used_data, beta, epoch, lf, synthetic_data, normalvariate, continous_data, is_hessian_penalty_activated)
                wandb_name = exp_name + "_" + train_tag
                model_name = exp_name + ".ckpt"
                attribute_name = exp_name + "_attributes.pickle"
                model_path = base_path + model_name
                attribute_path = base_path + attribute_name

                experiment_path = utils.create_experiment_directory()

                model_params = run_utils.create_model_params(experiment_path, epoch, lf, beta, int(epoch / 4), expanded_user_item, mixup,
                                                             no_generative_factors, epoch, is_hessian_penalty_activated, used_data)

                args.max_epochs = epoch

                wandb_logger = WandbLogger(project='recommender-xai', tags=['morpho', train_tag], name=wandb_name)
                trainer = pl.Trainer.from_argparse_args(args,
                                                        logger=wandb_logger, #False
                                                        gpus=0,
                                                        weights_summary='full',
                                                        checkpoint_callback = False,
                                                        callbacks = [ProgressBar(), EarlyStopping(monitor='train_loss')]
                )



                if(train):

                    print('<---------------------------------- VAE Training ---------------------------------->')
                    print("Running with the following configuration: \n{}".format(args))
                    if (synthetic_data):
                        model_params['synthetic_data'], model_params['syn_y'] = data_utils.create_synthetic_data(no_generative_factors,
                                                                                                                 experiment_path,
                                                                                                                 expanded_user_item,
                                                                                                                 continous_data,
                                                                                                                 normalvariate,
                                                                                                                 False)
                        generate_distribution_df()

                    model = VAE(model_params)
                    wandb_logger.watch(model, log='gradients', log_freq=100)

                    # utils.print_nn_summary(model, size =200)

                    print('------ Start Training ------')
                    trainer.fit(model)
                    kld_matrix = model.KLD


                    # print('% altering has provided information gain:{}'.format(
                    #     int(settings.ig_m_hat_cnt) / (int(settings.ig_m_cnt) + int(settings.ig_m_hat_cnt))))

                    # model.dis_KLD


                    print('------ Saving model ------')
                    trainer.save_checkpoint(model_path)
                    model.save_attributes(attribute_path)

                print('------ Load model -------')
                #Sample
                # z = torch.randn(64, lf)
                # sample = model.decode(z)
                # save_image(sample.view(64, 1, 28, 28), './results/mnist_imgs/sample_morpho_20' + '.png')

                # base_curated_test_path = "results/models/curated_vae/"
                # model_path = base_curated_test_path + "4_beta_10_epochs_10_lf_synt_True.ckpt"
                # model_path = base_curated_test_path + "4_beta_16_epochs_10_lf_synt_True.ckpt"
                # model_path = base_curated_test_path + "4_beta_4_epochs_10_lf_synt_True.ckpt"
                # model_path = base_curated_test_path + "4_beta_20_epochs_10_lf_synt_True.ckpt"
                # attribute_path = base_curated_test_path + "4_beta_10_epochs_10_lf_synt_True_attributes.pickle"
                # attribute_path = base_curated_test_path + "4_beta_16_epochs_10_lf_synt_True_attributes.pickle"
                # attribute_path = base_curated_test_path + "4_beta_4_epochs_10_lf_synt_True_attributes.pickle"
                # attribute_path = base_curated_test_path + "4_beta_20_epochs_10_lf_synt_True_attributes.pickle"

                test_model = VAE.load_from_checkpoint(model_path)#, load_saved_attributes=True, saved_attributes_path='attributes.pickle'
                # test_model.test_size = model_params['test_size']
                test_model.load_attributes_and_files(attribute_path)
                test_model.experiment_path_test = experiment_path

                # print("show np_z_train mean:{}, min:{}, max:{}".format(z_mean_train, z_min_train, z_max_train ))
                print('------ Start Test ------')
                start = time.time()
                trainer.test(test_model) #The test loop will not be used until you call.
                # print('Test time in seconds: {}'.format(time.time() - start))
                # print('% altering has provided information gain:{}'.format( int(settings.ig_m_hat_cnt)/(int(settings.ig_m_cnt)+int(settings.ig_m_hat_cnt) )))
                # print(results)

                dct_param ={'epochs':epoch, 'lf':lf,'beta':beta}

                # plot_utils.plot_results(test_model,
                #                    test_model.experiment_path_test,
                #                    test_model.experiment_path_train,
                #                    dct_param )

                disentangle_utils.run_disentanglement_eval(test_model, experiment_path, dct_param)

                artifact = wandb.Artifact('Plots', type='result')
                artifact.add_dir(experiment_path)#, name='images'
                wandb_logger.experiment.log_artifact(artifact)


                working_directory = os.path.abspath(os.getcwd())
                absolute_path = working_directory + "/" + experiment_path + "images/"
                ls_path_images = [absolute_path + file_name for file_name in os.listdir(absolute_path)]
                # wandb.log({"images": [wandb.Image(plt.imread(img_path)) for img_path in ls_path_images]})

                dct_images = {img_path.split(sep='_')[2].split(sep='/')[-1]: wandb.Image(plt.imread(img_path)) for img_path in ls_path_images}
                wandb.log(dct_images)

                # wandb.log({"example_1": wandb.Image(...), "example_2",: wandb.Image(...)})

                #TODO Bring back in

                # neptune_logger.experiment.log_image('MCEs',"./results/images/mce_epochs_"+str(max_epochs)+".png")
                # neptune_logger.experiment.log_artifact("./results/images/mce_epochs_"+str(max_epochs)+".png")
                print('Test done')

    exit()

matplotlib.pyplot.imread