from typing import Dict

import cv2
import numpy as np
import tensorflow as tf
from PIL import Image
from tensorflow import keras

RESOLUTION = 224

crop_layer = keras.layers.CenterCrop(RESOLUTION, RESOLUTION)
norm_layer = keras.layers.Normalization(
    mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
    variance=[(0.229 * 255) ** 2, (0.224 * 255) ** 2, (0.225 * 255) ** 2],
)
rescale_layer = keras.layers.Rescaling(scale=1.0 / 127.5, offset=-1)


def preprocess_image(orig_image: Image, model_type: str, size=RESOLUTION):
    """Image preprocessing utility."""
    # Turn the image into a numpy array and add batch dim.
    image = np.array(orig_image)
    image = tf.expand_dims(image, 0)

    # If model type is vit rescale the image to [-1, 1].
    if model_type == "original_vit":
        image = rescale_layer(image)

    # Resize the image using bicubic interpolation.
    resize_size = int((256 / 224) * size)
    image = tf.image.resize(image, (resize_size, resize_size), method="bicubic")

    # Crop the image.
    preprocessed_image = crop_layer(image)

    # If model type is DeiT or DINO normalize the image.
    if model_type != "original_vit":
        image = norm_layer(preprocessed_image)

    return orig_image, preprocessed_image.numpy()


def attention_rollout_map(
    image: Image, attention_score_dict: Dict[str, np.ndarray], model_type: str
):
    """Computes attention rollout results.

    Reference:
        https://arxiv.org/abs/2005.00928

    Code copied and modified from here:
        https://github.com/jeonsworld/ViT-pytorch/blob/main/visualize_attention_map.ipynb
    """
    num_cls_tokens = 2 if "distilled" in model_type else 1

    # Stack the individual attention matrices from individual transformer blocks.
    attn_mat = tf.stack(
        [attention_score_dict[k] for k in attention_score_dict.keys()]
    )
    attn_mat = tf.squeeze(attn_mat, axis=1)

    # Average the attention weights across all heads.
    attn_mat = tf.reduce_mean(attn_mat, axis=1)

    # To account for residual connections, we add an identity matrix to the
    # attention matrix and re-normalize the weights.
    residual_attn = tf.eye(attn_mat.shape[1])
    aug_attn_mat = attn_mat + residual_attn
    aug_attn_mat = (
        aug_attn_mat / tf.reduce_sum(aug_attn_mat, axis=-1)[..., None]
    )
    aug_attn_mat = aug_attn_mat.numpy()

    # Recursively multiply the weight matrices.
    joint_attentions = np.zeros(aug_attn_mat.shape)
    joint_attentions[0] = aug_attn_mat[0]

    for n in range(1, aug_attn_mat.shape[0]):
        joint_attentions[n] = np.matmul(
            aug_attn_mat[n], joint_attentions[n - 1]
        )

    # Attention from the output token to the input space.
    v = joint_attentions[-1]
    grid_size = int(np.sqrt(aug_attn_mat.shape[-1]))
    mask = v[0, num_cls_tokens:].reshape(grid_size, grid_size)
    mask = cv2.resize(mask / mask.max(), image.size)[..., np.newaxis]
    result = (mask * image).astype("uint8")
    return result