import tensorflow as tf
import matplotlib.pyplot as plt
from matplotlib.pyplot import imread
import numpy as np
from tensorflow import keras
from keras.layers import Input, Lambda, Dense, Flatten, Rescaling
from keras.models import Model
from sklearn.metrics import confusion_matrix
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn.metrics import auc
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import classification_report
from sklearn.metrics import roc_auc_score
from sklearn.metrics import roc_curve
from sklearn.metrics import RocCurveDisplay
from os import listdir, path
from skimage.util import random_noise
from skimage import img_as_float
import math
import cv2
import os
import PIL
import shutil
import matplotlib.cm as cm


#function for creating model
#returns model, its inputs, Xception's last conv output, the whole model's outputs
def create_model_mod():
    inputs = keras.Input(shape = (160,160,3))
    #normalizing pixel values
    r = Rescaling(scale = 1./255)(inputs)
    x = base_model(r, training = False)
    gap = keras.layers.GlobalAveragePooling2D()(x)
    outputs = keras.layers.Dense(1,activation = 'linear')(gap)
    model = keras.Model(inputs, outputs)
    
    model.compile(
    loss = keras.losses.BinaryCrossentropy(from_logits = True),
    optimizer = keras.optimizers.Adam(0.001),
    metrics = ["accuracy"]
    )
    
    return model, inputs, x, outputs

def create_heatmap(model, imgs):
    #predicting the images and getting the conv outputs and predictions
    with tf.GradientTape() as tape:
        maps, preds = model(imgs);
    
    #computing gradients of predictions w.r.t the feature maps
    grads = tape.gradient(preds, maps)
    
    # global average pooling of each feature map
    gap_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    
    #multiplying each pooled value with its correponding feature map
    # maps = maps[0]
    heatmap = maps @ gap_grads[..., tf.newaxis]
    
    #removing the extra dimension of value 1
    heatmap = tf.squeeze(heatmap)
    
    #applying relu activation
    heatmap = tf.keras.activations.relu(heatmap)
    
    return heatmap, preds.numpy()

def superimpose_single(heatmap, img, alpha = 0.4):
    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.utils.array_to_img(jet_heatmap)
    jet_heatmap = jet_heatmap.resize((160,160))
    jet_heatmap = keras.utils.img_to_array(jet_heatmap)
    
    # Superimpose the heatmap on original image
    superimposed_img = jet_heatmap * alpha + img
    # superimposed_img = keras.utils.array_to_img(superimposed_img)
    
    return superimposed_img

def gen_grad_img_single(weights, img, alpha = 0.4):
    model_mod, input, x, output = create_model_mod()
    model_mod.load_weights(weights)
    grad_model = Model(input, [x, output])
    heatmaps, y_pred = create_heatmap(grad_model, img)
    
    for i in range(len(y_pred)):
      if y_pred[i] > 0.5: y_pred[i] = 1
      else: y_pred[i] = 0
    
    img = superimpose_single(heatmaps, img[0])
    return np.array(img).astype('uint8'), y_pred


weights = "weights.h5"
# img, y_pred = gen_grad_img_single(weights, img)

demo = gr.Interface(
  fn = gen_grad_img_single,
  inputs = gr.Image(type = "pil", shape = (224,224)),
  # outputs = [gr.outputs.Label(num_top_classes = 2, label = 'Classifiaction'), gr.Textbox('infer_time', label = 'Inference Time(ms)')]
  outputs = ["image", gr.Textbox('y_pred', label = 'Prediction')]
   )
demo.launch()