"""
This example shows how to train a MEGAN model based on a chemistry dataset consisting of a CSV files
containing the SMILES representations of the molecules.

"""
import os
import warnings
warnings.filterwarnings("ignore")
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import csv
import random
from typing import List, Dict, Any

from pycomex.experiment import Experiment
from pycomex.util import Skippable

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow.keras as ks
from matplotlib.backends.backend_pdf import PdfPages
from kgcnn.data.moleculenet import OneHotEncoder
from kgcnn.data.utils import ragged_tensor_from_nested_numpy
from imageio.v2 import imread
from sklearn.metrics import r2_score

import graph_attention_student.typing as tc
from graph_attention_student.util import DATASETS_FOLDER
from graph_attention_student.data import create_molecule_eye_tracking_dataset
from graph_attention_student.data import load_eye_tracking_dataset
from graph_attention_student.data import load_eye_tracking_dataset_dict
from graph_attention_student.data import process_graph_dataset
from graph_attention_student.training import NoLoss
from graph_attention_student.training import LogProgressCallback
from graph_attention_student.models import MultiAttentionStudent
from graph_attention_student.models import Megan
from graph_attention_student.visualization import plot_node_importances
from graph_attention_student.visualization import plot_edge_importances

# == DATASET PARAMETERS ==
# [!] Insert the absolute path to the CSV file here
CSV_PATH = os.path.join(DATASETS_FOLDER, 'dataset-E.csv')
# CSV_PATH = os.path.join(DATASETS_FOLDER, 'Water', 'Water.csv')
CSV_PATH = os.path.join(DATASETS_FOLDER, 'Benzene', 'Benzene.csv')

# [!] Insert the column name of the smiles column here
SMILES_COLUMN = "SMILES"
# [!] Insert the column name of the target value here
TARGET_COLUMN = 'Solubility'
TRAIN_SPLIT = 0.80

# == MODEL PARAMETERS ==

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

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

# This is another hyperparameter of the explanation only train step. Usually should be between 1 and 10
IMPORTANCE_MULTIPLIER = 10

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

# The coefficient value of the explanation sparsity regularization that is applied to the network. Higher
# values should lead to sparser explanations.
SPARSITY_FACTOR = 1.0

# 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_LIMITS = [-6, +2]
REGRESSION_REFERENCE = 0

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

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

# == TRAINING PARAMETERS ==
DEVICE = 'gpu:0'
LOSS = ks.losses.MeanSquaredError()
LEARNING_RATE = 0.001
# LEARNING_RATE = 0.01
BATCH_SIZE = 32
EPOCHS = 500

# == EVALUATION PARAMETERS ==
LOG_STEP = 20
METRIC_KEY = 'mean_squared_error'
CHANNEL_NAMES = {
    0: 'Negative',
    1: 'Positive'
}
CHANNEL_DIRECTIONS = {
    0: -1,
    1: 1,
}

DEBUG = True
BASE_PATH = os.getcwd()
NAMESPACE = 'results/smiles_example'
with Skippable(), (e := Experiment(base_path=BASE_PATH, namespace=NAMESPACE, glob=globals())):

    # -- CONVERTING TO "EYE TRACKING DATASET"
    # First of all we convert the dataset into the "eye tracking" format. In this format, the dataset
    # is represented as a folder where each sample is associated with 2 files:
    # - One PNG file which visualizes that element
    # - One JSON file, which contains the graph representation of the molecule and other metadata
    folder_path = os.path.dirname(CSV_PATH)
    dataset_path = os.path.join(folder_path, 'eye_tracking_dataset')

    # We only do the conversion though if it does not already exists
    if not os.path.exists(dataset_path):
        e.info('starting to process dataset...')
        e.info('reading csv file...')
        molecule_infos: Dict[str, Dict[str, Any]] = {}
        with open(CSV_PATH, mode='r') as csv_file:
            reader = csv.DictReader(csv_file, delimiter=',', quotechar='"')
            for i, row in enumerate(reader):
                smile = row[SMILES_COLUMN]
                value = float(row[TARGET_COLUMN])
                element_id = str(i)

                molecule_info = {
                    'id': element_id,
                    'name': element_id,
                    'value': value,
                    'smiles': smile,
                    'input_type': 'regression',
                }
                molecule_infos[element_id] = molecule_info
                if i % LOG_STEP == 0:
                    e.info(f' * ({i:03d}) - smile: {smile}')

        e.info(f'creating eye tracking dataset in: {dataset_path}')
        os.mkdir(dataset_path)
        create_molecule_eye_tracking_dataset(
            molecule_infos=molecule_infos,
            dest_path=dataset_path,
            logger=e.logger,
            log_step=LOG_STEP,
            # This dictionary determines what kinds of chemical properties are going to be added as features
            # of the resulting dataset. There are two kinds of features: Ones associated with the
            # nodes (atoms) of the graph and the ones associated with the edges (bonds).
            # This dictionary will be passed almost exactly as it is to the "kgcnn.moleculenet.MoleculeNet"
            # class which does the actual interaction with RDKit, so look into that for more details about
            # what kinds of options are available.
            set_attributes_kwargs={
                'nodes': ['Symbol', 'TotalDegree', 'FormalCharge', 'NumRadicalElectrons', 'Hybridization',
                          'IsAromatic', 'IsInRing', 'TotalNumHs', 'CIPCode', 'ChiralityPossible',
                          'ChiralTag'],
                'encoder_nodes': {
                    'Symbol': OneHotEncoder(
                        ['B', 'C', 'N', 'O', 'F', 'Si', 'P', 'S', 'Cl', 'As', 'Se', 'Br', 'Te', 'I', 'At'],
                        dtype="str"
                    ),
                    'Hybridization': OneHotEncoder([2, 3, 4, 5, 6]),
                    'TotalDegree': OneHotEncoder([0, 1, 2, 3, 4, 5], add_unknown=False),
                    'TotalNumHs': OneHotEncoder([0, 1, 2, 3, 4], add_unknown=False),
                    'CIPCode': OneHotEncoder(['R', 'S'], add_unknown=False, dtype='str'),
                    'ChiralityPossible': OneHotEncoder(["1"], add_unknown=False, dtype='str'),
                },
                'edges': ['BondType', 'IsAromatic', 'IsConjugated', 'IsInRing', 'Stereo'],
                'encoder_edges': {
                    'BondType': OneHotEncoder([1, 2, 3, 12], add_unknown=False),
                    'Stereo': OneHotEncoder([0, 1, 2, 3], add_unknown=False)
                },
            }
        )
        e.info(f'created eye tracking dataset with {len(molecule_infos)} elements!')

    # -- LOADING DATASET
    # Now we need to load the dataset and do some pre-processing so that we can use it for model training
    e.info('loading dataset...')
    eye_tracking_dataset: tc.EtDatasetDict = load_eye_tracking_dataset_dict(dataset_path)
    e.info(f'loaded dataset with {len(eye_tracking_dataset)} elements')

    # the "eye_tracking_dataset" consists of dicts, which hold the information about the filepaths for the
    # image file, metadata file, the metadata dict itself, the graph dict. But for the training we only
    # need the graph dict, which is why we extract those here.
    dataset = []
    for c, (name, data) in enumerate(eye_tracking_dataset.items()):
        metadata = data['metadata']
        g = metadata['graph']

        # We need to fill this field with the target value! In the future for multiple regression targets
        # this may also be a numpy array / vector. The important thing is that the output shape of the
        # network matches the shape of this!
        g['graph_labels'] = np.array([data['metadata']['value']])
        metadata['index'] = c

        dataset.append(g)

    dataset_indices = list(range(len(dataset)))
    dataset_length = len(dataset)
    train_indices = random.sample(dataset_indices, k=int(TRAIN_SPLIT * len(dataset_indices)))
    test_indices = [index for index in dataset_indices if index not in train_indices]
    e['train_indices'] = train_indices
    e['test_indices'] = test_indices
    dataset_index_map = {d['metadata']['index']: d for d in eye_tracking_dataset.values()}

    # This turns the list of graph dicts into the final form which we need for the training of the model:
    # keras RaggedTensors which contain all the graphs.
    x_train, y_train, x_test, y_test = process_graph_dataset(dataset, test_indices)

    # -- TRAINING THE MODEL
    model = Megan(
        units=UNITS,
        final_units=FINAL_UNITS,
        use_bias=True,
        importance_factor=IMPORTANCE_FACTOR,
        sparsity_factor=SPARSITY_FACTOR,
        regression_limits=REGRESSION_LIMITS,
        regression_reference=REGRESSION_REFERENCE,
        dropout_rate=DROPOUT_RATE
    )

    model.compile(
        loss=[
            ks.losses.MeanSquaredError(),
            NoLoss(),
            NoLoss()
        ],
        loss_weights=[
            1,
            0,
            0
        ],
        metrics=[
            ks.metrics.MeanSquaredError(),
            ks.metrics.MeanAbsoluteError(),
        ],
        optimizer=ks.optimizers.Adam(learning_rate=LEARNING_RATE),
        run_eagerly=False
    )
    with tf.device('gpu:0'):
        history = model.fit(
            x_train,
            y_train,
            batch_size=BATCH_SIZE,
            epochs=EPOCHS,
            validation_data=(x_test, y_test),
            validation_freq=1,
            callbacks=LogProgressCallback(
                logger=e.logger,
                epoch_step=LOG_STEP,
                identifier=f'val_output_1_{METRIC_KEY}'
            ),
            verbose=0
        )

    e[f'history'] = history.history
    e[f'epochs'] = list(range(EPOCHS))
    param_count = model.count_params()
    e['param_count'] = param_count
    e.info(f'number of parameters: {param_count}')

    # -- EVALUATING TEST SET
    # Now we get to the reason why we converted the dataset into the "eye tracking" format first: This will
    # now simplify the visualization of the results a lot!
    # We now evaluate all the elements in the test set and visualize the results. We will do this by
    # creating one big PDF document which displays all the elements of the test set and then also displays
    # what kinds of explanations the model predicted for those.
    #e.info('Evaluating test set')

    e.info('Evaluating test set and drawing examples...')
    pdf_path = os.path.join(e.path, f'explanations.pdf')
    ncols = 1
    nrows = IMPORTANCE_CHANNELS
    with PdfPages(pdf_path) as pdf:

        e.info('generating test set predictions...')
        out_pred, ni_pred, ei_pred = [v.numpy() for v in model(x_test)]
        out_true, ni_true, ei_true = y_test

        # Saving the values for the current element to the internal experiment dictionary
        e.info('saving predictions...')
        for c, index in enumerate(test_indices):
            e[f'out/true/{index}'] = float(out_true[c][0])
            e[f'out/pred/{index}'] = float(out_pred[c][0])
            ni = ni_pred[c]
            ei = ei_pred[c]
            e[f'ni/{index}'] = ni
            e[f'ei/{index}'] = ei

        # Calculating the explanation fidelity* value
        if IMPORTANCE_FACTOR != 0:
            for k, direction in CHANNEL_DIRECTIONS.items():
                e.info(f'deviation for channel {k} with intended direction: {direction}')
                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, _, _ = model(x_test, node_importances_mask=masks_tensor)
                out_masked = out_masked.numpy()
                for c, index in enumerate(test_indices):
                    e[f'out_mod_fidelity/{index}/{k}'] = float(out_masked[c][0])

            e.info('calculating fidelity* values...')
            for index in test_indices:
                out = e[f'out/pred/{index}']
                fidelity = 0
                for k in CHANNEL_DIRECTIONS.keys():
                    out_ = e[f'out_mod_fidelity/{index}/{k}']
                    fidelity += CHANNEL_DIRECTIONS[k] * (out - out_)
                    e[f'fidelity/{index}'] = fidelity

        # Plotting the explanations for the current sample
        e.info('plotting the examples...')
        for c, index in enumerate(test_indices):
            data = dataset_index_map[index]
            g = data['metadata']['graph']
            image = np.asarray(imread(data['image_path']))
            fig, rows = plt.subplots(nrows=nrows, ncols=ncols, figsize=(8 * ncols, 8 * nrows), squeeze=False)
            fig.suptitle(f'Element {data["metadata"]["name"]}\n'
                         f'{data["metadata"]["smiles"]}\n'
                         f'True: {out_true[c]} - Pred: {out_pred[c]}\n'
                         f'Explanation Fidelity: {e[f"fidelity/{index}"]:.3f}')
            node_coordinates = np.array(g['node_coordinates'])
            for row_index in range(nrows):
                column_index = 0

                ax = rows[row_index][column_index]
                ax.set_title(f'fidelity deviation: {e[f"out_mod_fidelity/{index}/{row_index}"]:.3f}')
                ax.set_ylabel(f'Explanation Channel {row_index}:\n'
                              f'{CHANNEL_NAMES[row_index]}')
                ax.imshow(image, extent=(0, image.shape[0], 0, image.shape[1]))
                plot_node_importances(
                    g=g,
                    ax=ax,
                    vmax=np.max(ni),
                    node_importances=ni[:, row_index],
                    node_coordinates=node_coordinates
                )
                plot_edge_importances(
                    g=g,
                    ax=ax,
                    vmax=np.max(ei),
                    edge_importances=ei[:, row_index],
                    node_coordinates=node_coordinates
                )

            pdf.savefig(fig)
            plt.close(fig)

            if c % LOG_STEP == 0:
                e.info(f' * ({c}/{len(test_indices)})')


with Skippable(), e.analysis:

    # -- PLOTTING TRAINING PROCESS
    e.info('plotting training information...')
    pdf_path = os.path.join(e.path, 'training.pdf')
    with PdfPages(pdf_path) as pdf:
        epochs = list(range(EPOCHS))
        # Plotting the loss of the training and validation dataset
        fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(8, 8))
        fig.suptitle('Loss during Training')
        ax.set_ylabel('Loss')
        ax.set_xlabel('Epochs')
        ax.plot(epochs, e['history']['output_1_loss'],
                color='blue', ls='--', alpha=0.5, label='train')
        final_value = e['history/output_1_loss'][-1]
        ax.scatter(epochs[-1], final_value, color='blue', alpha=0.5, label=f'{final_value:.2f}')
        ax.plot(epochs, e['history']['val_output_1_loss'],
                color='blue', ls='-', alpha=1.0, label='test')
        final_value = e['history/val_output_1_loss'][-1]
        ax.scatter(epochs[-1], final_value, color='blue', alpha=1.0, label=f'{final_value:.2f}')
        ax.legend()
        pdf.savefig(fig)
        plt.close(fig)

        # Plotting the metric of the training and validation dataset
        fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(8, 8))
        fig.suptitle(f'"{METRIC_KEY}" during training')
        ax.set_ylabel(METRIC_KEY)
        ax.set_xlabel('Epochs')
        ax.plot(epochs, e[f'history/output_1_{METRIC_KEY}'],
                color='blue', ls='--', alpha=0.5, label='train')
        final_value = e[f'history/output_1_{METRIC_KEY}'][-1]
        ax.scatter(epochs[-1], final_value, color='blue', alpha=0.5, label=f'{final_value:.2f}')
        ax.plot(epochs, e[f'history/val_output_1_{METRIC_KEY}'],
                color='blue', ls='-', alpha=1.0, label='test')
        final_value = e[f'history/val_output_1_{METRIC_KEY}'][-1]
        ax.scatter(epochs[-1], final_value, color='blue', alpha=1.0, label=f'{final_value:.2f}')
        ax.legend()
        pdf.savefig(fig)
        plt.close(fig)

        # Calculating R2 Score for the test set and doing scatter plot
        r2 = r2_score(list(e['out/true'].values()), list(e['out/pred'].values()))
        fig, ax = plt.subplots(ncols=1, nrows=1, figsize=(8, 8))
        fig.suptitle(f'Scatter Plot - TEST SET\n'
                     f'R2: {r2:.3f}')
        ax.set_xlabel('True Value')
        ax.set_ylabel('Predicted Value')
        value_min = min(*e['out/true'].values(), *e['out/pred'].values())
        value_max = max(*e['out/true'].values(), *e['out/pred'].values())
        for index in e['out/true'].keys():
            ax.scatter(e[f'out/true/{index}'], e[f'out/pred/{index}'],
                       color='green', alpha=0.5)
        ax.plot([value_min, value_max], [value_min, value_max], color='black', alpha=0.5)
        pdf.savefig(fig)
        plt.close(fig)

