'''
Grad-CAM visualization utilities

- Based on https://keras.io/examples/vision/grad_cam/

---
- 2021-12-18 jkang first created 
- 2022-01-16 
    - copied from https://huggingface.co/spaces/jkang/demo-gradcam-imagenet/blob/main/utils.py
    - updated for artis/trend classifier
'''
import matplotlib.cm as cm

import os
import re
from glob import glob
import numpy as np
import tensorflow as tf
tfk = tf.keras
K = tfk.backend

# Disable GPU for testing
# os.environ['CUDA_VISIBLE_DEVICES'] = '-1'


def get_imagenet_classes():
    '''Retrieve all 1000 imagenet classes/labels as dictionaries'''
    classes = tfk.applications.imagenet_utils.decode_predictions(
        np.expand_dims(np.arange(1000), 0), top=1000
    )
    idx2lab = {cla[2]: cla[1] for cla in classes[0]}
    lab2idx = {idx2lab[idx]: idx for idx in idx2lab}
    return idx2lab, lab2idx


def search_by_name(str_part):
    '''Search imagenet class by partial matching string'''
    results = [key for key in list(lab2idx.keys()) if re.search(str_part, key)]
    if len(results) != 0:
        return [(key, lab2idx[key]) for key in results]
    else:
        return []


def get_xception_model():
    '''Get model to use'''
    base_model = tfk.applications.xception.Xception
    preprocessor = tfk.applications.xception.preprocess_input
    decode_predictions = tfk.applications.xception.decode_predictions
    last_conv_layer_name = "block14_sepconv2_act"

    model = base_model(weights='imagenet')
    grad_model = tfk.models.Model(
        inputs=[model.inputs],
        outputs=[model.get_layer(last_conv_layer_name).output,
                 model.output]
    )
    return model, grad_model, preprocessor, decode_predictions


def get_img_4d_array(image_file, image_size=(299, 299)):
    '''Load image as 4d array'''
    img = tfk.preprocessing.image.load_img(
        image_file, target_size=image_size)  # PIL obj
    img_array = tfk.preprocessing.image.img_to_array(
        img)  # float32 numpy array
    img_array = np.expand_dims(img_array, axis=0)  # 3d -> 4d (1,299,299,3)
    return img_array


def make_gradcam_heatmap(grad_model, img_array, pred_idx=None):
    '''Generate heatmap to overlay with
    - img_array: 4d numpy array
    - pred_idx: eg. index out of 1000 imagenet classes
        if None, argmax is chosen from prediction
    '''
    # Get gradient of pred class w.r.t. last conv activation
    with tf.GradientTape() as tape:
        last_conv_act, predictions = grad_model(img_array)
        if pred_idx == None:
            pred_idx = tf.argmax(predictions[0])
        class_channel = predictions[:, pred_idx]  # (1,1000) => (1,)

    # d(class_channel/last_conv_act)
    grads = tape.gradient(class_channel, last_conv_act)
    pooled_grads = tf.reduce_mean(grads, axis=(
        0, 1, 2))  # (1,10,10,2048) => (2048,)

    # (10,10,2048) x (2048,1) => (10,10,1)
    heatmap = last_conv_act[0] @ pooled_grads[..., tf.newaxis]
    heatmap = tf.squeeze(heatmap)  # (10,10)

    # Normalize heatmap between 0 and 1
    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
    return heatmap, pred_idx.numpy(), predictions.numpy().squeeze()


def align_image_with_heatmap(img_array, heatmap, alpha=0.3, cmap='jet'):
    '''Align the image with gradcam heatmap
    - img_array: 4d numpy array
    - heatmap: output of `def make_gradcam_heatmap()` as 2d numpy array
    '''
    img_array = img_array.squeeze()  # 4d => 3d

    # Rescale to 0-255 range
    heatmap_scaled = np.uint8(255 * heatmap)
    img_array_scaled = np.uint8(255 * img_array)

    colormap = cm.get_cmap(cmap)
    colors = colormap(np.arange(256))[:, :3]  # mapping RGB to heatmap
    heatmap_colored = colors[heatmap_scaled]  # ? still unclear

    # Make RGB colorized heatmap
    heatmap_colored = (tfk.preprocessing.image.array_to_img(heatmap_colored)  # array => PIL
                       .resize((img_array.shape[1], img_array.shape[0])))
    heatmap_colored = tfk.preprocessing.image.img_to_array(
        heatmap_colored)  # PIL => array

    # Overlay image with heatmap
    overlaid_img = heatmap_colored * alpha + img_array_scaled
    overlaid_img = tfk.preprocessing.image.array_to_img(overlaid_img)
    return overlaid_img


if __name__ == '__main__':
    # Test GradCAM
    examples = sorted(glob(os.path.join('examples', '*.jpg')))
    idx2lab, lab2idx = get_imagenet_classes()

    model, grad_model, preprocessor, decode_predictions = get_xception_model()

    img_4d_array = get_img_4d_array(examples[0])
    img_4d_array = preprocessor(img_4d_array)

    heatmap = make_gradcam_heatmap(grad_model, img_4d_array, pred_idx=None)

    img_pil = align_image_with_heatmap(
        img_4d_array, heatmap, alpha=0.3, cmap='jet')

    img_pil.save('test.jpg')
    print('done')