#  -------------------------------------------------------------------------------------------
#  Copyright (c) Microsoft Corporation. All rights reserved.
#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
#  -------------------------------------------------------------------------------------------

import math
from dataclasses import dataclass
from functools import partial
from typing import Any, Callable, Optional, Set, Tuple

import torch
import torch.nn as nn
from timm.models.layers import DropPath, Mlp, trunc_normal_


def torch_int_div(tensor1, tensor2):
    """
    A function that performs integer division across different versions of PyTorch.
    """
    return torch.div(tensor1, tensor2, rounding_mode="floor")

@dataclass
class MultiHeadAttentionOutput:
    mha_output: torch.Tensor
    attention: Optional[torch.Tensor] = None


class VisionTransformerPooler(nn.Module):
    """
    :param input_dim: Input feature dimension (i.e., channels in old CNN terminology)
    :param grid_shape: Shape of the grid of patches per image
    :param num_heads: Number of self-attention heads within the MHA block
    :param num_blocks: Number of blocks per attention layer
    :param norm_layer: Normalisation layer

    `self.type_embed`: Is used to characterise prior and current scans, and
                       create permutation variance across modalities/series.
    """

    def __init__(self,
                 input_dim: int,
                 grid_shape: Tuple[int, int],
                 num_heads: int = 8,
                 num_blocks: int = 3,
                 norm_layer: Any = partial(nn.LayerNorm, eps=1e-6)):
        super().__init__()

        block_kwargs = dict(dim=input_dim, num_heads=num_heads, mlp_ratio=1., drop=0.10, attn_drop=0.10,
                            drop_path=0.25, act_layer=nn.GELU, norm_layer=norm_layer)
        self.blocks = nn.ModuleList([Block(**block_kwargs) for _ in range(num_blocks)])
        self.norm_post = norm_layer(input_dim)
        self.grid_shape = grid_shape
        self.num_patches = grid_shape[0] * grid_shape[1]
        self.num_blocks = num_blocks

        # Temporal positional embeddings
        num_series: int = 2
        self.type_embed = nn.Parameter(torch.zeros(num_series, 1, input_dim))
        trunc_normal_(self.type_embed, std=.02)

        # Positional embeddings 1 x L x C (L: Sequence length, C: Feature dimension)
        self.pos_drop = nn.Dropout(p=0.10)
        pos_embed_class = SinePositionEmbedding(embedding_dim=input_dim // 2, normalize=True)
        pos_embed = pos_embed_class(mask=torch.ones([1, grid_shape[0], grid_shape[1]]))  # 1 x L x C
        self.register_buffer("pos_embed", pos_embed, persistent=False)

        # Initialisation
        self.apply(self._init_weights)

    def no_weight_decay(self) -> Set[str]:
        return {'type_embed'}

    def forward(self, current_image: torch.Tensor, previous_image: Optional[torch.Tensor] = None) -> torch.Tensor:
        B, C, H, W = current_image.shape
        assert H == self.grid_shape[0] and W == self.grid_shape[1], "Input and grid shapes do not match"

        # Flatten patch embeddings to have shape (B x L x C), L = H * W
        if previous_image is not None:
            assert previous_image.shape == current_image.shape, "current_image and previous_image shapes do not match"
            previous_image = previous_image.view(B, C, H * W).transpose(1, 2)
        current_image = current_image.view(B, C, H * W).transpose(1, 2)
        pos_embed = self.pos_embed.repeat(B, 1, 1)  # type: ignore

        # Final token activations (B x 2L x C)
        token_features = self.forward_after_reshape(x=current_image, pos_embed=pos_embed, x_previous=previous_image)

        # Extract the patch features of current image
        cur_img_token_id = 0
        current_token_features = token_features[:, cur_img_token_id:self.num_patches+cur_img_token_id]
        current_patch_features = current_token_features.transpose(1, 2).view(B, C, H, W)

        return current_patch_features

    def forward_after_reshape(self,
                              x: torch.Tensor,
                              pos_embed: torch.Tensor) -> torch.Tensor:
        B, L, _ = x.shape  # Batch, Sequence length, Feature dimension

        # Positional and type embeddings
        type_embed = self.type_embed[0].expand(B, L, -1)

        # Add positional and type embeddings (used in query and key matching)
        pos_and_type_embed = pos_embed + type_embed

        # Positional dropout
        x = self.pos_drop(x)

        # Multihead attention followed by MLP
        for block in self.blocks:
            x = block(x=x, pos_and_type_embed=pos_and_type_embed)
        x = self.norm_post(x)

        return x

    def _init_weights(self, m: nn.Module) -> None:
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)


class MultiHeadAttentionLayer(nn.Module):
    """
    Multi-head self attention module

    The content builds on top of the TIMM library (vision_transformer.py) and differs by the following:
        - Defines a custom `MultiHeadAttentionLayer` which does not only apply `self-attention` but it can be
            generalised to arbitrary (query, key, value) input tuples. This feature can be valuable to process
            more than 2 scans at a time.
        - `Self-attention` specific use-case can still be invoked by calling the `forward_as_mhsa` method.
    """

    def __init__(self,
                 dim: int,
                 num_heads: int = 8,
                 qkv_bias: bool = False,
                 attn_drop: float = 0.,
                 proj_drop: float = 0.) -> None:
        super().__init__()
        self.num_heads = num_heads
        assert dim % num_heads == 0, f"The embedding dim ({dim}) must be divisible by the number of heads ({num_heads})"
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5
        self.return_attention = False

        self.proj_q = nn.Linear(dim, dim, bias=qkv_bias)
        self.proj_k = nn.Linear(dim, dim, bias=qkv_bias)
        self.proj_v = nn.Linear(dim, dim, bias=qkv_bias)

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, k: torch.Tensor, q: torch.Tensor, v: torch.Tensor) -> MultiHeadAttentionOutput:
        B, N, C = v.shape
        assert C % self.num_heads == 0, \
            f"The embedding dim ({C}) must be divisible by the number of heads ({self.num_heads})"

        w_q = self.proj_q(q).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        w_k = self.proj_k(k).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        w_v = self.proj_v(v).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)

        attn = (w_q @ w_k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        o = (attn @ w_v).transpose(1, 2).reshape(B, N, C)
        o = self.proj(o)
        o = self.proj_drop(o)

        attention_output = attn if self.return_attention else None

        return MultiHeadAttentionOutput(mha_output=o, attention=attention_output)

    def forward_as_mhsa(self, input: torch.Tensor) -> MultiHeadAttentionOutput:
        return self(k=input, q=input, v=input)


class Block(nn.Module):
    """
    Encapsulates multi-layer perceptron and multi-head self attention modules into a block.

    The content builds on top of the TIMM library (vision_transformer.py) and differs by the following:
        - This implementation uses spatio-temporal positional embeddings instead of 2D positional embeddings only,
            and they are taken into account within the forward pass of each ViT block.
        - Utilises the custom defined `MultiHeadAttentionLayer` which does not apply `self-attention` only but can be
            generalised to arbitrary (query, key, value) tuples. This can be valuable to process more than 2 scans.

    Positional and type embeddings are handled in a similar fashion as DETR object localisation paper
    https://alcinos.github.io/detr_page/, where a fixed set of sine/cos positional embeddings are used
    in an additive manner to Q and K tensors.
    """

    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 1., qkv_bias: bool = False, drop: float = 0.,
                 attn_drop: float = 0., drop_path: float = 0., act_layer: Callable = nn.GELU,
                 norm_layer: Callable = nn.LayerNorm) -> None:
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = MultiHeadAttentionLayer(dim=dim, num_heads=num_heads, qkv_bias=qkv_bias,
                                            attn_drop=attn_drop, proj_drop=drop)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def with_pos_and_type_embed(self, tensor: torch.Tensor, emb: Optional[torch.Tensor]) -> torch.Tensor:
        # Add positional embeddings to key and query tensors
        return tensor if emb is None else tensor + emb

    def forward(self, x: torch.Tensor, pos_and_type_embed: Optional[torch.Tensor]) -> torch.Tensor:
        x_with_emb = self.with_pos_and_type_embed(self.norm1(x), emb=pos_and_type_embed)
        x = x + self.drop_path(self.attn.forward_as_mhsa(x_with_emb).mha_output)
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x


class SinePositionEmbedding():
    """
    This is a more standard version of the position embedding, very similar to the one used by the Attention is all you
    need paper, generalized to work on images.
    """

    def __init__(self,
                 embedding_dim: int = 64,
                 temperature: int = 10000,
                 normalize: bool = False,
                 scale: float = None) -> None:
        super().__init__()
        self.embedding_dim = embedding_dim
        self.temperature = temperature
        self.normalize = normalize
        if scale is not None and normalize is False:
            raise ValueError("normalize should be True if scale is passed")
        if scale is None:
            scale = 2 * math.pi
        self.scale = scale

    def __call__(self, mask: torch.Tensor) -> torch.Tensor:
        assert mask is not None, "No pixel mask provided"
        B, H, W = mask.shape
        y_embed = mask.cumsum(1, dtype=torch.float32)
        x_embed = mask.cumsum(2, dtype=torch.float32)
        if self.normalize:
            y_embed = y_embed / (y_embed[:, -1:, :] + 1e-6) * self.scale
            x_embed = x_embed / (x_embed[:, :, -1:] + 1e-6) * self.scale

        dim_t = torch.arange(self.embedding_dim, dtype=torch.float32)
        dim_t = self.temperature ** (2 * torch_int_div(dim_t, 2) / self.embedding_dim)

        pos_x = x_embed[:, :, :, None] / dim_t
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).view(B, H * W, self.embedding_dim * 2)

        return pos