import os
import numpy as np
import cv2
import imageio
from collections import defaultdict

import tensorflow as tf


def _raw_heatmap_image(img_tensor, model):
    """ gets gradients of predictions w.r.t final Conv2D layer
        and processes them into a heatmap """
    conv_layer = model.get_layer("block5_conv3")
    heatmap_model = tf.keras.models.Model([model.inputs], [conv_layer.output, model.output])

    with tf.GradientTape() as gtape:
        conv_output, predictions = heatmap_model(img_tensor)
        loss = predictions[:, np.argmax(predictions[0])]
        grads = gtape.gradient(loss, conv_output)
        pooled_grads = tf.keras.backend.mean(grads, axis=(0, 1, 2))

    heatmap = tf.reduce_mean(tf.multiply(pooled_grads, conv_output), axis=-1)
    heatmap = np.maximum(heatmap, 0)
    max_heat = np.max(heatmap)
    if max_heat == 0:
        max_heat = 1e-10
    heatmap /= max_heat
    heatmap = heatmap[0] # unpack along batch dim
    return heatmap

def _preprocess_image(img):
    """ outputs a processed image and an image tensor"""
    #img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    #img = cv2.resize(img, (224, 224)) 
    img = tf.keras.preprocessing.image.load_img(img, target_size=(224, 224))
    img_tensor = tf.keras.preprocessing.image.img_to_array(img)
    img_tensor = np.expand_dims(img_tensor, axis=0)
    img_tensor = img_tensor/255.
    return img_tensor, img

def super_imposed_image(img, heatmap, weight, threshold):
    """ Super imposes the heatmap onto the original image
        Weight: float between 0 and 1 for how heavily to use the heatmap
        Threshold: float between 0 and 1 for a threhold value of the heatmap to use
                    eg discards heatmap pixels below threshold """ 
    map_img = heatmap*(heatmap > threshold)
    map_img = cv2.resize(heatmap, (224, 224))
    map_img = np.uint8(255*map_img)
    map_img = cv2.applyColorMap(map_img, cv2.COLORMAP_JET)

    superimposed_img = np.uint8(weight*map_img + img)
    return superimposed_img

def get_frames_from_video(vid_path):
    """ Utility function to get frames from a given video """
    video_cap = cv2.VideoCapture(vid_path)
    count = 0
    images_out = []

    while video_cap.isOpened() and count < 500:
        success, img = video_cap.read()
        if success:
            if count%10==0:
                images_out.append(img)
            count += 1
        else:
            break
    cv2.destroyAllWindows()
    video_cap.release()
    return images_out

def video_to_heatmap_gif(vid_path, model, path_out='cul_eight.gif', duration=0.2, weight=0.15, threshold=0.75):
    """ Saves an animation of images processed with heatmap generated by the model """
    images = get_frames_from_video(vid_path)
    heatmap_applied_images = []
    for img in images:
        gen_image = heatmap_image(img, model, weight, threshold)
        heatmap_applied_images.append(gen_image)
  
    imageio.mimsave(path_out, heatmap_applied_images, duration=duration)
    return True

def heatmap_image(img, model, weight=0.2, threshold=0.75):
    """Overlays a heatmap generated by the image onto the image"""
    img_tensor, processed_img = _preprocess_image(img)
    heatmap = _raw_heatmap_image(img_tensor, model)
    applied = super_imposed_image(processed_img, heatmap, weight, threshold)
    return applied


if __name__ == '__main__':
    model = tf.keras.models.load_model('.\saved_models\rowing_1213.h5')
    video_path = 'cul_eight.MOV'
    video_to_heatmap_gif(video_path, model)