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 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


# == 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'


# == 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.0

# 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 = 5

# 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 = 5.0

# 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_LIMITS = [[-3, 3]]

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 = [60, 20, 1]

# This dropout is applied after EVERY layer in the final MLP. Using this should not be necessary
FINAL_DROPOUT = 0.0

# == TRAINING PARAMETERS ==
EPOCHS = 250
OPTIMIZER_CB = lambda: ks.optimizers.Nadam(learning_rate=0.01)


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


# == EXPERIMENT PARAMETERS ==
PATH = pathlib.Path(__file__).parent.absolute()
EXPERIMENT_PATH = os.path.join(PATH, 'vgd_single.py')
BASE_PATH = PATH
NAMESPACE = 'results/vgd_single_megan'
DEBUG = True
with Skippable(), (se := SubExperiment(EXPERIMENT_PATH, BASE_PATH, NAMESPACE, globals())):

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

    @se.hook('fit_model')
    def fit_model(e, model, x_train, y_train, x_test, y_test):
        history = model.fit(
            x_train,
            y_train,
            batch_size=e.parameters['BATCH_SIZE'],
            epochs=e.parameters['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
        )
        return history.history


    @se.hook('query_model')
    def query_model(e, model, x, y, include_importances: bool = True):
        e.info('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


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

        # ~ 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.parameters['CHANNEL_DIRECTIONS'][k] * (out - out_masked)

            fidelities.append(fidelity)

        return fidelities

