import tensorflow as tf
import numpy as np
from tensorflow import keras
from keras.layers import Input, Lambda, Dense, Flatten, Rescaling
from keras.models import Model
import PIL
from PIL import Image
import gradio as gr
import matplotlib.cm as cm

base_model = keras.applications.Xception(
    # weights = "../input/xception/xception_weights_tf_dim_ordering_tf_kernels_notop.h5",
    input_shape = (160,160,3),
    include_top = False,)

base_model.trainable = False


def img_pros(img):
  img = tf.keras.preprocessing.image.img_to_array(img)
  img = tf.image.resize(img, [160,160])
  img = tf.expand_dims(img, axis = 0)
  return img

#function for creating model
#returns model, its inputs, Xception's last conv output, the whole model's outputs
def create_model_mod(classes, activation):
    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(classes ,activation = activation)(gap)
    model = keras.Model(inputs, outputs)
    
    if activation == "linear":
        loss_s = keras.losses.BinaryCrossentropy(from_logits = True)
    else:
        loss_s = keras.losses.BinaryCrossentropy()
    
    model.compile(
    loss = loss_s,
    optimizer = keras.optimizers.Adam(0.001),
    metrics = ["accuracy"]
    )
    
    return model, inputs, x, outputs


#function that creates a gradcam model and returns it
def create_grad_model(weights, classes, activation):
    model_mod, input, x, output = create_model_mod(classes, activation)
    #lodaing weights of already trained model
    model_mod.load_weights(weights)
    grad_model = Model(input, [x, output])
    return grad_model
    
#create heatmaps of the given images
#returns the heatmaps and the raw score of predicted class of each image
def create_heatmap(model, imgs, class_index):
    model.layers[-1].activation = None
    
    #predicting the images and getting the conv outputs and predictions from the gradcam model
    with tf.GradientTape() as tape:
        maps, preds = model(imgs);
    
        # class_channel = tf.expand_dims(preds[:,class_index],axis = 1)
        class_channel = preds[:, class_index]
    
    #computing gradients of predictions w.r.t the feature maps
    if class_index == -1:
        grads = tape.gradient(preds, maps)
    else:
        grads = tape.gradient(class_channel, 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()


#superimpose function buth for a single input image
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

#for generating single gradcam image
def gen_grad_img_single(grad_model, img, class_index, alpha = 0.4):
    heatmaps, y_pred = create_heatmap(grad_model, img, class_index)
    
    # 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


def gen_grad_both(grad_model, img):
    img_c, y_pred_c = gen_grad_img_single(grad_model, img, 0)
    img_d, y_pred_d = gen_grad_img_single(grad_model, img, 1)
    y_pred_c = np.around(y_pred_c,3)
    y_pred_d = np.around(y_pred_d,3)
    # show_imgs([img_c, img_d], [y_true, y_true], [size[0], size[1]], cols, [y_pred_c, y_pred_d], font_size = font_size)
    
    infer = ""
    if y_pred_c[0][0] > y_pred_c[0][1]: infer = "cat"
    else: infer = "dog"
        
    return img_c, img_d, y_pred_c, infer

weights = "weights_nm.h5"

def get_grad(img):
    img = img_pros(img)
    grad_model = create_grad_model(weights, 2, "softmax")
    grad_img_c, grad_img_d, y_pred, infer = gen_grad_both(grad_model, img)
    # pred_class = ""
    # if y_pred[0] > 0.5: pred_class = "cat"
    # else: pred_class = "dog"
        
    text = "Raw Score: " + str(y_pred[0]) + "\nClassification: " + infer
    return grad_img_c, grad_img_d, text

# demo = gr.Interface(
#     fn = get_grad,
#     with gr.Row():
#         inputs = gr.Image(type = "pil", shape = (224,224), width = 320, height = 320),
#         outputs = [gr.Image(type = "numpy", width = 320, height = 320, label = "Grad_CAM w.r.t cat"), gr.Image(type = "numpy", width = 320, height = 320, label = "Grad_CAM w.r.t dog"), gr.Textbox(label = 'Prediction', info = '[P of cat, P of dog]')], 
#     description = "Visual Explanations from Deep Networks",
#     title = "Gradient-Weighted Class Activation Mapping (Grad-CAM)"
# )

with gr.Blocks() as demo:
    with gr.Row().style(equal_height=True):
        img_input = gr.Image(type = "pil", shape = (224,224), width = 320, height = 320)
        img_output1 = gr.Image(type = "numpy", width = 320, height = 320, label = "Grad_CAM w.r.t cat")
        img_output2 = gr.Image(type = "numpy", width = 320, height = 320, label = "Grad_CAM w.r.t dog")
        

demo.launch()