""" trainer python script

This script allows training the proposed lfbnet model.

This script requires to specify the directory path to the preprocessed PET MIP images. It could read the patient ids

from

the given directory path, or it could accept patient ids as .xls and .csv files. please provide the directory path to

the

csv or xls file. It assumes the csv/xls file have two columns with level 'train' and 'valid' indicating the training and

validation patient ids respectively.

Please see the _name__ == '__main__': as example which is equivalent to:

e.g.train_valid_data_dir = r"E:\LFBNet\data\remarc_default_MIP_dir/"

    train_valid_ids_path_csv = r'E:\LFBNet\data\csv\training_validation_indexs\remarc/'

    train_ids, valid_ids = get_training_and_validation_ids_from_csv(train_valid_ids_path_csv)

    trainer = NetworkTrainer(

        folder_preprocessed_train=train_valid_data_dir, folder_preprocessed_valid=train_valid_data_dir,

        ids_to_read_train=train_ids,

        ids_to_read_valid=valid_ids

        )

    trainer.train()

"""

# Import libraries

import os

import glob

import sys

import time

from datetime import datetime

import numpy as np

from numpy.random import seed

from random import randint

from tqdm import tqdm

from typing import Tuple, List

from numpy import ndarray

from copy import deepcopy

from medpy.metric import binary

import matplotlib.pyplot as plt

from keras import backend as K

import re

# make LFBNet as parent directory, for absolute import libraries. local application import.

p = os.path.abspath('../..')

if p not in sys.path:

    sys.path.append(p)

# import LFBNet modules

from src.LFBNet.data_loader import DataLoader

from src.LFBNet.network_architecture import lfbnet

from src.LFBNet.losses import losses

from src.LFBNet.preprocessing import save_nii_images

from src.LFBNet.utilities import train_valid_paths

from src.LFBNet.postprocessing import remove_outliers_in_sagittal

# choose cuda gpu

CUDA_VISIBLE_DEVICES = 1

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# set randomness repetable across experiments.

seed(1)

# Define the parameters of the data to process

K.set_image_data_format('channels_last')

def default_training_parameters(

        num_epochs: int = 5000, batch_size: int = 16, early_stop: int = None, fold_number: int = None,

        model_name_save: List[str] = None, loss: str = None, metric: str = None

        ) -> dict:

    """ Configure default parameters for training.

    Training parameters are setted here. For other options, the user should modifier these values.

    Parameters

    ----------

    num_epochs: int, maximum number of epochs to train the model.

    batch_size: int, number of images per batch

    early_stop: int, the number of training epochs the model should train while it is not improving the accuracy.

    fold_number: int, optional, fold number while applying cross-validation-based training.

    model_name_save: str, model name to save

    loss: str, loss funciton

    metric: str, specify the metric, such as dice

    Returns

    -------

    Returns configured dictionary for the training.

    """

    if early_stop is None:

        # early stop 50 % of the maximum number of epochs

        early_stop = int(num_epochs * 0.5)

    if fold_number is None:

        fold_number = 'fold_run_at_' + str(time.time())

    if model_name_save is None:

        model_name_save = ["forward_" + str(time.time()), "feedback_" + str(time.time())]

    if loss is None:

        loss = losses.LossMetric.dice_plus_binary_cross_entropy_loss

    if metric is None:

        metric = losses.LossMetric.dice_metric

    config_trainer = {'num_epochs': num_epochs, 'batch_size': batch_size, 'num_early_stop': early_stop,

                      'fold_number': fold_number, 'model_name_save_forward': model_name_save[0],

                      'model_name_save_feedback': model_name_save[1], "custom_loss": loss, "custom_dice": metric}

    return config_trainer

def get_training_and_validation_ids_from_csv(path):

    """ Get training and validation ids from a given csv or xls file. Assuming the training ids are given with column

    name 'train' and validation ids in 'valid'

    Parameters

    ----------

    path: directory path to the csv or xls file.

    Returns

    -------

    Returns training and validation patient ids.

    """

    ids = train_valid_paths.read_csv_train_valid_index(path)

    train, valid = ids[0], ids[1]

    return train, valid

def get_train_valid_ids_from_folder(path_train_valid, ratio_valid_data=0.25):

    """ Returns the randomly split training and validation patient ids. The percentage of validation is given by the

    ratio_valid_data.

    Parameters

    ----------

    path_train_valid

    ratio_valid_data

     Returns

    -------

    Returns training patient id and validation patient ids respectively as in two array.

    """

    # given training and validation data on one folder, random splitting with .25% : train, valid

    if len(path_train_valid) == 1:

        all_cases_id = os.listdir(str(path_train_valid))  # all patients id

        # make permutation in the given list

        case_ids = np.array(all_cases_id)

        indices = np.random.permutation(len(case_ids))

        num_valid_data = int(ratio_valid_data * len(all_cases_id))

        train, valid = indices[num_valid_data:], indices[:num_valid_data]

        return [train, valid]

class NetworkTrainer:

    """

    class to train the lfb net

    """

    # keep the best loss and dice while training : Value shared across all instances, methods

    BEST_METRIC_VALIDATION = 0  # KEEP THE BEST VALIDATION METRIC SUCH AS THE DICE METRIC (BEST_DICE)

    BEST_LOSS_VALIDATION = 100  # KEEP THE BEST VALIDATION LOSS SUCH AS THE LOSS VALUES (BEST_LOSS)

    EARLY_STOP_COUNT = 0  # COUNTS THE NUMBER OF TRAINING ITERATIONS THE MODEL DID NOT INCREASE, TO COMPARE WITH THE

    now = datetime.now()  # current time, date, month,

    TRAINED_MODEL_IDENTIFIER = re.sub('[ :]', "_", now.ctime())

    # EARLY STOP CRITERIA

    def __init__(

            self, config_trainer: dict = None, folder_preprocessed_train: str = '../data/train/',

            folder_preprocessed_valid: str = '../data/valid/', ids_to_read_train: ndarray = None,

            ids_to_read_valid: ndarray = None, task: str = 'valid', predicted_directory: str = '../data/predicted/',

            save_predicted: bool = False

            ):

        """

        :param config_trainer:

        :param folder_preprocessed_train:

        :param folder_preprocessed_valid:

        :param ids_to_read_train:

        :param ids_to_read_valid:

        :param task:

        :predicted_directory:

        :save_predicted

        """

        if config_trainer is None:

            self.config_trainer = deepcopy(default_training_parameters())

        # training data

        self.folder_preprocessed_train = folder_preprocessed_train

        if ids_to_read_train is None:

            ids_to_read_train = os.listdir(folder_preprocessed_train)

        self.ids_to_read_train = ids_to_read_train

        # validation data

        self.folder_preprocessed_valid = folder_preprocessed_valid

        if ids_to_read_valid is None:

            ids_to_read_valid = os.listdir(folder_preprocessed_valid)

        self.ids_to_read_valid = ids_to_read_valid

        # save predicted directory:

        self.save_all = save_predicted

        self.predicted_directory = predicted_directory

        # load the lfb_network architecture

        self.model = lfbnet.LfbNet()

        self.task = task

        # forward network decoder

        # latent feedback at zero time: means no feedback from feedback network

        self.latent_dim = self.model.latent_dim

        self.h_at_zero_time = np.zeros(

            (int(self.config_trainer['batch_size']), int(self.latent_dim[0]), int(self.latent_dim[1]),

             int(self.latent_dim[2])), np.float32

            )

    @staticmethod

    def load_dataset(directory_: str = None, ids_to_read: List[str] = None):

        """

        :param ids_to_read:

        :param directory_:

        """

        # load batch of data

        data_loader = DataLoader(data_dir=directory_, ids_to_read=ids_to_read)

        image_batch_ground_truth_batch = data_loader.get_batch_of_data()

        batch_input_data, batch_output_data = image_batch_ground_truth_batch[0], image_batch_ground_truth_batch[1]

        # expand dimension for the channel

        batch_output_data = np.expand_dims(batch_output_data, axis=-1)

        batch_input_data = np.expand_dims(batch_input_data, axis=-1)

        return batch_input_data, batch_output_data

    def load_latest_weight(self):

        """ loads the weights of the model with the latest saved weight in the folder ./weight

        """

        # load the last trained weight in the folder weight

        folder_path = r'./weight/'

        file_type = r'\*.h5'

        files = glob.glob(folder_path + file_type)

        try:

            max_file = max(files, key=os.path.getctime)

        except:

            raise Exception("weight could not found !")

        base_name = str(os.path.basename(max_file))

        print(base_name)

        self.model.combine_and_train.load_weights('./weight/forward_system' + str(base_name.split('system')[1]))

        # f

        self.model.fcn_feedback.load_weights('./weight/feedback_system' + str(base_name.split('system')[1]))

    def train(self):

        """Train the model

        """

        batch_size = self.config_trainer['batch_size']

        # self.load_latest_weight()

        # training

        if self.task == 'train':

            # training

            for current_epoch in range(self.config_trainer['num_epochs']):

                feedback_loss_dice = []

                forward_loss_dice = []

                forward_decoder_loss_dice = []

                # shuffle the index of the training data

                index_read = np.random.permutation(int(len(self.ids_to_read_train)))

                # read data

                for selected_patient in range(len(index_read)):

                    # get index of batch of data

                    start = selected_patient * batch_size

                    idx_list_batch = index_read[start:start + batch_size]

                    # if there are still elements in the given batch

                    if idx_list_batch.size > 0:

                        # get index of Why not ? kk = indx_list_batch

                        kk = [str(k) for i, k in enumerate(self.ids_to_read_train) if i in idx_list_batch]

                        batch_input_data, batch_output_data = self.load_dataset(

                            directory_=self.folder_preprocessed_train, ids_to_read=kk

                            )

                        assert len(batch_input_data) > 0, "batch of data not loaded correctly"

                        # shuffle within the batch

                        index_batch = np.random.permutation(int(batch_input_data.shape[0]))

                        batch_input_data = batch_input_data[index_batch]

                        batch_output_data = batch_output_data[index_batch]

                        # batches per epoch: Selected batch might as in id could have more images than the batch size

                        batch_per_epoch = int(batch_input_data.shape[0] / batch_size)

                        for batch_per_epoch_ in range(batch_per_epoch):

                            batch_input = batch_input_data[

                                          batch_per_epoch_ * batch_size:(batch_per_epoch_ + 1) * batch_size]

                            batch_output = batch_output_data[

                                           batch_per_epoch_ * batch_size:(batch_per_epoch_ + 1) * batch_size]

                            # Train forward models

                            if current_epoch % 2 == 0:

                                # step 1: train the forward network encoder and decoder

                                loss, dice = self.model.combine_and_train.train_on_batch(

                                    [batch_input, self.h_at_zero_time], [batch_output]

                                    )  # self.h_at_zero_time

                                forward_loss_dice.append([loss, dice])

                            else:

                                predicted_decoder = self.model.combine_and_train.predict(

                                    [batch_input, self.h_at_zero_time]

                                    )  # , self.h_at_zero_time

                                # step 2: train the feedback network, considering the output of the forward network

                                loss, dice = self.model.fcn_feedback.train_on_batch(predicted_decoder, batch_output)

                                feedback_loss_dice.append([loss, dice])

                                # Step 3: train the forward decoder, considering the trained

                                feedback_latent_result = self.model.feedback_latent.predict([predicted_decoder])

                                forward_encoder_output = self.model.forward_encoder.predict([batch_input])

                                # forward_encoder_output.insert(1, feedback_latent_result)

                                forward_encoder_output = forward_encoder_output[::-1]  # bottleneck should be first

                                forward_encoder_output.insert(1, feedback_latent_result)

                                loss, dice = self.model.forward_decoder.train_on_batch(

                                    [output for output in forward_encoder_output], [batch_output]

                                    )

                                forward_decoder_loss_dice.append([loss, dice])

                forward_loss_dice = np.array(forward_loss_dice)

                feedback_loss_dice = np.array(feedback_loss_dice)

                forward_decoder_loss_dice = np.array(forward_decoder_loss_dice)

                if current_epoch % 2 == 0:

                    loss, dice = np.mean(forward_loss_dice, axis=0)

                    print(

                        'Training_forward_system: >%d, '

                        ' fwd_loss = %.3f, fwd_dice=%0.3f, ' % (current_epoch, loss, dice)

                        )

                else:

                    loss_forward, dice_forward = np.mean(forward_decoder_loss_dice, axis=0)

                    loss_feedback, dice_feedback = np.mean(feedback_loss_dice, axis=0)

                    print(

                        'Training_forward_decoder_and_feedback_system: >%d, '

                        'fwd_decoder_loss=%03f, '

                        'fwd_decoder_dice=%0.3f '

                        'fdb_loss=%03f, '

                        'fdb_dice=%.3f ' % (current_epoch, loss_forward, dice_forward, loss_feedback, dice_feedback)

                        )

                # validation test:

                self.validation(current_epoch=current_epoch)

                # CHECK TRAINING STOPPING CRITERIA:  maximum number of epochs (epoch - 1), meet early stop

                if NetworkTrainer.EARLY_STOP_COUNT == self.config_trainer['num_early_stop']:

                    # save model with early stop identification

                    self.model.combine_and_train.save(

                        'weight/forward_system_early_stopped_' + str(NetworkTrainer.TRAINED_MODEL_IDENTIFIER) + '_.h5'

                        )

                    self.model.fcn_feedback.save(

                        'weight/feedback_system_early_stopped_' + str(NetworkTrainer.TRAINED_MODEL_IDENTIFIER) + '_.h5'

                        )

                    break  # STOP TRAINING WITH BREAK, OR EXIT TRAINING

        # if not training load the last saved weights, and check validation

        elif self.task == 'valid':

            # load the last trained weight in the folder weight

            self.load_latest_weight()

            self.validation(current_epoch=self.config_trainer['num_epochs'])

    def validation(self, verbose: int = 0, current_epoch: int = None):

        """

        Compute the validation dice, loss of the training from the validation data

        """

        # path to the validation data, if not specified, the default path ../data/valid/ would be considered

        folder_preprocessed = self.folder_preprocessed_valid

        # image folder names, or identifier: if not specified the default values would be the name of the folder inside

        # the directory "folder processed" or the self.folder_processed_valid :

        valid_identifier = self.ids_to_read_valid

        '''

        WE CAN IMPLEMENT THE EVALUATION METHOD AS BATCH BASED, PATIENT BASED, OR THE WHOLE-VALIDATION DATA BASED. FOR

        THE LAST OPTION WE NEED TO IMPLEMENT THE EVALUATION() FUNCTION HERE. 

        '''

        ''''

        declare variables to return: 

                        forward loss and dice with h0 (no feedback), 

                        feedback network loss and dice

                        forward decoder loss and dice,

                        forward loss and dice with ht (with feedback latent space)

        '''

        loss_dice = {'loss_fwd_h0': [], 'dice__fwd_h0': [], 'loss_fdb_h0': [], 'dice_fdb_h0': [],

                     'loss_fwd_decoder': [], 'dice_fwd_decoder': [], 'loss_fwd_ht': [], 'dice_fwd_ht': []}

        all_dice_sen_sep = {'dice': [], 'specificity': [], 'sensitivity': []}

        # load the dataset,

        # get the validation ids

        for id_to_validate in valid_identifier:

            try:

                id_to_validate = str(id_to_validate).split('.')[0]

            except:

                pass

            valid_input, valid_output = self.load_dataset(directory_=folder_preprocessed, ids_to_read=[id_to_validate])

            if len(valid_input) == 0:

                print("data %s not read" % id_to_validate)

                continue

            results, dice_sen_sep = self.evaluation(

                input_image=valid_input.copy(), ground_truth=valid_output.copy(), case_name=str(id_to_validate)

                )

            # append all loss to loss and dice to dice from all cases in valid identifiers

            for keys in results.keys():

                loss_dice[str(keys)].append(results[str(keys)][0])

            for keys in dice_sen_sep.keys():

                all_dice_sen_sep[str(keys)].append(dice_sen_sep[str(keys)][0])

        print("\n Dice, sensitivity, specificity \t")

        for k, v in all_dice_sen_sep.items():

            print('%s :  %0.3f ' % (k, np.mean(list(v), axis=0)), end=" ")

        print("\n")

        """

        print the mean of the validation loss and validation dice

        """

        # FOR STOPPING CRITERIA WE ARE USING THE MODEL AT THE 3RD STEP

        dice_mean = np.mean(loss_dice['dice_fwd_ht'])

        loss_mean = np.mean(loss_dice['loss_fwd_ht'])

        # at the first epoch

        if current_epoch == 0:

            NetworkTrainer.BEST_METRIC_VALIDATION = dice_mean

            NetworkTrainer.BEST_LOSS_VALIDATION = loss_mean

        # compare the current dice and loss with the previous epoch's loss and dice:

        # NOW CONSIDER DICE AS OPTIMIZATION METRIC

        print("Current validation loss and metrics at epoch %d: >> " % current_epoch, end=" ")

        for k, v in loss_dice.items():

            print('%s :  %0.3f ' % (k, np.mean(v)), end=" ")

        print("\n")

        if NetworkTrainer.BEST_METRIC_VALIDATION <= dice_mean:

            # reset early stop count, best dice, and best loss values

            NetworkTrainer.BEST_LOSS_VALIDATION = loss_mean

            NetworkTrainer.BEST_METRIC_VALIDATION = dice_mean

            NetworkTrainer.EARLY_STOP_COUNT = 0

            # save the best model weights

            if not os.path.exists('./weight'):

                os.mkdir('./weight')

            self.model.combine_and_train.save(

                'weight/forward_system_' + str(NetworkTrainer.TRAINED_MODEL_IDENTIFIER) + '_.h5'

                )

            self.model.fcn_feedback.save(

                'weight/feedback_system_' + str(NetworkTrainer.TRAINED_MODEL_IDENTIFIER) + '_.h5'

                )

        else:  # just print the current validation metric (dice) and loss, and count early stop

            # Increase the early stop count per epoch

            NetworkTrainer.EARLY_STOP_COUNT += 1

        print(

            '\n Best model on validation data : %0.3f :  Dice: %0.3f \n' % (

                NetworkTrainer.BEST_LOSS_VALIDATION, NetworkTrainer.BEST_METRIC_VALIDATION)

            )

    def evaluation(

            self, verbose: int = 0, input_image: ndarray = None, ground_truth: ndarray = None,

            validation_or_test: str = 'test', case_name: str = None

            ):

        """

        :param case_name:

        :param validation_or_test:

        :param verbose:

        :param input_image:

        :param ground_truth:

        Parameters

        ----------

        save_all

        """

        ''''

        declare variables to return: 

                        forward loss and dice with h0 (no feedback), 

                        feedback network loss and dice

                        forward decoder loss and dice,

                        forward loss and dice with ht (with feedback latent space)

        '''

        all_loss_dice = {'loss_fwd_h0': [], 'dice__fwd_h0': [], 'loss_fdb_h0': [], 'dice_fdb_h0': [],

                         'loss_fwd_decoder': [], 'dice_fwd_decoder': [], 'loss_fwd_ht': [], 'dice_fwd_ht': []}

        dice_sen_sp = {'dice': [], 'specificity': [], 'sensitivity': []}

        # latent  feedback variable h0

        # replace the first number of batches with the number of input images from the first channel

        h0_input = np.zeros(

            (len(input_image), int(self.latent_dim[0]), int(self.latent_dim[1]), int(self.latent_dim[2])), np.float32

            )

        # step 0:

        # Loss and dice on the validation of the forward system

        loss, dice = self.model.combine_and_train.evaluate([input_image, h0_input], [ground_truth], verbose=verbose)

        all_loss_dice['loss_fwd_h0'].append(loss), all_loss_dice['dice__fwd_h0'].append(dice)

        # predict from the forward system

        predicted = self.model.combine_and_train.predict([input_image, h0_input])

        # step 2:

        # Loss and dice on the validation of the feedback system

        loss, dice = self.model.fcn_feedback.evaluate([predicted], [ground_truth], verbose=verbose)

        all_loss_dice['loss_fdb_h0'].append(loss), all_loss_dice['dice_fdb_h0'].append(dice)

        # step 3:

        feedback_latent = self.model.feedback_latent.predict(predicted)  # feedback: hf

        forward_encoder_output = self.model.forward_encoder.predict([input_image])  # forward system's encoder output

        forward_encoder_output = forward_encoder_output[::-1]  # bottleneck should be first

        forward_encoder_output.insert(1, feedback_latent)

        loss, dice = self.model.forward_decoder.evaluate(

            [output for output in forward_encoder_output], [ground_truth], verbose=verbose

            )

        all_loss_dice['loss_fwd_decoder'].append(loss), all_loss_dice['dice_fwd_decoder'].append(dice)

        # loss and dice from the combined and feed back latent space : input  [input_image, fdb_latent_space]

        loss, dice = self.model.combine_and_train.evaluate(

            [input_image, feedback_latent], [ground_truth], verbose=verbose

            )

        all_loss_dice['loss_fwd_ht'].append(loss), all_loss_dice['dice_fwd_ht'].append(dice)

        """

        For the testing time, we use defined metrics on the predicted images instead of using model.evaluate during 

        the validation cases 

        """

        predicted = self.model.combine_and_train.predict([input_image, feedback_latent])

        # binary.dc, sen, and specificty works only on binary images

        dice_sen_sp['dice'].append(

            binary.dc(NetworkTrainer.threshold_image(predicted), NetworkTrainer.threshold_image(ground_truth))

            )

        dice_sen_sp['sensitivity'].append(

            binary.sensitivity(NetworkTrainer.threshold_image(predicted), NetworkTrainer.threshold_image(ground_truth))

            )

        dice_sen_sp['specificity'].append(

            binary.specificity(NetworkTrainer.threshold_image(predicted), NetworkTrainer.threshold_image(ground_truth))

            )

        # all = np.concatenate((ground_truth, predicted, input_image), axis=0)

        # display_image(all)

        # Sometimes save predictions

        if self.save_all:

            predicted = self.model.combine_and_train.predict([input_image, feedback_latent])

            save_nii_images(

                [predicted, ground_truth, input_image], identifier=str(case_name),

                name=[case_name + "_predicted", case_name + "_ground_truth", case_name + "_image"], path_save=self.predicted_directory

                )

        else:

            n = randint(0, 10)

            if n % 3 == 0:

                predicted = self.model.combine_and_train.predict([input_image, feedback_latent])

                save_nii_images(

                    [predicted, ground_truth, input_image], identifier=str(case_name),

                    name=[case_name + "_predicted", case_name + "_ground_truth", case_name + "_image"], path_save=self.predicted_directory

                    )

        return all_loss_dice, dice_sen_sp

    @staticmethod

    def display_image(im_display: ndarray):

        """ display given images

        :param all: 2D image arrays to display

        :returns: display images

        """

        plt.figure(figsize=(10, 8))

        plt.subplots_adjust(hspace=0.5)

        plt.suptitle("Daily closing prices", fontsize=18, y=0.95)

        # loop through the length of tickers and keep track of index

        for n, im in enumerate(im_display):

            # add a new subplot iteratively

            plt.subplot(3, 2, n + 1)

            plt.imshow(im)  # chart formatting

        plt.show()

    @staticmethod

    # binary.dc, sen, and specificty works only on binary images

    def threshold_image(im_: ndarray, thr_value: float = 0.5) -> ndarray:

        """ threshold given input array with the given thresholding value

        :param im_: ndarray of images

        :param thr_value: thresholding value

        :return: threshold array image

        """

        im_[im_ > thr_value] = 1

        im_[im_ < thr_value] = 0

        return im_

class ModelTesting:

    """ performs prediction on a given data set. It predicts the segmentation results, and save the results, calculate

    the clinical metrics such as TMTV, Dmax, sTMTV, sDmax.

    """

    now = datetime.now()  # current time, date, month,

    TRAINED_MODEL_IDENTIFIER = re.sub('[ :]', "_", now.ctime())

    print("current directory", os.getcwd())

    def __init__(

            self, config_test: dict = None, preprocessed_dir: str = '../data/test/', data_list: List[str] = None,

            predicted_dir: str = "../data/predicted"

            ):

        """

        :param config_trainer:

        :param folder_preprocessed_train:

        :param folder_preprocessed_valid:

        :param ids_to_read_train:

        :param ids_to_read_valid:

        :param task:

        :param predicted_dir:

        """

        if config_test is None:

            self.config_test = deepcopy(default_training_parameters())

        # training data

        self.preprocessed_dir = preprocessed_dir

        self.predicted_dir = predicted_dir

        # if the list of testing cases are not given, get from the directory

        if data_list is None:

            data_list = os.listdir(preprocessed_dir)

        self.data_list = data_list

        # load the lfb_network architecture

        self.model = lfbnet.LfbNet()

        # latent feedback at zero time: means no feedback from feedback network

        self.latent_dim = self.model.latent_dim

        # load the last trained weight in the folder weight

        print(os.getcwd())

        folder_path = os.path.join(os.getcwd(), 'src/weight')

        print(folder_path)

        full_path = [path_i for path_i in glob.glob(str(folder_path) + '/*.h5')]

        print("files \n", full_path)

        try:

            max_file = max(full_path, key=os.path.getctime)

        except:

            raise Exception("weight could not found !")

        base_name = str(os.path.basename(max_file))

        print(base_name)

        self.model.combine_and_train.load_weights(

            str(folder_path) + '/forward_system' + str(base_name.split('system')[1])

            )

        # f

        self.model.fcn_feedback.load_weights(str(folder_path) + '/feedback_system' + str(base_name.split('system')[1]))

        self.test()

    def test(self):

        """

                   Compute the validation dice, loss of the training from the validation data

           """

        # path to the validation data, if not specified, the default path ../data/valid/ would be considered

        #

        folder_preprocessed = self.preprocessed_dir

        # image folder names, or identifier: if not specified the default values would be the name of the folder inside

        # the directory "folder processed" or the self.folder_processed_valid :

        test_identifier = self.data_list

        ''''

        declare variables to return if there is a reference segmentation or ground truth : 

                        forward loss and dice with h0 (no feedback), 

                        feedback network loss and dice

                        forward decoder loss and dice,

                        forward loss and dice with ht (with feedback latent space)

        '''

        loss_dice = {'loss_fwd_h0': [], 'dice__fwd_h0': [], 'loss_fdb_h0': [], 'dice_fdb_h0': [],

                     'loss_fwd_decoder': [], 'dice_fwd_decoder': [], 'loss_fwd_ht': [], 'dice_fwd_ht': []}

        # get the validation ids

        test_output = []

        for id_to_test in tqdm(list(test_identifier)):

            test_input, test_output = NetworkTrainer.load_dataset(

                directory_=folder_preprocessed, ids_to_read=[id_to_test]

                )

            if len(test_input) == 0:

                print("data %s not read" % id_to_test)

                continue

            '''

            if there is a ground truth segmentation (gt), and you would like to compare with the predicted segmentation

            by the deep learning model

            '''

            if len(test_output):

                results = self.evaluation_test(

                    input_image=test_input.copy(), ground_truth=test_output.copy(), case_name=str(id_to_test)

                    )

                # append all loss to loss and dice to dice from all cases in valid identifiers

                for keys in results.keys():

                    loss_dice[str(keys)].append(results[str(keys)][0])

                print("Results (sagittal and coronal) for case id: %s  : >> " % id_to_test, end=" ")

                for k, v in loss_dice.items():

                    print('%s :  %0.3f ' % (k, np.mean(v)), end=" ")

                print("\n")

            # Predict the segmentation and save in the folder predicted, dataset identifier

            else:

                self.prediction(input_image=test_input.copy(), case_name=str(id_to_test))

        """

        print the mean of the testing loss and dice if there is a ground truth, for all cases 

        """

        if len(test_output):

            print("Total dataset metrics:  : >> ", end=" ")

            for k, v in loss_dice.items():

                print('%s :  %0.3f ' % (k, np.mean(v)), end=" ")

            print("\n")

    def evaluation_test(

            self, verbose: int = 0, input_image: ndarray = None, ground_truth: ndarray = None,

            validation_or_test: str = 'validate', case_name: str = None

            ):

        """

        :param case_name:

        :param validation_or_test:

        :param verbose:

        :param input_image:

        :param ground_truth:

        """

        ''''

        declare variables to return: 

                        forward loss and dice with h0 (no feedback), 

                        feedback network loss and dice

                        forward decoder loss and dice,

                        forward loss and dice with ht (with feedback latent space)

        '''

        all_loss_dice = {'loss_fwd_h0': [], 'dice__fwd_h0': [], 'loss_fdb_h0': [], 'dice_fdb_h0': [],

                         'loss_fwd_decoder': [], 'dice_fwd_decoder': [], 'loss_fwd_ht': [], 'dice_fwd_ht': []}

        # latent  feedback variable h0

        # replace the first number of batches with the number of input images from the first channel

        h0_input = np.zeros(

            (len(input_image), int(self.latent_dim[0]), int(self.latent_dim[1]), int(self.latent_dim[2])), np.float32

            )

        # step 0:

        # Loss and dice on the validation of the forward system

        loss, dice = self.model.combine_and_train.evaluate([input_image, h0_input], [ground_truth], verbose=verbose)

        all_loss_dice['loss_fwd_h0'].append(loss), all_loss_dice['dice__fwd_h0'].append(dice)

        # predict from the forward system

        predicted = self.model.combine_and_train.predict([input_image, h0_input])

        # step 2:

        # Loss and dice on the validation of the feedback system

        loss, dice = self.model.fcn_feedback.evaluate([predicted], [ground_truth], verbose=verbose)

        all_loss_dice['loss_fdb_h0'].append(loss), all_loss_dice['dice_fdb_h0'].append(dice)

        # step 3:

        feedback_latent = self.model.feedback_latent.predict(predicted)  # feedback: hf

        forward_encoder_output = self.model.forward_encoder.predict([input_image])  # forward system's encoder output

        forward_encoder_output = forward_encoder_output[::-1]  # bottleneck should be first

        forward_encoder_output.insert(1, feedback_latent)

        loss, dice = self.model.forward_decoder.evaluate(

            [output for output in forward_encoder_output], [ground_truth], verbose=verbose

            )

        all_loss_dice['loss_fwd_decoder'].append(loss), all_loss_dice['dice_fwd_decoder'].append(dice)

        # loss and dice from the combined and feed back latent space : input  [input_image, fdb_latent_space]

        loss, dice = self.model.combine_and_train.evaluate(

            [input_image, feedback_latent], [ground_truth], verbose=verbose

            )

        all_loss_dice['loss_fwd_ht'].append(loss), all_loss_dice['dice_fwd_ht'].append(dice)

        """

        For the testing time, we use defined metrics on the predicted images instead of using model.evaluate during 

        the validation cases 

        """

        if validation_or_test == "test":

            # return [dice, specificity, and sensitivity

            return {'dice': binary.dc(predicted, ground_truth),

                    'specificity': binary.specificity(predicted, ground_truth),

                    'sensitivity': binary.sensitivity(predicted, ground_truth)}

        predicted = self.model.combine_and_train.predict([input_image, feedback_latent])

        predicted = remove_outliers_in_sagittal(predicted)

        save_nii_images(

            [predicted, ground_truth, input_image], identifier=str(case_name),

            name=[case_name + "_predicted", case_name + "_ground_truth", case_name + "_pet"],

            path_save= os.path.join(str(self.predicted_dir), 'predicted_data')

            )

        return all_loss_dice

    def prediction(self, input_image: ndarray = None, case_name: str = None):

        """

        :param case_name:

        :param input_image:

        """

        # latent  feedback variable h0

        # replace the first number of batches with the number of input images from the first channel

        h0_input = np.zeros(

            (len(input_image), int(self.latent_dim[0]), int(self.latent_dim[1]), int(self.latent_dim[2])), np.float32

            )

        # STEP 1: forward system prediction

        # predict from the forward system

        predicted = self.model.combine_and_train.predict([input_image, h0_input])

        # step 2: Feedback system prediction

        feedback_latent = self.model.feedback_latent.predict(predicted)  # feedback: hf

        predicted = self.model.combine_and_train.predict([input_image, feedback_latent])

        predicted = remove_outliers_in_sagittal(predicted)

        save_nii_images(

            image=[predicted, input_image], identifier=str(case_name), name=[case_name + "_predicted",

                                                                             case_name + "_pet"],

            path_save= os.path.join(str(self.predicted_dir), 'predicted_data')

            )

if __name__ == '__main__':

    train_valid_data_dir = r"E:\LFBNet\data\remarc_default_MIP_dir/"

    train_valid_ids_path_csv = r'E:\LFBNet\data\csv\training_validation_indexs\remarc/'

    train_ids, valid_ids = get_training_and_validation_ids_from_csv(train_valid_ids_path_csv)

    trainer = NetworkTrainer(

        folder_preprocessed_train=train_valid_data_dir, folder_preprocessed_valid=train_valid_data_dir,

        ids_to_read_train=train_ids, ids_to_read_valid=valid_ids

        )

    trainer.train()