import os

import time

import numpy as np

import pandas as pd

import sklearn as sk

from sklearn.preprocessing import label_binarize

from util import util

from util import metrics

from torch.utils.tensorboard import SummaryWriter

class Visualizer:

    """

    This class print/save logging information

    """

    def __init__(self, param):

        """

        Initialize the Visualizer class

        """

        self.param = param

        self.output_path = os.path.join(param.checkpoints_dir, param.experiment_name)

        tb_dir = os.path.join(self.output_path, 'tb_log')

        util.mkdir(tb_dir)

        if param.isTrain:

            # Create a logging file to store training losses

            self.train_log_filename = os.path.join(self.output_path, 'train_log.txt')

            with open(self.train_log_filename, 'a') as log_file:

                now = time.strftime('%c')

                log_file.write('----------------------- Training Log ({:s}) -----------------------\n'.format(now))

            self.train_summary_filename = os.path.join(self.output_path, 'train_summary.txt')

            with open(self.train_summary_filename, 'a') as log_file:

                now = time.strftime('%c')

                log_file.write('----------------------- Training Summary ({:s}) -----------------------\n'.format(now))

            # Create log folder for TensorBoard

            tb_train_dir = os.path.join(self.output_path, 'tb_log', 'train')

            util.mkdir(tb_train_dir)

            util.clear_dir(tb_train_dir)

            # Create TensorBoard writer

            self.train_writer = SummaryWriter(log_dir=tb_train_dir)

        if param.isTest:

            # Create a logging file to store testing metrics

            self.test_log_filename = os.path.join(self.output_path, 'test_log.txt')

            with open(self.test_log_filename, 'a') as log_file:

                now = time.strftime('%c')

                log_file.write('----------------------- Testing Log ({:s}) -----------------------\n'.format(now))

            self.test_summary_filename = os.path.join(self.output_path, 'test_summary.txt')

            with open(self.test_summary_filename, 'a') as log_file:

                now = time.strftime('%c')

                log_file.write('----------------------- Testing Summary ({:s}) -----------------------\n'.format(now))

            # Create log folder for TensorBoard

            tb_test_dir = os.path.join(self.output_path, 'tb_log', 'test')

            util.mkdir(tb_test_dir)

            util.clear_dir(tb_test_dir)

            # Create TensorBoard writer

            self.test_writer = SummaryWriter(log_dir=tb_test_dir)

    def print_train_log(self, epoch, iteration, losses_dict, metrics_dict, load_time, comp_time, batch_size, dataset_size, with_time=True):

        """

        print train log on console and save the message to the disk

        Parameters:

            epoch (int)                     -- current epoch

            iteration (int)                 -- current training iteration during this epoch

            losses_dict (OrderedDict)       -- training losses stored in the ordered dict

            metrics_dict (OrderedDict)      -- metrics stored in the ordered dict

            load_time (float)               -- data loading time per data point (normalized by batch_size)

            comp_time (float)               -- computational time per data point (normalized by batch_size)

            batch_size (int)                -- batch size of training

            dataset_size (int)              -- size of the training dataset

            with_time (bool)                -- print the running time or not

        """

        data_point_covered = min((iteration + 1) * batch_size, dataset_size)

        if with_time:

            message = '[TRAIN] [Epoch: {:3d}   Iter: {:4d}   Load_t: {:.3f}   Comp_t: {:.3f}]   '.format(epoch, data_point_covered, load_time, comp_time)

        else:

            message = '[TRAIN] [Epoch: {:3d}   Iter: {:4d}]\n'.format(epoch, data_point_covered)

        for name, loss in losses_dict.items():

            message += '{:s}: {:.3f}   '.format(name, loss[-1])

        for name, metric in metrics_dict.items():

            message += '{:s}: {:.3f}   '.format(name, metric)

        print(message)  # print the message

        with open(self.train_log_filename, 'a') as log_file:

            log_file.write(message + '\n')  # save the message

    def print_train_summary(self, epoch, losses_dict, output_dict, train_time, current_lr):

        """

        print the summary of this training epoch

        Parameters:

            epoch (int)                             -- epoch number of this training model

            losses_dict (OrderedDict)               -- the losses dictionary

            output_dict (OrderedDict)               -- the downstream output dictionary

            train_time (float)                      -- time used for training this epoch

            current_lr (float)                      -- the learning rate of this epoch

        """

        write_message = '{:s}\t'.format(str(epoch))

        print_message = '[TRAIN] [Epoch: {:3d}]\n'.format(int(epoch))

        for name, loss in losses_dict.items():

            write_message += '{:.6f}\t'.format(np.mean(loss))

            print_message += name + ': {:.3f}   '.format(np.mean(loss))

            self.train_writer.add_scalar('loss_'+name, np.mean(loss), epoch)

        metrics_dict = self.get_epoch_metrics(output_dict)

        for name, metric in metrics_dict.items():

            write_message += '{:.6f}\t'.format(metric)

            print_message += name + ': {:.3f}   '.format(metric)

            self.train_writer.add_scalar('metric_'+name, metric, epoch)

        train_time_msg = 'Training time used: {:.3f}s'.format(train_time)

        print_message += '\n' + train_time_msg

        with open(self.train_log_filename, 'a') as log_file:

            log_file.write(train_time_msg + '\n')

        current_lr_msg = 'Learning rate for this epoch: {:.7f}'.format(current_lr)

        print_message += '\n' + current_lr_msg

        self.train_writer.add_scalar('lr', current_lr, epoch)

        with open(self.train_summary_filename, 'a') as log_file:

            log_file.write(write_message + '\n')

        print(print_message)

    def print_test_log(self, epoch, iteration, losses_dict, metrics_dict, batch_size, dataset_size):

        """

        print performance metrics of this iteration on console and save the message to the disk

        Parameters:

            epoch (int)                     -- epoch number of this testing model

            iteration (int)                 -- current testing iteration during this epoch

            losses_dict (OrderedDict)       -- training losses stored in the ordered dict

            metrics_dict (OrderedDict)      -- metrics stored in the ordered dict

            batch_size (int)                -- batch size of testing

            dataset_size (int)              -- size of the testing dataset

        """

        data_point_covered = min((iteration + 1) * batch_size, dataset_size)

        message = '[TEST] [Epoch: {:3d}   Iter: {:4d}]   '.format(int(epoch), data_point_covered)

        for name, loss in losses_dict.items():

            message += '{:s}: {:.3f}   '.format(name, loss[-1])

        for name, metric in metrics_dict.items():

            message += '{:s}: {:.3f}   '.format(name, metric)

        print(message)

        with open(self.test_log_filename, 'a') as log_file:

            log_file.write(message + '\n')

    def print_test_summary(self, epoch, losses_dict, output_dict, test_time):

        """

        print the summary of this testing epoch

        Parameters:

            epoch (int)                             -- epoch number of this testing model

            losses_dict (OrderedDict)               -- the losses dictionary

            output_dict (OrderedDict)               -- the downstream output dictionary

            test_time (float)                       -- time used for testing this epoch

        """

        write_message = '{:s}\t'.format(str(epoch))

        print_message = '[TEST] [Epoch: {:3d}]      '.format(int(epoch))

        for name, loss in losses_dict.items():

            # write_message += '{:.6f}\t'.format(np.mean(loss))

            print_message += name + ': {:.3f}   '.format(np.mean(loss))

            self.test_writer.add_scalar('loss_'+name, np.mean(loss), epoch)

        metrics_dict = self.get_epoch_metrics(output_dict)

        for name, metric in metrics_dict.items():

            write_message += '{:.6f}\t'.format(metric)

            print_message += name + ': {:.3f}   '.format(metric)

            self.test_writer.add_scalar('metric_' + name, metric, epoch)

        with open(self.test_summary_filename, 'a') as log_file:

            log_file.write(write_message + '\n')

        test_time_msg = 'Testing time used: {:.3f}s'.format(test_time)

        print_message += '\n' + test_time_msg

        print(print_message)

        with open(self.test_log_filename, 'a') as log_file:

            log_file.write(test_time_msg + '\n')

    def get_epoch_metrics(self, output_dict):

        """

        Get the downstream task metrics for whole epoch

        Parameters:

            output_dict (OrderedDict)  -- the output dictionary used to compute the downstream task metrics

        """

        if self.param.downstream_task == 'classification':

            y_true = output_dict['y_true'].cpu().numpy()

            y_true_binary = label_binarize(y_true, classes=range(self.param.class_num))

            y_pred = output_dict['y_pred'].cpu().numpy()

            y_prob = output_dict['y_prob'].cpu().numpy()

            if self.param.class_num == 2:

                y_prob = y_prob[:, 1]

            accuracy = sk.metrics.accuracy_score(y_true, y_pred)

            precision = sk.metrics.precision_score(y_true, y_pred, average='macro', zero_division=0)

            recall = sk.metrics.recall_score(y_true, y_pred, average='macro', zero_division=0)

            f1 = sk.metrics.f1_score(y_true, y_pred, average='macro', zero_division=0)

            try:

                auc = sk.metrics.roc_auc_score(y_true_binary, y_prob, multi_class='ovo', average='macro')

            except ValueError:

                auc = -1

                print('ValueError: ROC AUC score is not defined in this case.')

            return {'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1': f1, 'auc': auc}

        elif self.param.downstream_task == 'regression':

            y_true = output_dict['y_true'].cpu().numpy()

            y_pred = output_dict['y_pred'].cpu().detach().numpy()

            mse = sk.metrics.mean_squared_error(y_true, y_pred)

            rmse = sk.metrics.mean_squared_error(y_true, y_pred, squared=False)

            mae = sk.metrics.mean_absolute_error(y_true, y_pred)

            medae = sk.metrics.median_absolute_error(y_true, y_pred)

            r2 = sk.metrics.r2_score(y_true, y_pred)

            return {'mse': mse, 'rmse': rmse, 'mae': mae, 'medae': medae, 'r2': r2}

        elif self.param.downstream_task == 'survival':

            metrics_start_time = time.time()

            y_true_E = output_dict['y_true_E'].cpu().numpy()

            y_true_T = output_dict['y_true_T'].cpu().numpy()

            y_pred_risk = output_dict['risk'].cpu().numpy()

            y_pred_survival = output_dict['survival'].cpu().numpy()

            time_points = util.get_time_points(self.param.survival_T_max, self.param.time_num)

            try:

                c_index = metrics.c_index(y_true_T, y_true_E, y_pred_risk)

            except ValueError:

                c_index = -1

                print('ValueError: NaNs detected in input when calculating c-index.')

            try:

                ibs = metrics.ibs(y_true_T, y_true_E, y_pred_survival, time_points)

            except ValueError:

                ibs = -1

                print('ValueError: NaNs detected in input when calculating integrated brier score.')

            metrics_time = time.time() - metrics_start_time

            print('Metrics computing time: {:.3f}s'.format(metrics_time))

            return {'c-index': c_index, 'ibs': ibs}

        elif self.param.downstream_task == 'multitask':

            metrics_start_time = time.time()

            # Survival

            y_true_E = output_dict['y_true_E'].cpu().numpy()

            y_true_T = output_dict['y_true_T'].cpu().numpy()

            y_pred_risk = output_dict['risk'].cpu().numpy()

            y_pred_survival = output_dict['survival'].cpu().numpy()

            time_points = util.get_time_points(self.param.survival_T_max, self.param.time_num)

            try:

                c_index = metrics.c_index(y_true_T, y_true_E, y_pred_risk)

            except ValueError:

                c_index = -1

                print('ValueError: NaNs detected in input when calculating c-index.')

            try:

                ibs = metrics.ibs(y_true_T, y_true_E, y_pred_survival, time_points)

            except ValueError:

                ibs = -1

                print('ValueError: NaNs detected in input when calculating integrated brier score.')

            # Classification

            y_true_cla = output_dict['y_true_cla'].cpu().numpy()

            y_true_cla_binary = label_binarize(y_true_cla, classes=range(self.param.class_num))

            y_pred_cla = output_dict['y_pred_cla'].cpu().numpy()

            y_prob_cla = output_dict['y_prob_cla'].cpu().numpy()

            if self.param.class_num == 2:

                y_prob_cla = y_prob_cla[:, 1]

            accuracy = sk.metrics.accuracy_score(y_true_cla, y_pred_cla)

            precision = sk.metrics.precision_score(y_true_cla, y_pred_cla, average='macro', zero_division=0)

            recall = sk.metrics.recall_score(y_true_cla, y_pred_cla, average='macro', zero_division=0)

            f1 = sk.metrics.f1_score(y_true_cla, y_pred_cla, average='macro', zero_division=0)

            '''

            try:

                auc = sk.metrics.roc_auc_score(y_true_cla_binary, y_prob_cla, multi_class='ovo', average='macro')

            except ValueError:

                auc = -1

                print('ValueError: ROC AUC score is not defined in this case.')

            '''

            # Regression

            y_true_reg = output_dict['y_true_reg'].cpu().numpy()

            y_pred_reg = output_dict['y_pred_reg'].cpu().detach().numpy()

            # mse = sk.metrics.mean_squared_error(y_true_reg, y_pred_reg)

            rmse = sk.metrics.mean_squared_error(y_true_reg, y_pred_reg, squared=False)

            mae = sk.metrics.mean_absolute_error(y_true_reg, y_pred_reg)

            medae = sk.metrics.median_absolute_error(y_true_reg, y_pred_reg)

            r2 = sk.metrics.r2_score(y_true_reg, y_pred_reg)

            metrics_time = time.time() - metrics_start_time

            print('Metrics computing time: {:.3f}s'.format(metrics_time))

            return {'c-index': c_index, 'ibs': ibs, 'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1': f1, 'rmse': rmse, 'mae': mae, 'medae': medae, 'r2': r2}

        elif self.param.downstream_task == 'alltask':

            metrics_start_time = time.time()

            # Survival

            y_true_E = output_dict['y_true_E'].cpu().numpy()

            y_true_T = output_dict['y_true_T'].cpu().numpy()

            y_pred_risk = output_dict['risk'].cpu().numpy()

            y_pred_survival = output_dict['survival'].cpu().numpy()

            time_points = util.get_time_points(self.param.survival_T_max, self.param.time_num)

            try:

                c_index = metrics.c_index(y_true_T, y_true_E, y_pred_risk)

            except ValueError:

                c_index = -1

                print('ValueError: NaNs detected in input when calculating c-index.')

            try:

                ibs = metrics.ibs(y_true_T, y_true_E, y_pred_survival, time_points)

            except ValueError:

                ibs = -1

                print('ValueError: NaNs detected in input when calculating integrated brier score.')

            # Classification

            accuracy = []

            f1 = []

            auc = []

            for i in range(self.param.task_num - 2):

                y_true_cla = output_dict['y_true_cla'][i].cpu().numpy()

                y_true_cla_binary = label_binarize(y_true_cla, classes=range(self.param.class_num[i]))

                y_pred_cla = output_dict['y_pred_cla'][i].cpu().numpy()

                y_prob_cla = output_dict['y_prob_cla'][i].cpu().numpy()

                if self.param.class_num[i] == 2:

                    y_prob_cla = y_prob_cla[:, 1]

                accuracy.append(sk.metrics.accuracy_score(y_true_cla, y_pred_cla))

                f1.append(sk.metrics.f1_score(y_true_cla, y_pred_cla, average='macro', zero_division=0))

                try:

                    auc.append(sk.metrics.roc_auc_score(y_true_cla_binary, y_prob_cla, multi_class='ovo', average='macro'))

                except ValueError:

                    auc.append(-1)

                    print('ValueError: ROC AUC score is not defined in this case.')

            # Regression

            y_true_reg = output_dict['y_true_reg'].cpu().numpy()

            y_pred_reg = output_dict['y_pred_reg'].cpu().detach().numpy()

            # mse = sk.metrics.mean_squared_error(y_true_reg, y_pred_reg)

            rmse = sk.metrics.mean_squared_error(y_true_reg, y_pred_reg, squared=False)

            # mae = sk.metrics.mean_absolute_error(y_true_reg, y_pred_reg)

            # medae = sk.metrics.median_absolute_error(y_true_reg, y_pred_reg)

            r2 = sk.metrics.r2_score(y_true_reg, y_pred_reg)

            metrics_time = time.time() - metrics_start_time

            print('Metrics computing time: {:.3f}s'.format(metrics_time))

            return {'c-index': c_index, 'ibs': ibs, 'accuracy_1': accuracy[0], 'f1_1': f1[0], 'auc_1': auc[0], 'accuracy_2': accuracy[1], 'f1_2': f1[1], 'auc_2': auc[1], 'accuracy_3': accuracy[2], 'f1_3': f1[2], 'auc_3': auc[2], 'accuracy_4': accuracy[3], 'f1_4': f1[3], 'auc_4': auc[3], 'accuracy_5': accuracy[4], 'f1_5': f1[4], 'auc_5': auc[4], 'rmse': rmse, 'r2': r2}

    def save_output_dict(self, output_dict):

        """

        Save the downstream task output to disk

        Parameters:

            output_dict (OrderedDict)  -- the downstream task output dictionary to be saved

        """

        down_path = os.path.join(self.output_path, 'down_output')

        util.mkdir(down_path)

        if self.param.downstream_task == 'classification':

            # Prepare files

            index = output_dict['index'].numpy()

            y_true = output_dict['y_true'].cpu().numpy()

            y_pred = output_dict['y_pred'].cpu().numpy()

            y_prob = output_dict['y_prob'].cpu().numpy()

            sample_list = self.param.sample_list[index]

            # Output files

            y_df = pd.DataFrame({'sample': sample_list, 'y_true': y_true, 'y_pred': y_pred}, index=index)

            y_df_path = os.path.join(down_path, 'y_df.tsv')

            y_df.to_csv(y_df_path, sep='\t')

            prob_df = pd.DataFrame(y_prob, columns=range(self.param.class_num), index=sample_list)

            y_prob_path = os.path.join(down_path, 'y_prob.tsv')

            prob_df.to_csv(y_prob_path, sep='\t')

        elif self.param.downstream_task == 'regression':

            # Prepare files

            index = output_dict['index'].numpy()

            y_true = output_dict['y_true'].cpu().numpy()

            y_pred = np.squeeze(output_dict['y_pred'].cpu().detach().numpy())

            sample_list = self.param.sample_list[index]

            # Output files

            y_df = pd.DataFrame({'sample': sample_list, 'y_true': y_true, 'y_pred': y_pred}, index=index)

            y_df_path = os.path.join(down_path, 'y_df.tsv')

            y_df.to_csv(y_df_path, sep='\t')

        elif self.param.downstream_task == 'survival':

            # Prepare files

            index = output_dict['index'].numpy()

            y_true_E = output_dict['y_true_E'].cpu().numpy()

            y_true_T = output_dict['y_true_T'].cpu().numpy()

            y_pred_risk = output_dict['risk'].cpu().numpy()

            survival_function = output_dict['survival'].cpu().numpy()

            y_out = output_dict['y_out'].cpu().numpy()

            sample_list = self.param.sample_list[index]

            time_points = util.get_time_points(self.param.survival_T_max, self.param.time_num)

            # Output files

            y_df = pd.DataFrame({'sample': sample_list, 'true_T': y_true_T, 'true_E': y_true_E, 'pred_risk': y_pred_risk}, index=index)

            y_df_path = os.path.join(down_path, 'y_df.tsv')

            y_df.to_csv(y_df_path, sep='\t')

            survival_function_df = pd.DataFrame(survival_function, columns=time_points, index=sample_list)

            survival_function_path = os.path.join(down_path, 'survival_function.tsv')

            survival_function_df.to_csv(survival_function_path, sep='\t')

            y_out_df = pd.DataFrame(y_out, index=sample_list)

            y_out_path = os.path.join(down_path, 'y_out.tsv')

            y_out_df.to_csv(y_out_path, sep='\t')

        elif self.param.downstream_task == 'multitask':

            # Survival

            index = output_dict['index'].numpy()

            y_true_E = output_dict['y_true_E'].cpu().numpy()

            y_true_T = output_dict['y_true_T'].cpu().numpy()

            y_pred_risk = output_dict['risk'].cpu().numpy()

            survival_function = output_dict['survival'].cpu().numpy()

            y_out_sur = output_dict['y_out_sur'].cpu().numpy()

            sample_list = self.param.sample_list[index]

            time_points = util.get_time_points(self.param.survival_T_max, self.param.time_num)

            y_df_sur = pd.DataFrame(

                {'sample': sample_list, 'true_T': y_true_T, 'true_E': y_true_E, 'pred_risk': y_pred_risk}, index=index)

            y_df_sur_path = os.path.join(down_path, 'y_df_survival.tsv')

            y_df_sur.to_csv(y_df_sur_path, sep='\t')

            survival_function_df = pd.DataFrame(survival_function, columns=time_points, index=sample_list)

            survival_function_path = os.path.join(down_path, 'survival_function.tsv')

            survival_function_df.to_csv(survival_function_path, sep='\t')

            y_out_sur_df = pd.DataFrame(y_out_sur, index=sample_list)

            y_out_sur_path = os.path.join(down_path, 'y_out_survival.tsv')

            y_out_sur_df.to_csv(y_out_sur_path, sep='\t')

            # Classification

            y_true_cla = output_dict['y_true_cla'].cpu().numpy()

            y_pred_cla = output_dict['y_pred_cla'].cpu().numpy()

            y_prob_cla = output_dict['y_prob_cla'].cpu().numpy()

            y_df_cla = pd.DataFrame({'sample': sample_list, 'y_true': y_true_cla, 'y_pred': y_pred_cla}, index=index)

            y_df_cla_path = os.path.join(down_path, 'y_df_classification.tsv')

            y_df_cla.to_csv(y_df_cla_path, sep='\t')

            prob_cla_df = pd.DataFrame(y_prob_cla, columns=range(self.param.class_num), index=sample_list)

            y_prob_cla_path = os.path.join(down_path, 'y_prob_classification.tsv')

            prob_cla_df.to_csv(y_prob_cla_path, sep='\t')

            # Regression

            y_true_reg = output_dict['y_true_reg'].cpu().numpy()

            y_pred_reg = np.squeeze(output_dict['y_pred_reg'].cpu().detach().numpy())

            y_df_reg = pd.DataFrame({'sample': sample_list, 'y_true': y_true_reg, 'y_pred': y_pred_reg}, index=index)

            y_df_reg_path = os.path.join(down_path, 'y_df_regression.tsv')

            y_df_reg.to_csv(y_df_reg_path, sep='\t')

        elif self.param.downstream_task == 'alltask':

            # Survival

            index = output_dict['index'].numpy()

            y_true_E = output_dict['y_true_E'].cpu().numpy()

            y_true_T = output_dict['y_true_T'].cpu().numpy()

            y_pred_risk = output_dict['risk'].cpu().numpy()

            survival_function = output_dict['survival'].cpu().numpy()

            y_out_sur = output_dict['y_out_sur'].cpu().numpy()

            sample_list = self.param.sample_list[index]

            time_points = util.get_time_points(self.param.survival_T_max, self.param.time_num)

            y_df_sur = pd.DataFrame(

                {'sample': sample_list, 'true_T': y_true_T, 'true_E': y_true_E, 'pred_risk': y_pred_risk}, index=index)

            y_df_sur_path = os.path.join(down_path, 'y_df_survival.tsv')

            y_df_sur.to_csv(y_df_sur_path, sep='\t')

            survival_function_df = pd.DataFrame(survival_function, columns=time_points, index=sample_list)

            survival_function_path = os.path.join(down_path, 'survival_function.tsv')

            survival_function_df.to_csv(survival_function_path, sep='\t')

            y_out_sur_df = pd.DataFrame(y_out_sur, index=sample_list)

            y_out_sur_path = os.path.join(down_path, 'y_out_survival.tsv')

            y_out_sur_df.to_csv(y_out_sur_path, sep='\t')

            # Classification

            for i in range(self.param.task_num - 2):

                y_true_cla = output_dict['y_true_cla'][i].cpu().numpy()

                y_pred_cla = output_dict['y_pred_cla'][i].cpu().numpy()

                y_prob_cla = output_dict['y_prob_cla'][i].cpu().numpy()

                y_df_cla = pd.DataFrame({'sample': sample_list, 'y_true': y_true_cla, 'y_pred': y_pred_cla}, index=index)

                y_df_cla_path = os.path.join(down_path, 'y_df_classification_'+str(i+1)+'.tsv')

                y_df_cla.to_csv(y_df_cla_path, sep='\t')

                prob_cla_df = pd.DataFrame(y_prob_cla, columns=range(self.param.class_num[i]), index=sample_list)

                y_prob_cla_path = os.path.join(down_path, 'y_prob_classification_'+str(i+1)+'.tsv')

                prob_cla_df.to_csv(y_prob_cla_path, sep='\t')

            # Regression

            y_true_reg = output_dict['y_true_reg'].cpu().numpy()

            y_pred_reg = np.squeeze(output_dict['y_pred_reg'].cpu().detach().numpy())

            y_df_reg = pd.DataFrame({'sample': sample_list, 'y_true': y_true_reg, 'y_pred': y_pred_reg}, index=index)

            y_df_reg_path = os.path.join(down_path, 'y_df_regression.tsv')

            y_df_reg.to_csv(y_df_reg_path, sep='\t')

    def save_latent_space(self, latent_dict, sample_list):

        """

            save the latent space matrix to disc

            Parameters:

                latent_dict (OrderedDict)          -- the latent space dictionary

                sample_list (ndarray)               -- the sample list for the latent matrix

        """

        reordered_sample_list = sample_list[latent_dict['index'].astype(int)]

        latent_df = pd.DataFrame(latent_dict['latent'], index=reordered_sample_list)

        output_path = os.path.join(self.param.checkpoints_dir, self.param.experiment_name, 'latent_space.tsv')

        print('Saving the latent space matrix...')

        latent_df.to_csv(output_path, sep='\t')

    @staticmethod

    def print_phase(phase):

        """

        print the phase information

        Parameters:

            phase (int)             -- the phase of the training process

        """

        if phase == 'p1':

            print('PHASE 1: Unsupervised Phase')

        elif phase == 'p2':

            print('PHASE 2: Supervised Phase')

        elif phase == 'p3':

            print('PHASE 3: Supervised Phase')