import tensorflow as tf
from tensorflow import keras
from keras.models import load_model
from keras.preprocessing import image
import numpy as np
from PIL import Image

# Display
#from IPython.display import Image, display
#import matplotlib.pyplot as plt
import matplotlib.cm as cm

class ModeloCovid:

    #Constructor de clase, importa el modelo y guarda un diccionario con las etiquetas de clase
    def __init__(self):
        self.model = load_model("./modeloGrey9M.h5")
        self.num2label = {0:'COVID-19', 1:'NORMAL', 2:'Viral Pneumonia'}
        self.dims = self.model.input_shape[1:3]

    #CON UNA IMAGEN Y LAS DIMENSIONES LA LEE DE FORMA ADECUADA
    def _get_img_array(self, img_path, size):
        # `img` is a PIL image of size 299x299
        img = keras.preprocessing.image.load_img(img_path, target_size=size, color_mode="grayscale")
        # `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)
        array = np.expand_dims(array, axis=0)
        return array
    
    def _make_gradcam_heatmap(self,img_array, 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)
        return heatmap.numpy()
    
    
    #A partir de una imagen retorna su clase mas problable tras la prediccion del modelo	
    def predict(self,img_path):
        img = image.load_img(img_path,target_size=self.dims, color_mode="grayscale")
        img = np.asarray(img)/255.0
        img = np.expand_dims(img, axis=0)
        img = np.expand_dims(img, axis=3)

        output = self.model.predict(img)#Aqui esta el vector de probabilidades
        prediction = np.argmax(output, axis=1).astype(np.int)

        return self.num2label[prediction[0]],output
    def gradCam(self, img_path, cam_path="./static/GRADCAM_img/cam.jpg"):
        model_builder = keras.applications.xception.Xception
        preprocess_input = keras.applications.xception.preprocess_input
        decode_predictions = keras.applications.xception.decode_predictions
        last_conv_layer_name = "conv2d_9"
        
        img_array = preprocess_input(self._get_img_array(img_path, size=self.dims))
        # Remove last layer's softmax
        _model=self.model
        _model.layers[-1].activation = None
        
        heatmap = self._make_gradcam_heatmap(img_array, _model, last_conv_layer_name)
        superimposed_img = self._save_and_display_gradcam(img_path,heatmap,cam_path)
        return heatmap, superimposed_img
        
    def _save_and_display_gradcam(self,img_path, heatmap, cam_path, alpha=0.4):
        # Load the original image

        img = keras.preprocessing.image.load_img(img_path, color_mode="grayscale")
        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)

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

        # Save the superimposed image
        superimposed_img.save(cam_path)

        # Display Grad CAM
        return superimposed_img