import matplotlib.cm as cm
import numpy as np
import tensorflow as tf

try:
    from tensorflow.keras.utils import array_to_img
    from tensorflow.keras.utils import img_to_array
except ImportError:
    from tensorflow.keras.preprocessing.image import array_to_img
    from tensorflow.keras.preprocessing.image import img_to_array


def process_image(img, size=(224, 224)):
    img_array = tf.keras.applications.imagenet_utils.preprocess_input(
        img, mode="torch"
    )
    img_array = tf.image.resize(img_array, size,)[
        None,
    ]
    return img_array


def get_gradcam_model(model):
    inp = tf.keras.Input(shape=(224, 224, 3))
    feats = model.forward_features(inp)
    preds = model.forward_head(feats)
    return tf.keras.models.Model(inp, [preds, feats])


def get_gradcam_prediction(
    img,
    grad_model,
    process=True,
    decode=True,
    pred_index=None,
    cmap="jet",
    alpha=0.4,
):
    """Grad-CAM for a single image

    Args:
        img (np.ndarray): raw image without batch_shape, e.g. (224, 224, 3)
        grad_model (tf.keras.Model): model with feature map and prediction
        process (bool, optional): imagenet pre-processing. Defaults to True.
        pred_index (int, optional): for particular calss. Defaults to None.
        cmap (str, optional): colormap. Defaults to 'jet'.
        alpha (float, optional): opacity. Defaults to 0.4.

    Returns:
        preds_decode: top5 predictions
        heatmap: gradcam heatmap
    """
    # process image for inference
    if process:
        img_array = process_image(img)
    else:
        img_array = tf.convert_to_tensor(img)[
            None,
        ]
        if img.min() != img.max():
            img = (img - img.min()) / (img.max() - img.min())
            img = np.uint8(img * 255.0)
    # get prediction
    with tf.GradientTape(persistent=True) as tape:
        preds, feats = grad_model(img_array)
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]
    # compute heatmap
    grads = tape.gradient(class_channel, feats)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    feats = feats[0]
    heatmap = feats @ pooled_grads[..., tf.newaxis]
    heatmap = tf.squeeze(heatmap)
    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
    heatmap = heatmap.numpy()
    heatmap = np.uint8(255 * heatmap)
    # colorize heatmap
    cmap = cm.get_cmap(cmap)
    colors = cmap(np.arange(256))[:, :3]
    heatmap = colors[heatmap]
    heatmap = array_to_img(heatmap)
    heatmap = heatmap.resize((img.shape[1], img.shape[0]))
    heatmap = img_to_array(heatmap)
    overlay = img + heatmap * alpha
    overlay = array_to_img(overlay)
    # decode prediction
    preds_decode = (
        tf.keras.applications.imagenet_utils.decode_predictions(preds.numpy())[
            0
        ]
        if decode
        else preds
    )
    return preds_decode, overlay
