"""

CHANGELOG

0.1.0
"""
import os
import pathlib
import typing as t

import tensorflow as tf
import tensorflow.keras as ks
import numpy as np
# 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 graph_attention_student.models import Megan
from graph_attention_student.training import NoLoss, mse
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


# == 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 = [0]
REGRESSION_WEIGHTS = [[1, 1]]

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.Nadam(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,
        fidelity_factor=e.FIDELITY_FACTOR,
        fidelity_funcs=e.FIDELITY_FUNCS,
        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_graph_attributes=False,
        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...')
    history = model.fit(
        x_train,
        y_train,
        batch_size=e.BATCH_SIZE,
        epochs=e.EPOCHS,
        validation_data=(x_test, y_test),
        validation_freq=1,
        callbacks=[
            LogProgressCallback(
                logger=e.logger,
                epoch_step=5,
                identifier='val_output_1_mean_squared_error'
            )
        ],
        verbose=0
    )
    print(model.bias)
    return history.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):
    e.log('saving the MEGAN model...')
    model_path = os.path.join(e.path, 'model')
    e['model_path'] = model_path

    model.save(model_path)


experiment.run_if_main()