from numpy.core.fromnumeric import size
from base_explainer import BaseExplainer
import tensorflow.keras as keras
import tensorflow as tf
import numpy as np
from PIL import Image
import matplotlib.cm as cm
import cv2

class GradCAMExplainer(BaseExplainer):

    #implementacao do metodo abstrato
    def get_explanation(self, img, model, img_size, props, preprocess_input = None, index=None):
        #clona o modelo e remove a softmax da ultima camada
        clone = tf.keras.models.clone_model(model)
        clone.layers[-1].activation = None
        
        #cria modelo grad-cam
        grad_model = tf.keras.models.Model([clone.inputs], [clone.get_layer(props["conv_layer"]).output, clone.output])
        #transforma a imagem em array
        img_array = self.get_img_array(img, size = img_size)
        #pre processa a imagem
        if preprocess_input:
            img_procecessed_array = preprocess_input(img_array)
        else:
            img_procecessed_array = img_array
        #faz o heatmap
        heatmap = self.__make_gradcam_heatmap(img_procecessed_array, grad_model, props["conv_layer"], pred_index=index)
        #poe o heatmap na imagem
        heat, mask = self.__save_and_display_gradcam(img, heatmap)

        return keras.preprocessing.image.array_to_img(heat)


    #transforma a imagem em array
    def get_img_array(self, img_path, size, expand=True):
        # `img` is a PIL image of size 299x299
        img = keras.preprocessing.image.load_img(img_path, target_size=size)
        # `array` is a float32 Numpy array of shape (299, 299, 3)
        array = keras.preprocessing.image.img_to_array(img)
        # We add a dimension to transform our array into a "batch"
        # of size (1, 299, 299, 3)
        if expand:
            array = np.expand_dims(array, axis=0)
        return array

    def __make_gradcam_heatmap(self, img_array, grad_model, last_conv_layer_name, pred_index=None):
        # First, we create a model that maps the input image to the activations
        # of the last conv layer as well as the output predictions
        #grad_model = tf.keras.models.Model(
            #[model.inputs], [model.get_layer(last_conv_layer_name).output, model.output]
        #)

        # Then, we compute the gradient of the top predicted class for our input image
        # with respect to the activations of the last conv layer
        with tf.GradientTape() as tape:
            last_conv_layer_output, preds = grad_model(img_array)
            if pred_index is None:
                pred_index = tf.argmax(preds[0])
            class_channel = preds[:, pred_index]

        # This is the gradient of the output neuron (top predicted or chosen)
        # with regard to the output feature map of the last conv layer
        grads = tape.gradient(class_channel, last_conv_layer_output)

        # This is a vector where each entry is the mean intensity of the gradient
        # over a specific feature map channel
        pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))

        # We multiply each channel in the feature map array
        # by "how important this channel is" with regard to the top predicted class
        # then sum all the channels to obtain the heatmap class activation
        last_conv_layer_output = last_conv_layer_output[0]
        heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
        heatmap = tf.squeeze(heatmap)

        # For visualization purpose, we will also normalize the heatmap between 0 & 1
        heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
        heatmap = heatmap.numpy()
        return heatmap

    def __save_and_display_gradcam(self, img_path, heatmap, cam_path="cam.jpg", alpha=0.4):
        # Load the original image
        img = keras.preprocessing.image.load_img(img_path)
        img = keras.preprocessing.image.img_to_array(img)

        # Rescale heatmap to a range 0-255
        heatmap = np.uint8(255 * heatmap)

        # Use jet colormap to colorize heatmap
        jet = cm.get_cmap("jet")

        # Use RGB values of the colormap
        jet_colors = jet(np.arange(256))[:, :3]
        jet_heatmap = jet_colors[heatmap]

        # Create an image with RGB colorized heatmap
        jet_heatmap = keras.preprocessing.image.array_to_img(jet_heatmap)
        jet_heatmap = jet_heatmap.resize((img.shape[1], img.shape[0]))
        jet_heatmap = keras.preprocessing.image.img_to_array(jet_heatmap)

        im = Image.fromarray(heatmap)
        im = im.resize((img.shape[1], img.shape[0]))
        
        im = np.asarray(im)
        im = np.where(im > 0, 1, im)


        # Superimpose the heatmap on original image
        superimposed_img = jet_heatmap * alpha + img
        
        return superimposed_img, im