import math
from typing import Any, Callable, Dict, Optional, Sequence, Set, Tuple, Type, Union, List

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.jit import Final

from timm.layers import Mlp
from timm.models.vision_transformer import LayerScale

from timm.layers import PatchEmbed, Mlp, DropPath, AttentionPoolLatent, RmsNorm, PatchDropout, SwiGLUPacked, \
    trunc_normal_, lecun_normal_, resample_patch_embed, resample_abs_pos_embed, use_fused_attn
from utils.drop_utils import drop
from utils.backward_utils import EntropyGrad, LoggingHookGrad



def stop_gradient(x, indices=None):
    """
        Stop graident for tokens that are dropped.
    """
    breakpoint()
    if len(x.shape) == 4: # attention
        mask = torch.zeros(x.shape[2]).to(x.device) # (T,)
        mask[indices] = 1
        mask = mask.view(1, 1, -1, 1)
    else: # feature (B, T, C)
        mask = torch.zeros(x.shape[1]).to(x.device) # (T,)
        mask[indices] = 1
        mask = mask.view(1, -1, 1)
    x = x.detach() * mask  + x * (1-mask)
    return x 



# Efficient implementation equivalent to the following:
def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, return_attention=False, return_cross_attention=False, weight_apply=False) -> torch.Tensor:
    # Efficient implementation equivalent to the following:
    L, S = query.size(-2), key.size(-2)
    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
    attn_bias = torch.zeros(L, S, dtype=query.dtype).to(query.device)
    if is_causal:
        assert attn_mask is None
        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0).to(qeury.device)
        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
        attn_bias.to(query.dtype)

    if attn_mask is not None:
        if attn_mask.dtype == torch.bool:
            attn_mask.masked_fill_(attn_mask.logical_not(), float("-inf"))
        else:
            attn_bias += attn_mask
    if return_cross_attention:
        cross_query = query.transpose(1,0).view(1, query.shape[1], -1, query.shape[-1])
        cross_key = key.transpose(1,0).view(1, key.shape[1], -1, key.shape[-1])
        cross_attn_weight = cross_query @ cross_key.transpose(-2, -1) * scale_factor
        cross_attn_weight = cross_attn_weight.view(query.shape[1], -1, L, cross_attn_weight.shape[-1]).transpose(1,0)
        normed_query = cross_query / cross_query.norm(dim=-1, keepdim=True)
        normed_key = cross_key / cross_key.norm(dim=-1, keepdim=True)
        normed_catt = normed_query @ normed_key.transpose(-2, -1)
        normed_catt = normed_catt.view(query.shape[1], -1, L, normed_catt.shape[-1]).transpose(1,0)

    else:
        cross_attn_weight = None
        normed_key = F.normalize(key, dim=-1)
        normed_query = F.normalize(query, dim=-1)
        normed_catt = normed_query @ normed_key.transpose(-2, -1)

    attn_weight = query @ key.transpose(-2, -1) * scale_factor
    attn_weight += attn_bias
    attn_logit = attn_weight
    attn_weight = torch.softmax(attn_weight, dim=-1)
    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
    
    # apply entropy regularization. 
    if weight_apply:
        attn_weight = EntropyGrad.apply(attn_weight, attn_weight)

    out = attn_weight @ value
    if return_attention:
        return out, attn_weight, cross_attn_weight, normed_catt, attn_logit
    else:
        return out


class Attention(nn.Module):
    fused_attn: Final[bool]

    def __init__(
            self,
            dim: int,
            num_heads: int = 8,
            qkv_bias: bool = False,
            qk_norm: bool = False,
            attn_drop: float = 0.,
            proj_drop: float = 0.,
            norm_layer: nn.Module = nn.LayerNorm,
    ) -> None:
        super().__init__()
        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.fused_attn = use_fused_attn()

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x: torch.Tensor, return_attention=False, return_cross_attention=False, weight_apply=False, fixed_tokens=None) -> torch.Tensor:
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        self.qkv_feature = qkv
        if self.qkv_feature.requires_grad:
            self.qkv_feature.retain_grad()
#        qkv = LoggingHookGrad.apply(qkv)
        q, k, v = qkv.unbind(0)
        q, k = self.q_norm(q), self.k_norm(k)
        if fixed_tokens is not None:
            q = stop_gradient(q, fixed_tokens)
            k = stop_gradient(k, fixed_tokens)
            v = stop_gradient(v, fixed_tokens)
        if self.fused_attn:
            x = scaled_dot_product_attention(
                q, k, v,
                dropout_p=self.attn_drop.p if self.training else 0.,
                return_attention=return_attention,
                return_cross_attention=return_cross_attention,
                weight_apply=weight_apply
            )
            if return_attention:
                x, attn, cross_attn, normed_catt, attn_logit = x
        else:
            q = q * self.scale
            attn = q @ k.transpose(-2, -1)
            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
            x = attn @ v
        
        x = x.transpose(1, 2).reshape(B, -1, C)
        x = self.proj(x)
        x = self.proj_drop(x)
#        x = EntropyGrad.apply(x, attn)
        if return_attention:
            return x, attn, cross_attn, normed_catt, attn_logit, (q.detach(), k.detach(), v.detach())
        return x

class Block(nn.Module):
    def __init__(
            self,
            dim: int,
            num_heads: int,
            mlp_ratio: float = 4.,
            qkv_bias: bool = False,
            qk_norm: bool = False,
            proj_drop: float = 0.,
            attn_drop: float = 0.,
            init_values: Optional[float] = None,
            drop_path: float = 0.,
            act_layer: nn.Module = nn.GELU,
            norm_layer: nn.Module = nn.LayerNorm,
            mlp_layer: nn.Module = Mlp,
    ) -> None:
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_norm=qk_norm,
            attn_drop=attn_drop,
            proj_drop=proj_drop,
            norm_layer=norm_layer,
        )
        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        self.norm2 = norm_layer(dim)
        self.mlp = mlp_layer(
            in_features=dim,
            hidden_features=int(dim * mlp_ratio),
            act_layer=act_layer,
            drop=proj_drop,
        )
        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()


    def forward(self, x: torch.Tensor, drop_tokens=0, return_attention=False, return_cross_attention=False, drop_option='drop', indices=None, weight_apply=False, return_features=False, fix_tokens=None) -> torch.Tensor:
        att = self.attn(self.norm1(x), return_attention=return_attention, return_cross_attention=return_cross_attention,weight_apply=weight_apply, fixed_tokens=fix_tokens)
        if return_attention:
            att, attn_weight, cross_attn_weight, normed_catt_weight, attn_logit, features = att 
        
        if fix_tokens is not None:
            x = stop_gradient(x, fix_tokens)

        x = x + self.drop_path1(self.ls1(att))
#        x = self.drop_path1(self.ls1(att))
        if drop_tokens != 0:
            # x (B, H, N, C)
            x, indices = drop(x.transpose(0,1), attn_weight, drop_tokens, drop_option, indices=indices)
            x = x.transpose(0, 1)
        else:
            indices = None
        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
        if return_attention:
            if return_features:
                return x, attn_weight, cross_attn_weight, normed_catt_weight, indices, attn_logit, features
            else:
                return x, attn_weight, cross_attn_weight, normed_catt_weight, indices, attn_logit, features

        return x


