import numpy as np
import tensorflow as tf

from tensorflow.keras import backend as K, models

# Import DeepExplain
from deepexplain.tensorflow import DeepExplain


# Tested
def grad_cam(input_model, input_x, input_y, target_layer_name):
    """
    Compute the grad-cam for a text input
    """
    # shape output (batch_size, sequence_length, num_filters)
    # shape grads_val (batch_size, sequence_length, num_filters)
    output, grads_val = get_gradients(input_model, input_x, input_y, target_layer_name)

    # get the maximum gradient for each gram of words
    # shape (batch_size, sequence_length)
    weights = np.max(grads_val, axis=2)

    # shape cam (batch_size, sequence_length)
    cams = np.einsum('ijk,ij->ij', output, weights)

    # Process cam
    new_cams = np.transpose(np.transpose(cams) - np.mean(cams, axis=1))
    new_cams = np.maximum(new_cams, 0)
    max_value = np.max(new_cams, axis=1)
    max_value[max_value == 0] = 0.000001  # Replace the zero value by 0.000001
    new_cams = np.transpose(np.transpose(new_cams) / max_value)

    # shape cam (batch_size, sequence_length)
    return new_cams


# Tested
def saliency_map(input_model, input_x, input_y, target_layer_name):
    """
    Compute the saliency map for a text input
    """
    # shape output (batch_size, sequence_length, num_filters)
    # shape grads_val (batch_size, sequence_length, num_filters)
    _, grads_val = get_gradients(input_model, input_x, input_y, target_layer_name)

    # get the maximum gradient for each gram of words
    # shape (batch_size, sequence_length)
    s_maps = np.max(grads_val, axis=2)

    # Process cam
    new_cams = np.transpose(np.transpose(s_maps) - np.mean(s_maps, axis=1))
    new_cams = np.maximum(new_cams, 0)
    new_cams = np.transpose(np.transpose(new_cams) / np.max(new_cams, axis=1))

    return new_cams


# Tested
def get_gradients(input_model, input_x, input_y, target_layer_name):
    """
    Compute the gradients from the output layer to the target layer
    """
    # convert input_y from one-hot format into index format
    # input_y = np.argmax(input_y, axis=1)

    # get the output of the logits not the softmax
    output_y = tf.gather_nd(input_model.layers[-2].output, np.dstack([range(input_x.shape[0]), input_y])[0])

    conv_output = input_model.get_layer(target_layer_name).output

    # Calculate the gradients from the target class to the target layer
    grads = K.gradients(output_y, conv_output)[0]

    gradient_function = K.function([input_model.input, K.learning_phase()], [conv_output, grads])

    # shape output (batch_size, sequence_length, num_filters)
    # shape grads_val (batch_size, sequence_length, num_filters)
    output, grads_val = gradient_function([input_x, 0])

    return output, grads_val


# Tested
def word_level_interpretation(heatmap, window_size):
    """
    distribute the heatmap values from the filter level to the word level of the text input
    :param heatmap: the heatmap values of the input text                # shape [batch_size, sequence_length]
    :param window_size: the window size of the convolutional filter
    """
    l = list()
    sequence_length = np.shape(heatmap)[1]

    l.append(heatmap)
    for k in range(1, window_size):
        l.append(np.concatenate([heatmap[:, :k], heatmap[:, :sequence_length - k]], axis=1))
    word_level_heatmap = np.sum(l, axis=0) / window_size

    # shape [batch_size, sequence_length]
    return word_level_heatmap


# Tested
def unify_dim(tensor, window_size):
    """
    unify the tensor to the specific dimension length
    """
    t = np.zeros(shape=(tensor.shape[0], window_size - 1))
    return np.concatenate([tensor, t], axis=1)


# Tested
def get_interpretation(interpretation_method, input_model, input_x, input_y, window_size, target_layer_name):
    """
    Get the heatmap generated by interpretation method for each word per input text
    :param interpretation_method: the interpretation method to be used
    :param input_model:
    :param input_x:
    :param input_y:
    :param window_size:
    :param target_layer_name:
    """
    predicted_features = []
    for idx, window_size in enumerate(window_size):
        # use the interpretation method function to generate the heatmap (ex. grad_cam)
        if interpretation_method == 'grad_cam':
            heatmap = grad_cam(input_model, input_x, input_y, target_layer_name[idx])
        elif interpretation_method == 'saliency_map':
            heatmap = saliency_map(input_model, input_x, input_y, target_layer_name[idx])
        word_level_heatmap = word_level_interpretation(heatmap, window_size)
        word_level_heatmap = unify_dim(word_level_heatmap, window_size)
        predicted_features.append(word_level_heatmap)

    ans = np.sum(predicted_features, axis=0) / len(target_layer_name)

    return ans


# Tested
def get_interpretation_de(input_model_path, xs, y_train, num_classes, interpretation_method='grad*input'):
    """
    :param input_model_path: the labels for the input data
    :param interpretation_method: {'saliency', 'grad*input', 'intgrad', 'elrp', 'deeplift'}
    :param xs
    :param y_train
    :param num_classes
    """
    with DeepExplain(session=tf.compat.v1.keras.backend.get_session()) as de:
        new_model = models.load_model(input_model_path)

        # convert the index labels to one-hot format
        ys = np.zeros(shape=(len(y_train), num_classes))
        for i, label in enumerate(y_train):
            ys[i, label] = 1

        y_prediction = new_model.predict(xs)
        y_prediction = np.argmax(y_prediction, axis=1)

        # Interpretation ===============================================================================================

        # Evaluate the embedding tensor on the model input
        # shape [None, 50, 128]
        embedding_tensor = new_model.layers[1].output
        embedding_function = K.function([new_model.input], [embedding_tensor])
        embedding_output = embedding_function([xs])

        # Get tensor before the final activation
        # shape [batch_size, num_classes]
        logits_tensor = new_model.layers[-2].output

        # Run DeepExplain with the embedding as input
        heat_map = de.explain(interpretation_method, logits_tensor * ys, embedding_tensor, embedding_output[0])

        # sum the values for the embedding dimension to get the relevance value of each word
        interpretation_heatmap = np.sum(heat_map, axis=-1)
        interpretation_heatmap = np.maximum(interpretation_heatmap, 0)

        return interpretation_heatmap, y_prediction
