# python modules

import argparse, os, pickle

import logging

import errno

from pprint import pprint

import pandas as pd

import json

# os modifications

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

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

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

# tensorflow imports

import tensorflow

from tensorflow.keras.models import Model

physical_devices = tensorflow.config.list_physical_devices('GPU')

if len(physical_devices) > 0:

    USE_GPU = 1

else:

    USE_GPU = 0

for device in physical_devices:

    tensorflow.config.experimental.set_memory_growth(device, True)

# import utils and models    

from utils.metrics import compute_scores

from models import *

from modules.image_encoder import load_encoded_vecs

from utils import *

from utils.dataset import Dataset, IuXrayDataset, ImageCLEFDataset

# import nltk

import nltk

nltk.download('punkt', quiet=True)

# store dataset as well as results path

DATASET_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')

RESULTS_PATH = os.path.join(DATASET_PATH, 'results')

def make_dir(str_path:str) -> None:

    """ Try to make directory properly

    Args:

        str_path (str): The str path to create our directory

    """

    try:

        os.mkdir(str_path)

    except OSError as exc:

        if exc.errno != errno.EEXIST:

            raise

        pass

# make results directory    

make_dir(RESULTS_PATH)

# begin loggings

logging.basicConfig(format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.INFO)

class DiagnosticCaptioning:

    def __init__(self):

        """ Main application to implement my created systems

        """

        # fetch user cmd selections

        self.parser = argparse.ArgumentParser()

        self.parse_agrs()

    def parse_agrs(self) -> None:

        """ Parse all arguments selected in execution from the user

        """

        # Data loader settings

        self.parser.add_argument("--dataset", type=str, default="iu_xray", choices=["iu_xray", "imageclef"], help="the dataset to be used.")

        # Employing model

        self.parser.add_argument("--model_choice", type=str, default="cnn_rnn", choices=["cnn_rnn", "knn"], help="Which model to employ for testing.")

        self.parser.add_argument("--k", type=int, default=5, help="k for K-NN")

        # Captions settings

        self.parser.add_argument("--max_length", type=int, default=40, help="the maximum sequence length of the reports.")

        self.parser.add_argument("--threshold", type=int, default=3, help="the cut off frequency for the words.")

        # Model settings (for layers)

        self.parser.add_argument("--image_encoder", type=str, default="densenet121", help="the visual encoder to be used.")

        self.parser.add_argument("--embedding_dim", type=int, default=100, help="the embedding dimension for Embedding Layers.")

        self.parser.add_argument("--ling_model", type=str, default="gru", choices=["gru", "lstm", "bigru"], help="the Linguistig Model (RNN) for Decoder module as well as Text encoder.")

        # Model settings

        self.parser.add_argument("--multi_modal", type=bool, default=False, help="if to use multi_modal as our model for CNN-RNN only.")

        self.parser.add_argument("--dropout", type=float, default=0.2, help="the dropout rate of our model.")

        # Generate text apporach related

        self.parser.add_argument("--sample_method", type=str, default="greedy", choices=["greedy", "beam_3", "beam_5", "beam_7"], help="the sample methods to sample a report.")

        # Trainer settings

        self.parser.add_argument("--batch_size", type=int, default=8, help="the number of samples for a batch",)

        self.parser.add_argument("--n_gpu", type=int, default=USE_GPU, help="the number of gpus to be used.")

        self.parser.add_argument("--epochs", type=int, default=100, help="the number of training epochs.")

        self.parser.add_argument("--save_dir",type=str, default="cnn_rnn",help="the path to save the models.")

        self.parser.add_argument("--early_stop", type=int, default=10, help="the patience of training.")

    def __init_device(self) -> tuple[bool, bool, bool]: 

        """ Private method to initialize the GPU usage if available else CPU

        Returns:

            tuple[bool, bool, bool]: Bool variables whether to use sinlge or multiple GPUs if available else CPU

        """

        use_CPU, use_GPU, use_multiGPU = False, False, False

        n_gpus = self.parser.parse_args().n_gpu

        # case GPU available

        if n_gpus > 0:

            if n_gpus == 1:

                use_GPU = True

            else:

                use_multiGPU = True

        else:

            # case CPU available

            use_CPU = True

        return use_CPU, use_GPU, use_multiGPU

    def __load_iuxray_data(self) -> tuple[dict, dict, dict]:

        """ Loads IU X-Ray dataset from directory

        Returns:

            tuple[dict, dict, dict]: Image vectors, captions and tags in dictionary format, with keys to be the Image IDs.

        """

        # get dataset path

        iu_xray_data_path = os.path.join(DATASET_PATH, 'iu_xray')

        iu_xray_images_data_path = os.path.join(iu_xray_data_path, 'two_images.json')

        iu_xray_captions_data_path = os.path.join(iu_xray_data_path, 'two_captions.json')

        iu_xray_tags_data_path = os.path.join(iu_xray_data_path, 'two_tags.json')

        # fetch images, captions, tags

        with open(iu_xray_images_data_path) as json_file:

            images = json.load(json_file)

        with open(iu_xray_captions_data_path) as json_file:

            captions = json.load(json_file)

        with open(iu_xray_tags_data_path) as json_file:

            tags = json.load(json_file)

        encoder = self.parser.parse_args().image_encoder

        image_encoded_vectors_path = os.path.join(iu_xray_data_path, f"{encoder}.pkl")

        # load image embeddings for the employed encoder      

        image_vecs = load_encoded_vecs(image_encoded_vectors_path)

        return image_vecs, captions, tags

    def __load_imageclef_data(self) -> tuple[dict, dict]:

        """ Loads ImageCLEF dataset from directory

        Returns:

            tuple[dict, dict]: Image vectors, captions in dictionary format, with keys to be the Image IDs.

        """

         # get dataset path

        imageclef_data_path = os.path.join(DATASET_PATH, 'imageCLEF')

        # fetch images, captions

        imageclef_image_captions_pairs = os.path.join(imageclef_data_path, 'Imageclef2022_dataset_all.csv')

        clef_df = pd.read_csv(imageclef_image_captions_pairs, sep='\t')

        captions = dict( zip( clef_df.ID.to_list(), clef_df.caption.to_list() ) )

        encoder = self.parser.parse_args().image_encoder

        image_encoded_vectors_path = os.path.join(imageclef_data_path, f"{encoder}.pkl")

        # load image embeddings for the employed encoder   

        image_vecs = load_encoded_vecs(image_encoded_vectors_path)

        return image_vecs, captions

    def __create_iu_xray_dataset(self, images:dict, captions:dict, tags:dict) -> IuXrayDataset:

        """ Builds the IU X-Ray dataset using the IuXrayDataset loader class

        Args:

            images (dict): Dictionary with keys to be the ImageIDs and values the image embeddings.

            captions (dict): Dictionary with keys to be the ImageIDs and values the captions.

            tags (dict): Dictionary with keys to be the ImageIDs and values the tags embeddings.

        Returns:

            IuXrayDataset: the employed IuXrayDataset object

        """

        iu_xray_dataset = IuXrayDataset(image_vectors=images, captions_data=captions, tags_data=tags)

        logging.info('IU-XRay dataset created.')

        logging.info(iu_xray_dataset)

        return iu_xray_dataset

    def __create_imageCLEF_dataset(self, images:dict, captions:dict) -> ImageCLEFDataset:

        """ Builds the ImageCLEF dataset using the ImageCLEFDataset loader class

        Args:

            images (dict): Dictionary with keys to be the ImageIDs and values the image embeddings.

            captions (dict): Dictionary with keys to be the ImageIDs and values the captions.

        Returns:

            ImageCLEFDataset: the employed ImageCLEFDataset object

        """

        imageCLEF_dataset = ImageCLEFDataset(image_vectors=images, captions_data=captions)

        logging.info('ImageCLEF dataset created.')

        logging.info(imageCLEF_dataset)

        return imageCLEF_dataset

    def train_cnn_rnn(self, dataset:Dataset) -> tuple[CNN_RNN, Model]:

        """ Begins the training process for the implemented CNN-RNN model

        More details are provided in my Thesis

        Args:

            dataset (Dataset): The employed dataset, i.e. IU X-Ray or ImageCLEF

        Returns:

            CNN_RNN, Model: The created CNN-RNN and the trained model

        """

        # fetch important args

        which_dataset = self.parser.parse_args().dataset

        epochs = self.parser.parse_args().epochs

        encoder = self.parser.parse_args().image_encoder

        max_length = self.parser.parse_args().max_length

        embedding_dim = self.parser.parse_args().embedding_dim

        ling_model = self.parser.parse_args().ling_model

        multi_modal = self.parser.parse_args().multi_modal

        logging.info(multi_modal)

        batch_size = self.parser.parse_args().batch_size

        # create the save directory for the model

        saved_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), self.parser.parse_args().save_dir)

        make_dir(saved_dir)

        # get the created vocabulary for our CNN-RNN

        _, tokenizer, word2idx, idx2word = dataset.get_tokenizer_utils()

        # make the model name according to arguments

        model_name = f'{which_dataset}_enc{encoder}_epochs{epochs}_maxlen{max_length}_embed{embedding_dim}_lingmodel{ling_model}_multimodal{multi_modal}'

        saved_model_name = os.path.join(saved_dir, model_name)

        logging.info(f'CNN-RNN model will be saved at: {saved_model_name}.h5')

        # build the CNN-RNN model

        SnT = CNN_RNN(tokenizer=tokenizer, word_to_idx=word2idx, 

                        idx_to_word=idx2word, max_length=max_length, 

                        embedding_dim=embedding_dim, ling_model=ling_model, 

                        multi_modal=multi_modal, loss="categorical_crossentropy")

        logging.info(f'Utilized vocabulary contains {SnT.vocab_size} words!')

        # get dataset splits   

        train, dev, test = dataset.get_splits_sets()

        # case IU X-Ray

        if which_dataset == 'iu_xray':

            # fetch all tags

            all_tags = dict(train[2], **dev[2])

            all_tags = dict(all_tags, **test[2])

            print('TAGS:', len(all_tags))

            # initialize the Multi-Modal version if user selected this kind of network

            tags_patient_pair = SnT.build_multimodal_encoder(all_tags)

            train_tags = {

                    key:value for key,value in tags_patient_pair.items() if key in train[1].keys()

            }      

            # store training data we want to utilise

            # 1st index --> image vectors

            # 2nd index --> captions

            # 3rd index --> tags

            train_data = [train[0], train[1], train_tags]

        else:

             # case ImageCLEF

             # store training data we want to utilise

             # 1st index --> image vectors

             # 2nd index --> captions

            train_data = [train[0], train[1]]

        # we use Adam as our optimizer for our training procedure  

        optimizer = tensorflow.keras.optimizers.Adam()

        # case IU X-Ray

        if which_dataset == 'iu_xray':

            # get the image embedding input shape. Every patient in IU X-Ray has 2 medical images. Thus, we read the shape from the first one.

            image_input_shape = list(train[0].values())[0][0].shape[1]

            # start train

            trained_model = SnT.train_iuxray_model(train_data=train_data, 

                                                    input_shape=(image_input_shape,), 

                                                    optimizer=optimizer, 

                                                    model_name=saved_model_name, 

                                                    n_epochs=epochs, 

                                                    batch_size=batch_size)

        else:

            # case ImageCLEF

              # get the image embedding input shape.

            image_input_shape = list(train[0].values())[0].shape[1]

            # start train

            trained_model = SnT.train_imageclef_model(train_data=train_data, 

                                                    input_shape=(image_input_shape,), 

                                                    optimizer=optimizer, 

                                                    model_name=saved_model_name, 

                                                    n_epochs=epochs, 

                                                    batch_size=batch_size)

        return SnT, trained_model

    def eval_cnn_rnn(self, cnn_rnn:CNN_RNN, model_to_eval:Model, dataset:Dataset) -> None:

        """ Begins the evaluation process for the trained model in the given dataset

        Args:

            cnn_rnn (CNN_RNN): The created CNN-RNN object that we will employ to apply our evaluation method

            model_to_eval (Model): The trained model that will be assessed

            dataset (Dataset): The employed dataset (IU X-Ray, ImageCLEF)

        """

        # fetch the generation algorithm (Greedy or Beam Search)

        generate_choice = self.parser.parse_args().sample_method

        which_dataset = self.parser.parse_args().dataset

        # fetch dev, test set

        _, dev, test = dataset.get_splits_sets()

        # first evaluate our model in validation set

        if which_dataset == 'iu_xray':

            gold, predicted = cnn_rnn.evaluate_model(model=model_to_eval, 

                                                            test_captions=dev[1], 

                                                            test_images=dev[0], 

                                                            test_tags=dev[2], 

                                                            evaluator_choice=generate_choice)

        else:

            gold, predicted = cnn_rnn.evaluate_model(model=model_to_eval, 

                                                            test_captions=dev[1], 

                                                            test_images=dev[0], 

                                                            test_tags=None, 

                                                            evaluator_choice=generate_choice)

        # get the results path for our results dataframe

        dev_gold_path = os.path.join(RESULTS_PATH, 'dev_gold.csv')

        dev_pred_path = os.path.join(RESULTS_PATH, 'dev_pred.csv')

        # save gold truth captions

        df_gold = pd.DataFrame.from_dict(gold, orient="index")

        df_gold.to_csv(dev_gold_path, sep='|', header=False)

        # save predicted captions  

        df_pred = pd.DataFrame.from_dict(predicted, orient="index")

        df_pred.to_csv(dev_pred_path, sep='|', header=False)

        # score

        scores = compute_scores(gts=dev_gold_path, res=dev_pred_path, scores_filename='dev_set_cnn_rnn_scores', save_scores=True)

        print('CNN_RNN scores in Validation set')

        pprint(scores)

        # Now evaluate our model in test set

        if which_dataset == 'iu_xray':

            gold, predicted = cnn_rnn.evaluate_model(model=model_to_eval, 

                                                            test_captions=test[1], 

                                                            test_images=test[0], 

                                                            test_tags=test[2],

                                                            eval_dataset=which_dataset,

                                                            evaluator_choice=generate_choice)

        else:

            gold, predicted = cnn_rnn.evaluate_model(model=model_to_eval, 

                                                            test_captions=test[1], 

                                                            test_images=test[0], 

                                                            test_tags=None,

                                                            eval_dataset=which_dataset, 

                                                            evaluator_choice=generate_choice)

        # get the results path for our results dataframe

        dev_gold_path = os.path.join(RESULTS_PATH, 'test_gold.csv')

        dev_pred_path = os.path.join(RESULTS_PATH, 'test_pred.csv')

         # save gold truth captions 

        df_gold = pd.DataFrame.from_dict(gold, orient="index")

        df_gold.to_csv(dev_gold_path, sep='|', header=False)

        # save predicted captions  

        df_pred = pd.DataFrame.from_dict(predicted, orient="index")

        df_pred.to_csv(dev_pred_path, sep='|', header=False)

        # score

        scores = compute_scores(gts=dev_gold_path, res=dev_pred_path, scores_filename='test_set_cnn_rnn_scores', save_scores=True)

        print('CNN_RNN scores in Test set')

        pprint(scores)

    def run_process(self) -> None:

        """ Begins the whole process according to the user settings.

        It employes the selected dataset in the selected model.

        For the latter we have CNN-RNN and kNN. More details for each of these models are provided in my Thesis.

        """

        which_dataset = self.parser.parse_args().dataset

        employed_model = self.parser.parse_args().model_choice

        # case IU X-Ray

        if which_dataset == "iu_xray":

            image_vecs, captions, tags = self.__load_iuxray_data()

            iu_xray_dataset = self.__create_iu_xray_dataset(image_vecs, captions, tags)

            # case CNN-RNN

            if employed_model == 'cnn_rnn':

                # Train CNN-RNN model

                cnn_rnn, trained_model = self.train_cnn_rnn(dataset=iu_xray_dataset)

                # Evaluate in model in Validation and Test set

                self.eval_cnn_rnn(cnn_rnn=cnn_rnn, model_to_eval=trained_model, dataset=iu_xray_dataset)

            else:

                 # case k-NN

                k = self.parser.parse_args().k

                multi_modal = self.parser.parse_args().multi_modal

                kNN = KNN(dataset=iu_xray_dataset, k=k, similarity_function='cosine', text_model='clinical_bert')

                # init the results path

                results_path = os.path.join(RESULTS_PATH, 'iuxray_{k}-NN_test_captions.csv')

                # and execute the k-NN algorithm

                kNN.run_algo(multi_modal = multi_modal, results_dir_path=results_path)

        else:

            # case ImageCLEF

            image_vecs, captions = self.__load_imageclef_data()

            imageCLEF_dataset = self.__create_imageCLEF_dataset(image_vecs, captions)

            # case CNN-RNN

            if employed_model == 'cnn_rnn':

                # Train CNN-RNN model

                cnn_rnn, trained_model = self.train_cnn_rnn(dataset=imageCLEF_dataset)

                # Evaluate in model in Validation and Test set

                self.eval_cnn_rnn(cnn_rnn=cnn_rnn, model_to_eval=trained_model, dataset=imageCLEF_dataset)

            else:

                 # case k-NN

                k = self.parser.parse_args().k

                kNN = KNN(dataset=imageCLEF_dataset, k=k, similarity_function='cosine', text_model='clinical_bert')

                # init the results path

                results_path = os.path.join(RESULTS_PATH, 'imageclef_{k}-NN_test_captions.csv')

                 # and execute the k-NN algorithm

                kNN.run_algo(results_dir_path=results_path)

    def main(self) -> None:

        """ Begins the process for this application

        """

        # flags for GPU and CPU usage

        use_CPU, use_GPU, _ = self.__init_device()

        if use_CPU:

            logging.info('Using CPU')

            with tensorflow.device("/device:GPU:0"):

                self.run_process()

        elif use_GPU:

            logging.info('Using single GPU')

            with tensorflow.device("/device:GPU:0"):

                self.run_process()

        else:

            logging.info('Using multi GPU')

            tensorflow.debugging.set_log_device_placement(True)

            gpus = tensorflow.config.list_logical_devices("GPU")

            strategy = tensorflow.distribute.MirroredStrategy(gpus)

            with strategy.scope():

                self.run_process()

if __name__ == '__main__':

    logging.info(DATASET_PATH)

    dc = DiagnosticCaptioning()

    dc.main()