# -*- coding: utf-8 -*-
"""grad_cam.ipynb

**
 * This file is part of Hybrid CNN-LSTM for COVID-19 Severity Score Prediction paper.
 *
 * Written by Ankan Ghosh Dastider and Farhan Sadik.
 *
 * Copyright (c) by the authors under Apache-2.0 License. Some rights reserved, see LICENSE.
 */
 
"""
import tensorflow as tf
import keras
print('tensorflow version: ', tf.__version__)
print('keras version: ', keras.__version__)
import numpy as np
import cv2

#For images
from keras.preprocessing.image import array_to_img, img_to_array, load_img
#For model loading
from keras.models import load_model
#For Grad-CAM calculation
from tensorflow.keras import models
import tensorflow as tf

IMAGE_SIZE  = (128, 128)

def grad_cam(input_model, x, layer_name):
    """
    Args: 
        input_model(object):Model object
        x(ndarray):image
        layer_name(string):The name of the convolution layer
    Returns:
        output_image(ndarray):Colored image of the original image
    """

    #Image preprocessing
    #Since there is only one image to read, mode must be increased..I can't predict
    X = np.expand_dims(x, axis=0)
    preprocessed_input = X.astype('float32')
    # / 255.0    

    grad_model = models.Model([input_model.inputs], [input_model.get_layer(layer_name).output, input_model.output])
    
    with tf.GradientTape() as tape:
        conv_outputs, predictions = grad_model(preprocessed_input)
        class_idx = np.argmax(predictions[0])
        loss = predictions[:, class_idx]

    #Calculate the gradient
    output = conv_outputs[0]
    grads = tape.gradient(loss, conv_outputs)[0]

    gate_f = tf.cast(output > 0, 'float32')
    gate_r = tf.cast(grads > 0, 'float32')

    guided_grads = gate_f * gate_r * grads

    #Average the weights and multiply by the output of the layer
    weights = np.mean(guided_grads, axis=(0, 1))
    cam = np.dot(output, weights)

    #Scale the image to the same size as the original image
    cam = cv2.resize(cam, IMAGE_SIZE, cv2.INTER_LINEAR)
    #Instead of ReLU
    cam  = np.maximum(cam, 0)
    #Calculate heatmap
    heatmap = cam / cam.max()

    #Pseudo-color monochrome images
    jet_cam = cv2.applyColorMap(np.uint8(255.0*heatmap), cv2.COLORMAP_JET)
    #Convert to RGB
    rgb_cam = cv2.cvtColor(jet_cam, cv2.COLOR_BGR2RGB)
    #Combined with the original image
    output_image = (np.float32(rgb_cam) + x / 2)  
    
    return output_image

target_layer = 'concatenate_5'
cam = grad_cam(model, x, target_layer)
out_1=rgb2gray(z)*cam[:,:,0]
out_2=rgb2gray(z)*cam[:,:,1]
out_3=rgb2gray(z)*cam[:,:,2]
out=np.stack((out_1,out_2,out_3),axis=2)
out=out/255.