import tensorflow as tf
from tensorflow.keras.models import load_model, Model
import cv2
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.cm as cm

class GradCam:
    def __init__(self, model, img, last_conv_layer_name, pred_index=None):
        self.model = model
        self.img_path = img
        self.last_conv_layer_name = last_conv_layer_name


    def make_gradcam_heatmap(self, 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
        img_array= self.img_path
        grad_model = tf.keras.models.Model(
            [self.model.inputs], [self.model.get_layer(self.last_conv_layer_name).output, self.model.output]
        )

        # 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
        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)
        return heatmap.numpy()
    

    def save_and_display_gradcam(self, cam_path="cam.jpg", alpha=0.4):
        heatmap = self.make_gradcam_heatmap()
        # Load the original image
        
        img = self.img_path

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

        # Use the jet colormap to colorize the heatmap
        jet = cm.get_cmap("jet")
        jet_colors = jet(np.arange(512))[:, :3]
        jet_heatmap = jet_colors[heatmap]

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

        # Superimpose the heatmap on the original image
        superimposed_img = jet_heatmap * alpha + img
        superimposed_img = tf.keras.preprocessing.image.array_to_img(superimposed_img)

        # Save and display the image
        superimposed_img.save(cam_path)
        plt.imshow(superimposed_img)
        plt.axis('off')
        plt.show()
        plt.savefig(path)