"""
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
EMBEDDING_UNITS: t.List[int] = [64, 32, 16]
# 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] = [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'

# 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,
}

# == 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__ = True


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


@experiment.hook('create_model', default=False)
def create_model(e):
    
    e.log(f'constructing model "Megan2"...')
    e.log(f' * final dropout: {e.FINAL_DROPOUT}')
    model = Megan2(
        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,
        # v2 specific
        fidelity_factor=e.FIDELITY_FACTOR,
        fidelity_funcs=e.FIDELITY_FUNCS,
        embedding_units=e.EMBEDDING_UNITS,
    )
    model.compile(
        loss=[
            mse,
            NoLoss(),
            NoLoss(),
        ],
        loss_weights=[
            1, 0, 0
        ],
        metrics=[mse],
        optimizer=e.OPTIMIZER_CB(),
        run_eagerly=False,
    )
    return model
    


experiment.run_if_main()