"""
This module provides and example of how to load a model which was previously saved to the disk. This can
be done easily using the keras "load_model" function.

Additionally, this module will show how to visualize multiple explanations generated by such a loaded MEGAN
model in a single PDF file using the visualization utilities of ``visual_graph_datasets``
"""
import os
import pathlib
import random

import tensorflow.keras as ks
from pycomex.functional.experiment import Experiment
from pycomex.utils import folder_path, file_namespace
from visual_graph_datasets.data import load_visual_graph_dataset
from visual_graph_datasets.visualization.importances import create_importances_pdf

from graph_attention_student.data import process_index_data_map
from graph_attention_student.keras import CUSTOM_OBJECTS

PATH = pathlib.Path(__file__).parent.absolute()
ASSETS_PATH = os.path.join(PATH, 'assets')
MODEL_PATH = os.path.join(ASSETS_PATH, 'aqsoldb_model')

# ! NOTE: To run the example locally, you will have to download the corresponding visual graph dataset and
# insert the local path here.
VISUAL_GRAPH_DATASET_PATH = '/media/ssd/.visual_graph_datasets/datasets/aqsoldb'
# The number of examples from the dataset to be visualized in the PDF file.
NUM_EXAMPLES = 100

__DEBUG__ = True


@Experiment(base_path=folder_path(__file__),
            namespace=file_namespace(__file__),
            glob=globals())
def experiment(e: Experiment):

    # ~ Loading dataset
    e.log('loading the dataset...')
    metadata_map, index_data_map = load_visual_graph_dataset(
        path=e.VISUAL_GRAPH_DATASET_PATH,
        logger=e.logger,
        log_step=1000,
        metadata_contains_index=True
    )

    dataset_indices, dataset = process_index_data_map(index_data_map)

    # == LOADING MODEL ==
    # Model loading is really simple via the "load_model" function. However, one important aspect that
    # HAS TO be considered is the requirement for the custom object scope as seen here! Without this custom
    # object scope, the loading process would not work properly!
    e.log('loading the model...')
    with ks.utils.custom_object_scope(CUSTOM_OBJECTS):
        model = ks.models.load_model(e.MODEL_PATH)

    # == VISUALIZING RESULTS ==
    # Here we will query the model with a certain number of elements from the dataset and then visualize
    # the predictions as well as the explanations for those predictions all together in one PDF file, which
    # has one page per element.
    e.log('visualizing example predictions...')

    # First of all we need to determine which elements to query the model with and then we need to
    # actually query the model. T
    example_indices = random.sample(dataset_indices, k=e.NUM_EXAMPLES)
    example_graphs = [dataset[i] for i in example_indices]
    predictions = model.predict_graphs(example_graphs)

    output_path = os.path.join(e.path, 'explanations.pdf')
    create_importances_pdf(
        graph_list=example_graphs,
        image_path_list=[index_data_map[i]['image_path'] for i in example_indices],
        node_positions_list=[g['node_positions'] for g in example_graphs],
        importances_map={
            'predicted': (
                [pred[1] for pred in predictions],
                [pred[2] for pred in predictions]
            )
        },
        # It is optionally possible to supply a list of string labels, which assigns additional text that
        # will be printed in the title of each corresponding element.
        labels_list=[f'Prediction: {pred[0]}' for pred in predictions],
        # The absolute path of the final PDF file
        output_path=output_path,
        importance_channel_labels=['negative', 'positive'],
        # Since this can be a somewhat long process, it is possible to log the progress to a custom Logger
        # instance.
        logger=e.logger,
        log_step=10,
    )


experiment.run_if_main()