"""
The BASE experiment for the training and evaluation of MEGAN model on a visual graph dataset (VGD).
New MEGAN model parameters are introduced in this sub experiment and default values are provided.

CHANGELOG

0.1.0 - initial version
"""
import os
import pathlib
import typing as t

import tensorflow as tf
import tensorflow.keras as ks
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
# from pycomex.util import Skippable
# from pycomex.experiment import SubExperiment
from pycomex.functional.experiment import Experiment
from pycomex.utils import folder_path, file_namespace
from kgcnn.data.utils import ragged_tensor_from_nested_numpy
from umap import UMAP

from graph_attention_student.models import Megan
from graph_attention_student.models import Megan2
from graph_attention_student.models.mappers import IdentityMapper
from graph_attention_student.training import NoLoss, mse, mae
from graph_attention_student.training import LogProgressCallback
from graph_attention_student.fidelity import leave_one_out_analysis
from graph_attention_student.visualization import plot_leave_one_out_analysis
from graph_attention_student.keras import load_model


# == DATASET PARAMETERS ==
# For MEGAN we can use dual channel explanations and thus we can also compare with the dual channel
# ground truth explanations
NODE_IMPORTANCES_KEY: t.Optional[str] = 'node_importances_2'
EDGE_IMPORTANCES_KEY: t.Optional[str] = 'edge_importances_2'
NUM_TARGETS: int = 1


# == MODEL PARAMETERS ==
# These paremeters can be used to configure the model

MODEL_NAME: str = 'MEGAN'

# This list defines how many graph convolutional layers to configure the network with. Every element adds
# one layer. The numbers themselves are the hidden units to be used for those layers.
UNITS: t.List[int] = [32, 32, 32]
# The dropout rate which is applied after EVERY graph convolutional layer of the network. Especially for
# large networks (>20 hidden units and multiple importance channels, this dropout proves very useful)
DROPOUT_RATE: float = 0.1
# This is the weight of the additional explanation-only step which is being applied to the network.
# This explanation only step is important to develop actually interpretable explanations. Refer to the
# paper for more details about this.
IMPORTANCE_FACTOR: float = 1.0
# This is another hyperparameter of the explanation only train step. Usually should be between 1 and 10
IMPORTANCE_MULTIPLIER: float = 2
# This is the number of explanation channels that are generated by the model. This is also the number of
# attention heads used in the graph convolutional layers of the network. So to get a "wider" network, this
# parameter can also be increased. However, note that the value can only be != 2 if importance factor is
# set to exactly 0.0!
IMPORTANCE_CHANNELS: int = 2
# The coefficient value of the explanation sparsity regularization that is applied to the network. Higher
# values should lead to sparser explanations.
SPARSITY_FACTOR: float = 0.3
# If this is true, the results of the individual attention heads will be concatenated, otherwise the
# average will be used to combine the results.
CONCAT_HEADS: bool = False

# We need to supply the range of possible target values and a reference target value a priori to the
# network for the regression case. The regression limits should be as complete as possible. The reference
# does not have to be in the middle, changing it will influence the development of the explanations quite
# a bit.
REGRESSION_REFERENCE: float = 0.0
REGRESSION_WEIGHTS: t.List[float] = [1.0, 1.0]

FIDELITY_FACTOR = 0.1
FIDELITY_FUNCS = [
    lambda org, mod: tf.nn.leaky_relu(mod - org),
    lambda org, mod: tf.nn.leaky_relu(org - mod),
]

CHANNEL_DIRECTIONS = {
    0: -1,
    1: +1,
}

# At the tail end of the network there is a MLP, which does the final prediction. This list defines the
# depth of that MLP. Each element adds a layer and the values themselves dictate the hidden values.
# In the regression case a final entry with "1" hidden layer as the output layer will implicitly be
# added.
FINAL_UNITS: t.List[int] = [32, 16, 1]
# This dropout is applied after EVERY layer in the final MLP. Using this should not be necessary
FINAL_DROPOUT: float = 0.0
FINAL_ACTIVATION: str = 'linear'

# == TRAINING PARAMETERS ==
EPOCHS = 50
BATCH_SIZE = 32
OPTIMIZER_CB = lambda: ks.optimizers.Adam(learning_rate=0.001)


# == EVALUATION PARAMETERS ==
IMPORTANCE_CHANNEL_LABELS = ['negative', 'positive']


# == EXPERIMENT PARAMETERS ==
__DEBUG__ = True
__TESTING__ = False


experiment = Experiment.extend(
    'vgd_single.py',
    base_path=folder_path(__file__),
    namespace=file_namespace(__file__),
    glob=globals()
)


@experiment.hook('create_model')
def create_model(e):
    
    e.log('using model: MEGAN')
    model = Megan(
        units=e.UNITS,
        dropout_rate=e.DROPOUT_RATE,
        importance_factor=e.IMPORTANCE_FACTOR,
        importance_multiplier=e.IMPORTANCE_MULTIPLIER,
        importance_channels=e.IMPORTANCE_CHANNELS,
        sparsity_factor=e.SPARSITY_FACTOR,
        regression_reference=e.REGRESSION_REFERENCE,
        regression_weights=e.REGRESSION_WEIGHTS,
        final_units=e.FINAL_UNITS,
        final_activation=e.FINAL_ACTIVATION,
        final_dropout_rate=e.FINAL_DROPOUT,
        use_edge_features=True,
        concat_heads=e.CONCAT_HEADS,
    )
    model.compile(
        loss=[
            mse,
            NoLoss(),
            NoLoss(),
        ],
        loss_weights=[
            1, 0, 0
        ],
        metrics=[mse],
        optimizer=e.OPTIMIZER_CB(),
        run_eagerly=False,
    )
    return model


@experiment.hook('fit_model')
def fit_model(e, model, x_train, y_train, x_test, y_test):
    e.log('training MEGAN model...')
    try:    
        history = model.fit(
            x_train,
            y_train,
            batch_size=e.BATCH_SIZE,
            epochs=e.EPOCHS,
            validation_data=(x_test, y_test),
            validation_freq=1,
        )
        history = history.history
    except KeyboardInterrupt:
        history = {'loss': []}
        e.log('stopping model training due to keyboard interrupt')
        
    num_params = model.count_params()
    e.log(f'finished training model with {num_params} parameters')
        
    return history


@experiment.hook('query_model')
def query_model(e, model, x, y, include_importances: bool = True):
    e.log('querying the model...')
    out_pred, ni_pred, ei_pred = model(x)

    if include_importances:
        return out_pred, ni_pred, ei_pred
    else:
        return out_pred


@experiment.hook('calculate_fidelity')
def calculate_fidelity(e, model, dataset, indices_true, x_true, y_true, out_pred, ni_pred, ei_pred):
    rep = e['rep']

    if e.IMPORTANCE_CHANNELS != len(e.CHANNEL_DIRECTIONS):
        return [0 for _ in indices_true]

    e.log('leave-one-out analysis...')
    graphs = [dataset[index] for index in indices_true]
    results = leave_one_out_analysis(
        model=model,
        graphs=graphs,
        num_targets=e.NUM_TARGETS,
        num_channels=e.IMPORTANCE_CHANNELS,
    )
    fig = plot_leave_one_out_analysis(
        results=results,
        num_targets=e.NUM_TARGETS,
        num_channels=e.IMPORTANCE_CHANNELS
    )
    fig.savefig(os.path.join(e.path, 'leave_one_out.pdf'))

    # ~ fidelity
    # For each importance channel we construct a mask which only masks out that very channel
    # and then we query the model using that mask, which effectively means that this channel
    # has absolutely no effect on the prediction. We record the outputs generated by these
    # masked predictions and then afterwards calculate the fidelity from that.
    for k in range(IMPORTANCE_CHANNELS):
        # First of all we need to construct the masks
        masks = []
        for ni in ni_pred:
            mask = np.ones_like(ni)
            mask[:, k] = 0
            masks.append(mask)

        masks_tensor = ragged_tensor_from_nested_numpy(masks)
        out_masked, _, _ = [v.numpy() for v in
                            model(x_true, node_importances_mask=masks_tensor)]

        for c, index in enumerate(indices_true):
            e[f'out/pred/{rep}/{index}'] = float(out_pred[c][0])
            e[f"out/masked/{rep}/{index}/{k}"] = float(out_masked[c][0])

    fidelities = []
    for index in indices_true:

        fidelity = 0
        for k in range(IMPORTANCE_CHANNELS):
            out = e[f"out/pred/{rep}/{index}"]
            out_masked = e[f"out/masked/{rep}/{index}/{k}"]
            fidelity += e.CHANNEL_DIRECTIONS[k] * (out - out_masked)

        fidelities.append(fidelity)

    return fidelities


@experiment.hook('save_model')
def save_model(e, model):
    """
    This hook will be called at the very end of the experiment and it's purpose is to implement the saving 
    of the model into a persistent file format, such that it can be re-used later on.
    
    For the MEGAN model, this saving process is relatively simple by calling the keras model "save" method.
    """
    e.log('saving the MEGAN model...')
    model_path = os.path.join(e.path, 'model')
    e['model_path'] = model_path

    model.save(model_path)
    
    
def load_model(e):
    """
    This hook will be called in the experiment analysis to load the model from it's persistent representation 
    in case it will be needed.
    """
    # From the "save_model" we have saved the path into the experiment registry
    model_path = e['model_path']
    model = load_model(model_path)
    
    e.log(f'MEGAN final bias: {model.bias}')
    
    return model
    
    
@experiment.hook('model_post_process')
def model_post_process(e, model, index_data_map):
    """
    This hook will be executed after the model training and the default evaluation but before the 
    default visualization. The purpose of this hook is to implement any additionaly evaluation and 
    visualization functionality.
    
    This implementation implements the visualization of the graph embeddings. All the embeddings for the 
    different explanation channels will be plotted into the same 2D plot and color coded. If the the 
    graph embedding is not inherently 2D, UMAP mapper will be used to transform it into lower dimensional 
    2D space for plotting. 
    """
    e.log('visualizing the graph embedding space for the MEGAN model...')
    
    graphs = [data['metadata']['graph'] for data in index_data_map.values()]
    graph_embeddings = model.embedd_graphs(graphs)
    e.log(f'generated graph embeddings with the shape: {graph_embeddings.shape}')
    
    mapper = IdentityMapper()
    if model.graph_embedding_shape[-1] != 2:
        mapper = UMAP(
            n_neighbors=50,
            min_dist=0, 
            metric='cosine',
            repulsion_strength=2.0,
        )
        
        # To fit the mapping we need a two-dimensional array - specifically we want to consider all the 
        # of the embeddings of the individual channels as individual elements which is why we do the 
        # reshaping here.
        n, k, d = graph_embeddings.shape
        mapper.fit(graph_embeddings.reshape((n * k, d)))
    
    e.log('visualizing the embeddings...')
    fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(12, 12))
    ax.set_title(f'MEGAN Graph Embeddings\n'
                 f'{model.graph_embedding_shape[-1]}D to 2D ({mapper.__class__.__name__})')
    
    for channel_index, color in zip(range(model.importance_channels), mcolors.TABLEAU_COLORS):
        channel_embeddings = mapper.transform(graph_embeddings[:, channel_index, :])
        ax.scatter(
            channel_embeddings[:, 0], channel_embeddings[:, 1],
            c=color,
            s=1,
            label=f'ch. {channel_index}'
        )
        
    ax.legend()
    fig_path = os.path.join(e.path, 'graph_embeddings.pdf')
    fig.savefig(fig_path)
    plt.close(fig)
    


experiment.run_if_main()