""" Vision Transformer (ViT) in PyTorch
Hacked together by / Copyright 2020 Ross Wightman
"""
import torch
import torch.nn as nn
from einops import rearrange
from modules.layers_ours import *

from baselines.ViT.helpers import load_pretrained
from baselines.ViT.weight_init import trunc_normal_
from baselines.ViT.layer_helpers import to_2tuple


def _cfg(url='', **kwargs):
    return {
        'url': url,
        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
        'crop_pct': .9, 'interpolation': 'bicubic',
        'first_conv': 'patch_embed.proj', 'classifier': 'head',
        **kwargs
    }


default_cfgs = {
    # patch models
    'vit_small_patch16_224': _cfg(
        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/vit_small_p16_224-15ec54c9.pth',
    ),
    'vit_base_patch16_224': _cfg(
        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_p16_224-80ecf9dd.pth',
        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5),
    ),
    'vit_large_patch16_224': _cfg(
        url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_p16_224-4ee7a4dc.pth',
        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
}

def compute_rollout_attention(all_layer_matrices, start_layer=0):
    # all_layer_matrices List[layers x (batch, tokens, tokens)]
    # adding residual consideration
    num_tokens = all_layer_matrices[0].shape[1]
    batch_size = all_layer_matrices[0].shape[0]
    eye = torch.eye(num_tokens).expand(batch_size, num_tokens, num_tokens).to(all_layer_matrices[0].device) # (batch, tokens, tokens)
    all_layer_matrices = [all_layer_matrices[i] + eye for i in range(len(all_layer_matrices))] # List[layers x (batch, tokens, tokens)]
    # all_layer_matrices = [all_layer_matrices[i] / all_layer_matrices[i].sum(dim=-1, keepdim=True)
    #                       for i in range(len(all_layer_matrices))]
    joint_attention = all_layer_matrices[start_layer] # (batch, tokens, tokens)
    for i in range(start_layer+1, len(all_layer_matrices)):
        joint_attention = all_layer_matrices[i].bmm(joint_attention)
    return joint_attention

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = Linear(in_features, hidden_features)
        self.act = GELU()
        self.fc2 = Linear(hidden_features, out_features)
        self.drop = Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

    def relprop(self, cam, **kwargs):
        cam = self.drop.relprop(cam, **kwargs)
        cam = self.fc2.relprop(cam, **kwargs)
        cam = self.act.relprop(cam, **kwargs)
        cam = self.fc1.relprop(cam, **kwargs)
        return cam


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False,attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = head_dim ** -0.5

        # A = Q*K^T
        self.matmul1 = einsum('bhid,bhjd->bhij')
        # attn = A*V
        self.matmul2 = einsum('bhij,bhjd->bhid')

        self.qkv = Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = Dropout(attn_drop)
        self.proj = Linear(dim, dim)
        self.proj_drop = Dropout(proj_drop)
        self.softmax = Softmax(dim=-1)

        self.attn = None
        self.attn_cam = None
        self.v = None
        self.v_cam = None
        self.attn_gradients = None

    def get_attn(self):
        return self.attn

    def save_attn(self, attn):
        self.attn = attn

    def save_attn_cam(self, cam):
        self.attn_cam = cam

    def get_attn_cam(self):
        return self.attn_cam

    def get_v(self):
        return self.v

    def save_v(self, v):
        self.v = v

    def save_v_cam(self, cam):
        self.v_cam = cam

    def get_v_cam(self):
        return self.v_cam

    def save_attn_gradients(self, attn_gradients):
        self.attn_gradients = attn_gradients

    def get_attn_gradients(self):
        return self.attn_gradients

    def forward(self, x):
        b, n, _, h = *x.shape, self.num_heads
        qkv = self.qkv(x)
        q, k, v = rearrange(qkv, 'b n (qkv h d) -> qkv b h n d', qkv=3, h=h)

        self.save_v(v)

        dots = self.matmul1([q, k]) * self.scale

        attn = self.softmax(dots)
        attn = self.attn_drop(attn)

        self.save_attn(attn)
        attn.register_hook(self.save_attn_gradients)

        out = self.matmul2([attn, v]) # (batch, heads, tokens, head_dim)
        out = rearrange(out, 'b h n d -> b n (h d)') # (batch, tokens, embed_dim)

        out = self.proj(out) # (batch, tokens, embed_dim)
        out = self.proj_drop(out)
        return out

    def relprop(self, cam, **kwargs):
        cam = self.proj_drop.relprop(cam, **kwargs)
        cam = self.proj.relprop(cam, **kwargs)
        cam = rearrange(cam, 'b n (h d) -> b h n d', h=self.num_heads)

        # attn = A*V
        (cam1, cam_v)= self.matmul2.relprop(cam, **kwargs)
        cam1 /= 2
        cam_v /= 2

        self.save_v_cam(cam_v)
        self.save_attn_cam(cam1)

        cam1 = self.attn_drop.relprop(cam1, **kwargs)
        cam1 = self.softmax.relprop(cam1, **kwargs)

        # A = Q*K^T
        (cam_q, cam_k) = self.matmul1.relprop(cam1, **kwargs)
        cam_q /= 2
        cam_k /= 2

        cam_qkv = rearrange([cam_q, cam_k, cam_v], 'qkv b h n d -> b n (qkv h d)', qkv=3, h=self.num_heads)

        return self.qkv.relprop(cam_qkv, **kwargs)


class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.):
        super().__init__()
        self.norm1 = LayerNorm(dim, eps=1e-6)
        self.attn = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
        self.norm2 = LayerNorm(dim, eps=1e-6)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop)

        self.add1 = Add()
        self.add2 = Add()
        self.clone1 = Clone()
        self.clone2 = Clone()

        # self.x = None
        # self.mhsa = None

    def forward(self, x):
        # self.x = x.detach().clone() # SAVE X
        x1, x2 = self.clone1(x, 2)
        mhsa = self.attn(self.norm1(x2))
        # self.mhsa = mhsa.detach().clone() # SAVE MHSA
        x = self.add1([x1, mhsa])
        x1, x2 = self.clone2(x, 2)
        x = self.add2([x1, self.mlp(self.norm2(x2))])
        return x
        # x1, x2 = self.clone1(x, 2)
        # x = self.add1([x1, self.attn(self.norm1(x2))])
        # x1, x2 = self.clone2(x, 2)
        # x = self.add2([x1, self.mlp(self.norm2(x2))])
        # return x

    def relprop(self, cam, **kwargs):
        (cam1, cam2) = self.add2.relprop(cam, **kwargs)
        cam2 = self.mlp.relprop(cam2, **kwargs)
        cam2 = self.norm2.relprop(cam2, **kwargs)
        cam = self.clone2.relprop((cam1, cam2), **kwargs)

        (cam1, cam2) = self.add1.relprop(cam, **kwargs)
        cam2 = self.attn.relprop(cam2, **kwargs)
        cam2 = self.norm1.relprop(cam2, **kwargs)
        cam = self.clone1.relprop((cam1, cam2), **kwargs)
        return cam


class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x

    def relprop(self, cam, **kwargs):
        cam = cam.transpose(1,2)
        cam = cam.reshape(cam.shape[0], cam.shape[1],
                     (self.img_size[0] // self.patch_size[0]), (self.img_size[1] // self.patch_size[1]))
        return self.proj.relprop(cam, **kwargs)


class VisionTransformer(nn.Module):
    """ Vision Transformer with support for patch or hybrid CNN input stage
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=False, mlp_head=False, drop_rate=0., attn_drop_rate=0.):
        super().__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.patch_embed = PatchEmbed(
                img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

        self.blocks = nn.ModuleList([
            Block(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
                drop=drop_rate, attn_drop=attn_drop_rate)
            for i in range(depth)])

        self.norm = LayerNorm(embed_dim)
        if mlp_head:
            # paper diagram suggests 'MLP head', but results in 4M extra parameters vs paper
            self.head = Mlp(embed_dim, int(embed_dim * mlp_ratio), num_classes)
        else:
            # with a single Linear layer as head, the param count within rounding of paper
            self.head = Linear(embed_dim, num_classes)

        # FIXME not quite sure what the proper weight init is supposed to be,
        # normal / trunc normal w/ std == .02 similar to other Bert like transformers
        trunc_normal_(self.pos_embed, std=.02)  # embeddings same as weights?
        trunc_normal_(self.cls_token, std=.02)
        self.apply(self._init_weights)

        self.pool = IndexSelect()
        self.add = Add()

        self.inp_grad = None

    def save_inp_grad(self,grad):
        self.inp_grad = grad

    def get_inp_grad(self):
        return self.inp_grad


    def _init_weights(self, m):
        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)

    @property
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token'}

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.pool(x, dim=1, indices=torch.tensor(0, device=x.device)) # (1, 1, embed_dim)
        x = x.squeeze(1)
        x = self.head(x)
        return x
    
    def predmap_softmax(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        x = x.softmax(dim=-1) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        # weight = x.softmax(dim=-1)
        x = weight * x.softmax(dim=-1)
        # x = weight * x
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap2(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (B, tokens, embed_dim)
        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        #############################
        x = x * weight # (B, tokens, embed_dim) * (1, tokens, 1) = (B, tokens, embed_dim)


        x = self.head(x) # (1, tokens, classes)


        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap3(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        x = weight * x
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap4(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        weight = weight + 1 # (tokens, 1)
        x = weight * x
        #############################
        
        x = self.head(x) # (1, tokens, classes)

        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap5(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        x = weight * x.softmax(dim=-2) # SOFTMAX over tokens dim
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap6(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        # Apply softmax over tokens dim excluding CLS token
        weight = weight[1:, ...] # (tokens-1, 1)
        x = x[:, 1:, :] # (1, tokens-1, classes)
        x = weight * x.softmax(dim=-2) # (1, tokens-1, classes)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens-1)
        return x
    
    def predmap_temperature1_1(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 1.1
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap_temperature0_9(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 0.9
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap_temperature1_2(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 1.2
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap_temperature1_3(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 1.3
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap_temperature1_35(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 1.35
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap_temperature1_4(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 1.4
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap_temperature1_45(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 1.45
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap_temperature1_5(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 1.5
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap_temperature1_6(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        temperature = 1.6
        x = x / temperature
        x = weight * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap7(self, x, idx):
        input = x.clone()
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        from ViT_explanation_generator import Baselines, LRP
        from ViT_new import vit_base_patch16_224 as shchuna
        model = shchuna(pretrained=True).cuda()
        model.eval()
        baselines = Baselines(model)
        rollout = baselines.generate_rollout(input, start_layer=1) # (?, tokens-1, )
        rollout = rollout.unsqueeze(-1) # (?, tokens-1, 1)
        x = x[:, 1:, :] # Strip cls token # (1, tokens-1, classes)
        x = rollout / rollout.max(dim=-2, keepdims=True).values  * x.softmax(dim=-1)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        return x
    
    def predmap8(self, x, idx):
        input = x.clone()
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        # from ViT_explanation_generator import Baselines, LRP
        # from ViT_new import vit_base_patch16_224 as shchuna
        # model = shchuna(pretrained=True).cuda()
        # baselines = Baselines(model)
        blocks = self.blocks
        all_layer_attentions = []
        for blk in blocks:
            attn_heads = blk.attn.get_attn() # (B, H, N, N)
            avg_heads = (attn_heads.sum(dim=1) / attn_heads.shape[1]).detach() # (B, N, N)
            all_layer_attentions.append(avg_heads)
        all_layer_attentions = torch.stack(all_layer_attentions, dim=1) # (B, L, N, N)
        all_layer_attentions[..., 0, :] = all_layer_attentions[..., 0, :] + 1
        y = all_layer_attentions[..., 0, :] # (B, L, N)
        weight = y.prod(dim=1) # (B, N)
        weight = weight.unsqueeze(-1) # (B, N, 1)

        x = weight * x.softmax(dim=-1)
        x = x[:, 1:, :] # Strip cls token # (1, tokens-1, classes)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        return x
    
    def predmap9(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[...,0,:], dim=-1).item()

        ############################
        attn = self.blocks[-2].attn.get_attn() # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        # weight = x.softmax(dim=-1)
        x = weight * x.softmax(dim=-1)
        # x = weight * x
        #############################
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap9_all_layers(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        if idx is None:
            idx = torch.argmax(x[...,0,:], dim=-1).item()

        ############################
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attn = all_attn.mean(dim=1) # (blocks, tokens, tokens)
        all_attn = all_attn[:, 0, :] # (blocks, tokens)
        all_attn = all_attn.unsqueeze(-1) # (blocks, tokens, 1)
        x = all_attn * x.softmax(dim=-1) # (blocks, tokens, classes)

        #############################
        x = x[..., idx] # (blocks, tokens)
        x = x[..., 1:] # (blocks, tokens-1) Strip cls token
        return x
    
    def predmap10(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn1 = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        attn2 = self.blocks[-2].attn.get_attn() # (B, heads, tokens, tokens)
        attn = attn1 * attn2 # (B, heads, tokens, tokens)
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        # weight = x.softmax(dim=-1)
        x = weight * x.softmax(dim=-1)
        # x = weight * x
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap11(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        # self.blocks[-1].x # (B, tokens, embed_dim)
        # self.blocks[-1].mhsa # (B, tokens, embed_dim)
        x_norm_last = self.blocks[-1].x.flatten(1).norm(dim=-1) # (B, )
        mhsa_norm_last = self.blocks[-1].mhsa.flatten(1).norm(dim=-1) # (B, )
        ratio_last = mhsa_norm_last / x_norm_last

        x_norm_prev = self.blocks[-2].x.flatten(1).norm(dim=-1) # (B, )
        mhsa_norm_prev = self.blocks[-2].mhsa.flatten(1).norm(dim=-1) # (B, )
        ratio_prev = mhsa_norm_prev / x_norm_prev

        if ratio_prev > ratio_last:
            blk_idx = -2
        else:
            blk_idx = -1
        # print(f"{blk_idx} : {ratio_prev.item()} <> {ratio_last.item()}")

        attn_last = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        attn_prev = self.blocks[-2].attn.get_attn() # (B, heads, tokens, tokens)
        attn = ratio_last*attn_last + ratio_prev*attn_prev
        weight = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        weight = weight.unsqueeze(-1) # (tokens, 1)
        # weight = x.softmax(dim=-1)
        x = weight * x.softmax(dim=-1)
        # x = weight * x
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap12_attn_prev(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        attn = self.blocks[-2].attn.get_attn() # (B, heads, tokens, tokens)
        attn = attn[0].clamp(min=0).mean(dim=0)[0, :] # (tokens, )
        attn = attn[1:] # (tokens-1, ) Strip cls token
        return attn
    
    def attn_all_layers(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attn = all_attn.mean(dim=1) # (blocks, tokens, tokens)
        all_attn = all_attn[:, 0, :] # (blocks, tokens)
        x = all_attn # (blocks, tokens)
        #############################
        x = x[..., 1:] # (blocks, tokens-1) Strip cls token
        return x
    
    def predmap13(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-2].attn.get_attn() # (B, heads, tokens, tokens)
        attn = attn[0].clamp(min=0) # (heads, tokens, tokens)
        attn = attn[:, 0, :] # (heads, tokens, )
        attn = attn.unsqueeze(-1).unsqueeze(0) # (1, heads, tokens, 1)

        x = attn * x # (1, heads, tokens, classes)
        x = x.mean(dim=1) # (1, tokens, classes)
        x = x.softmax(dim=-1) # (1, tokens, classes)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap14(self, x, idx):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)

        ############################
        attn = self.blocks[-1].attn.get_attn() # (B, heads, tokens, tokens)
        attn = attn[0].clamp(min=0) # (heads, tokens, tokens)
        attn = attn[:, 0, :] # (heads, tokens, )
        attn = attn.unsqueeze(-1).unsqueeze(0) # (1, heads, tokens, 1)

        x = attn * x # (1, heads, tokens, classes)
        x = x.mean(dim=1) # (1, tokens, classes)
        x = x.softmax(dim=-1) # (1, tokens, classes)
        #############################
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        x = x[..., idx] # (1, tokens)
        x = x[..., 1:] # Strip cls token
        return x
    
    def predmap15(self, x, idx, return_extras=False):
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        
        if return_extras:
            extras["xpredmap"] = x.squeeze(0) # (tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        predmap = x[0, 1:, :] # (tokens-1, classes)

        projections = all_attnmap @ predmap # (blocks*heads, classes)
        # TODO: is this needed?
        projections = projections.softmax(dim=0) # (blocks*heads, classes)

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, classes)
        x = weighted_attnmap * predmap # (tokens-1, classes)
        #############################
        x = x[..., idx] # (tokens-1, )
        x = x.unsqueeze(0) # (1, tokens-1)
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, tokens-1)
            extras["projections"] = projections.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, classes)
            extras["weighted_attnmap"] = weighted_attnmap # (tokens-1, classes)
            extras["predmap"] = predmap # (tokens-1, classes)
            extras["CLS_self_attend"] = all_attn[..., 0, 0] # (blocks, heads)
            return x, extras
        return x # (1, tokens-1)
    
    
    def predmap15_layer(self, x, idx, return_extras=False, layer=-1):
        layer = len(self.blocks) + layer if layer < 0 else layer
        assert 0<=layer<len(self.blocks), f"Layer index out of range: {layer}"
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for l, blk in enumerate(self.blocks):
            x = blk(x)
            if l == layer:
                # calc predmap
                x_layer = x
                x_layer = self.norm(x_layer) # (1, tokens, embed_dim)
                x_layer = self.head(x_layer) # (1, tokens, classes)
                predmap = x_layer[0, 1:, :] # (tokens-1, classes)


        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        
        if return_extras:
            extras["xpredmap"] = x.squeeze(0) # (tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        projections = all_attnmap @ predmap # (blocks*heads, classes)
        # TODO: is this needed?
        projections = projections.softmax(dim=0) # (blocks*heads, classes)

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, classes)
        x = weighted_attnmap * predmap # (tokens-1, classes)
        #############################
        x = x[..., idx] # (tokens-1, )
        x = x.unsqueeze(0) # (1, tokens-1)
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, tokens-1)
            extras["projections"] = projections.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, classes)
            extras["weighted_attnmap"] = weighted_attnmap # (tokens-1, classes)
            extras["predmap"] = predmap # (tokens-1, classes)
            extras["CLS_self_attend"] = all_attn[..., 0, 0] # (blocks, heads)
            return x, extras
        return x # (1, tokens-1)    
    
    def predmap15_batched_layer(self, x, idx, return_extras=False, layer=-1, apply_softmax_classes=False, apply_softmax_tokens=False):
        layer = len(self.blocks) + layer if layer < 0 else layer
        assert 0<=layer<len(self.blocks), f"Layer index out of range: {layer}"
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for l, blk in enumerate(self.blocks):
            x = blk(x)
            if l == layer:
                # calc predmap
                x_layer = x
                x_layer = self.norm(x_layer) # (B, tokens, embed_dim)
                x_layer = self.head(x_layer) # (B, tokens, classes)
                predmap = x_layer[:, 1:, :] # (B, tokens-1, classes)
                if apply_softmax_classes:
                    predmap = predmap.softmax(dim=-1) # (B, tokens-1, classes)


        x = self.norm(x) # (B, tokens, embed_dim)
        x = self.head(x) # (B, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1) # (B, )
        
        if return_extras:
            extras["xpredmap"] = x # (B, tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn.clamp(min=0) # (B, blocks, heads, tokens, tokens)
        all_attnmap = all_attn[..., 0, 1:] # (B, blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(1,2) # (B, blocks*heads, tokens-1)

        projections = all_attnmap @ predmap # (B, blocks*heads, classes)
        # TODO: is this needed?
        projections = projections.softmax(dim=1) # (B, blocks*heads, classes)

        weighted_attnmap = all_attnmap.transpose(-1,-2) @ projections # (B, tokens-1, classes)
        x = weighted_attnmap * predmap # (B, tokens-1, classes)
        # if apply_softmax_classes:
            # x = x.softmax(dim=-1) # (B, tokens-1, classes)
        x = x[range(B), ..., idx] # (B, tokens-1, )
        if apply_softmax_tokens:
            x = x.softmax(dim=-1) # (B, tokens-1, )
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(1, (len(self.blocks), self.blocks[0].attn.num_heads)) # (B, blocks, heads, tokens-1)
            extras["projections"] = projections.unflatten(1, (len(self.blocks), self.blocks[0].attn.num_heads)) # (B, blocks, heads, classes)
            extras["weighted_attnmap"] = weighted_attnmap # (B, tokens-1, classes)
            extras["predmap"] = predmap # (B, tokens-1, classes)
            extras["CLS_self_attend"] = all_attn[..., 0, 0] # (B, blocks, heads)
            return x, extras
        return x # (B, tokens-1)
    
    
    def predmap15_softmax_classes_batched_layer(self, x, idx, return_extras=False, layer=-1):
        return self.predmap15_batched_layer(x, idx, return_extras, layer=layer, apply_softmax_classes=True)
    
    def predmap15_softmax_classes_tokens_batched_layer(self, x, idx, return_extras=False, layer=-1):
        return self.predmap15_batched_layer(x, idx, return_extras, layer=layer, apply_softmax_classes=True, apply_softmax_tokens=True)
    
    def predmap15_softmax_tokens_batched_layer(self, x, idx, return_extras=False, layer=-1):
        return self.predmap15_batched_layer(x, idx, return_extras, layer=layer, apply_softmax_tokens=True)

    def predmap16(self, x, idx, return_extras=False):
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        
        if return_extras:
            extras["xpredmap"] = x.squeeze(0) # (tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        predmap = x[0, 1:, :] # (tokens-1, classes)

        projections = all_attnmap @ predmap # (blocks*heads, classes)
        # Normalize
        # projections = projections / all_attnmap.norm(dim=-1, keepdim=True) # (blocks*heads, classes)
        # projections = projections / predmap.norm(dim=-2, keepdim=True) # (blocks*heads, classes)
        # # Set softmax temperature as in scaled-dot product attention
        # projections = projections / torch.sqrt(torch.tensor(predmap.shape[-2], dtype=torch.float32)) # (blocks*heads, classes)
        projections = projections.softmax(dim=0) # (blocks*heads, classes)
        
        cls_weight = 1-all_attn[:, :, 0, 0] # (blocks, heads)
        projections = projections * cls_weight.flatten(-2).unsqueeze(-1) # (blocks*heads, classes)

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, classes)
        x = weighted_attnmap * predmap # (tokens-1, classes)
        #############################
        x = x[..., idx] # (tokens-1, )
        x = x.unsqueeze(0) # (1, tokens-1)
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, tokens-1)
            extras["projections"] = projections.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, classes)
            extras["weighted_attnmap"] = weighted_attnmap # (tokens-1, classes)
            extras["predmap"] = predmap # (tokens-1, classes)
            return x, extras
        return x # (1, tokens-1)

    def predmap17(self, x, idx, return_extras=False):
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        
        if return_extras:
            extras["xpredmap"] = x.squeeze(0) # (tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        predmap = x[0, 1:, :] # (tokens-1, classes)

        projections = all_attnmap @ predmap # (blocks*heads, classes)
        projections = projections.softmax(dim=0) # (blocks*heads, classes)
        
        cls_weight = 1-all_attn[:, :, 0, 0] # (blocks, heads)
        projections = projections * cls_weight.flatten(-2).unsqueeze(-1) # (blocks*heads, classes)

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, classes)
        x = weighted_attnmap * predmap # (tokens-1, classes)
        #############################
        x = x[..., idx] # (tokens-1, )
        x = x.unsqueeze(0) # (1, tokens-1)
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, tokens-1)
            extras["projections"] = projections.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, classes)
            extras["weighted_attnmap"] = weighted_attnmap # (tokens-1, classes)
            extras["predmap"] = predmap # (tokens-1, classes)
            return x, extras
        return x # (1, tokens-1)

    def predmap18(self, x, idx, return_extras=False):
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        
        if return_extras:
            extras["xpredmap"] = x.squeeze(0) # (tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        predmap = x[0, 1:, :] # (tokens-1, classes)

        cls_weight = 1-all_attn[:, :, 0, 0] # (blocks, heads)
        cls_weight = cls_weight.flatten(-2).unsqueeze(-1) # (blocks*heads, 1)

        weighted_attnmap = all_attnmap.T @ cls_weight # (tokens-1, 1)
        x = weighted_attnmap * predmap # (tokens-1, classes)
        #############################
        x = x[..., idx] # (tokens-1, )
        x = x.unsqueeze(0) # (1, tokens-1)
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, tokens-1)
            extras["weighted_attnmap"] = weighted_attnmap # (tokens-1, classes)
            extras["predmap"] = predmap # (tokens-1, classes)
            return x, extras
        return x # (1, tokens-1)

    def predmap19(self, x, idx, return_extras=False):
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        
        if return_extras:
            extras["xpredmap"] = x.squeeze(0) # (tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)
        all_attnmap = all_attnmap / all_attnmap.sum(dim=-1, keepdim=True) # (blocks*heads, tokens-1)

        predmap = x[0, 1:, :] # (tokens-1, classes)

        projections = all_attnmap @ predmap # (blocks*heads, classes)
        projections = projections.softmax(dim=0) # (blocks*heads, classes)
        
        # cls_weight = 1/(all_attn[:, :, 0, 0]+1e-8) # (blocks, heads)
        # projections = projections * cls_weight.flatten(-2).unsqueeze(-1) # (blocks*heads, classes)

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, classes)
        x = weighted_attnmap * predmap # (tokens-1, classes)
        #############################
        x = x[..., idx] # (tokens-1, )
        x = x.unsqueeze(0) # (1, tokens-1)
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, tokens-1)
            extras["projections"] = projections.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, classes)
            extras["weighted_attnmap"] = weighted_attnmap # (tokens-1, classes)
            extras["predmap"] = predmap # (tokens-1, classes)
            return x, extras
        return x # (1, tokens-1)
    
    def predmap20(self, x, idx, return_extras=False, temperature=None):
        assert temperature is not None
        assert temperature > 0
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        
        if return_extras:
            extras["xpredmap"] = x.squeeze(0) # (tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        predmap = x[0, 1:, :] # (tokens-1, classes)

        projections = all_attnmap @ predmap # (blocks*heads, classes)
        # Normalize
        projections = projections / all_attnmap.norm(dim=-1, keepdim=True) # (blocks*heads, classes)
        projections = projections / predmap.norm(dim=-2, keepdim=True) # (blocks*heads, classes)
        # Set softmax temperature
        projections = projections / temperature # (blocks*heads, classes)
        projections = projections.softmax(dim=0) # (blocks*heads, classes)
        
        # cls_weight = 1-all_attn[:, :, 0, 0] # (blocks, heads)
        # projections = projections * cls_weight.flatten(-2).unsqueeze(-1) # (blocks*heads, classes)

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, classes)
        x = weighted_attnmap * predmap # (tokens-1, classes)
        #############################
        x = x[..., idx] # (tokens-1, )
        x = x.unsqueeze(0) # (1, tokens-1)
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, tokens-1)
            extras["projections"] = projections.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, classes)
            extras["weighted_attnmap"] = weighted_attnmap # (tokens-1, classes)
            extras["predmap"] = predmap # (tokens-1, classes)
            return x, extras
        return x # (1, tokens-1)
    
    def predmap21(self, x, idx, return_extras=False, temperature=None):
        assert temperature is not None
        assert temperature > 0
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()
        
        if return_extras:
            extras["xpredmap"] = x.squeeze(0) # (tokens, classes)

        ############################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        predmap = x[0, 1:, :] # (tokens-1, classes)

        projections = all_attnmap @ predmap # (blocks*heads, classes)
        # Normalize
        # projections = projections / all_attnmap.norm(dim=-1, keepdim=True) # (blocks*heads, classes)
        projections = projections / predmap.norm(dim=-2, keepdim=True) # (blocks*heads, classes)
        # Set softmax temperature
        projections = projections / temperature # (blocks*heads, classes)
        projections = projections.softmax(dim=0) # (blocks*heads, classes)
        
        # cls_weight = 1-all_attn[:, :, 0, 0] # (blocks, heads)
        # projections = projections * cls_weight.flatten(-2).unsqueeze(-1) # (blocks*heads, classes)

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, classes)
        x = weighted_attnmap * predmap # (tokens-1, classes)
        #############################
        x = x[..., idx] # (tokens-1, )
        x = x.unsqueeze(0) # (1, tokens-1)
        if return_extras:
            extras["all_attnmap"] = all_attnmap.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, tokens-1)
            extras["projections"] = projections.unflatten(0, (len(self.blocks), self.blocks[0].attn.num_heads)) # (blocks, heads, classes)
            extras["weighted_attnmap"] = weighted_attnmap # (tokens-1, classes)
            extras["predmap"] = predmap # (tokens-1, classes)
            return x, extras
        return x # (1, tokens-1)

    def predmap22(self, x, idx, return_extras=False, layer=-1):
        layer = len(self.blocks) + layer if layer < 0 else layer
        assert 0 <= layer < len(self.blocks)
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        for l,blk in enumerate(self.blocks):
            x = blk(x)
            if l == layer:
                x_layer = x

        x = self.norm(x) # (1, tokens, embed_dim)
        x = self.head(x) # (1, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1).item()

        x_layer = self.norm(x_layer) # (1, tokens, embed_dim)
        x_layer = self.head(x_layer) # (1, tokens, classes)
        
        predmap = x_layer[0, 1:, :] # (tokens-1, classes)
        predmap = predmap[..., idx] # (tokens-1, )
        predmap = predmap.unsqueeze(0) # (1, tokens-1)
        
        return predmap


    def predmap23(self, x, idx, return_extras=False, layer=-1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)

        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)

        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()
        x = x[..., idx] # (B, layers, tokens)

        # Strip CLS token
        x = x[..., 1:] # (B, layers, tokens-1)

        x = x.mean(dim=1) # (B, tokens-1)
        # x = x.max(dim=1) # (B, tokens-1)
        # x = x.prod(dim=1).sqrt(n_layers) # (B, tokens-1)
        # x = x.min(dim=1) # (B, tokens-1)

        return x  # (1, tokens-1)

    def predmap23_softmax_classes(self, x, idx, return_extras=False, layer=-1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)
        x = x.softmax(dim=-1) # (B*layers, tokens, classes)

        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)

        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()
        x = x[..., idx] # (B, layers, tokens)

        # Strip CLS token
        x = x[..., 1:] # (B, layers, tokens-1)

        x = x.mean(dim=1) # (B, tokens-1)
        # x = x.max(dim=1) # (B, tokens-1)
        # x = x.prod(dim=1).sqrt(n_layers) # (B, tokens-1)
        # x = x.min(dim=1) # (B, tokens-1)

        return x  # (1, tokens-1)
    
    
    def predmap24(self, x, idx, return_extras=False, layer=-1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)

        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)

        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()
        x = x[..., idx] # (B, layers, tokens)

        # Strip CLS token
        x = x[..., 1:] # (B, layers, tokens-1)

        x = x.max(dim=1).values # (B, tokens-1)
        # x = x.prod(dim=1).sqrt(n_layers) # (B, tokens-1)
        # x = x.min(dim=1) # (B, tokens-1)

        return x  # (1, tokens-1)
    
    
    def predmap25(self, x, idx, return_extras=False, layer=-1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)

        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)

        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()
        x = x[..., idx] # (B, layers, tokens)

        # Strip CLS token
        x = x[..., 1:] # (B, layers, tokens-1)

        x = x.prod(dim=1).pow(1/n_layers) # (B, tokens-1)
        # x = x.min(dim=1) # (B, tokens-1)

        return x  # (1, tokens-1)
    
    def predmap26(self, x, idx, return_extras=False, layer=-1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)

        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)

        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()
        x = x[..., idx] # (B, layers, tokens)

        # Strip CLS token
        x = x[..., 1:] # (B, layers, tokens-1)

        # x = x.prod(dim=1).pow(1/n_layers) # (B, tokens-1)
        x = x.min(dim=1).values # (B, tokens-1)

        return x  # (1, tokens-1)
    

    def predmap27(self, x, idx, return_extras=False, layer=-1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)

        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)
        # calc Entropy
        n_classes = x.shape[-1]
        ent = x.softmax(dim=-1) # (B, layers, tokens, classes)
        ent = ent * ent.log() # (B, layers, tokens, classes)
        ent = -ent.sum(dim=-1, keepdims=True) # (B, layers, tokens, 1)
        ent = ent / torch.tensor(n_classes, dtype=torch.float32).log() # (B, layers, tokens, 1)
        focus = 1 - ent # (B, layers, tokens, 1)
        focus = focus.softmax(dim=1) # (B, layers, tokens, 1)
        y = focus * x # (B, layers, tokens, classes)
        y = y.sum(dim=1) # (B, tokens, classes)


        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()
        y = y[..., idx] # (B, tokens)

        # Strip CLS token
        y = y[..., 1:] # (B, tokens-1)

        return y  # (1, tokens-1)
    
    def predmap28(self, x, idx, return_extras=False, temperature=1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)

        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)

        # Select predicted class
        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()


        x = x[..., idx] # (B, layers, tokens)
        
        # calc Entropy
        n_tokens = x.shape[-1]
        ent = x.softmax(dim=-1) # (B, layers, tokens)
        ent = ent * ent.log() # (B, layers, tokens)
        ent = -ent.sum(dim=-1, keepdims=True) # (B, layers, 1)
        ent = ent / torch.tensor(n_tokens, dtype=torch.float32).log() # (B, layers, 1)
        focus = 1 - ent # (B, layers, 1)
        focus = focus / temperature # (B, layers, 1)
        focus[:,-1,:] = -99999
        focus = focus.softmax(dim=1) # (B, layers, 1)
        y = focus * x # (B, layers, tokens)
        y = y.sum(dim=1) # (B, tokens)
        predmap = y # (B, tokens)
        predmap = predmap[..., 1:] # (1, tokens-1)

        
        ##########################################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        # (blocks*heads, tokens-1) * (1, tokens-1) -> (blocks*heads, classes)
        projections = all_attnmap * predmap # (blocks*heads, tokens-1)
        projections = projections.sum(dim=-1, keepdims=True) # (blocks*heads, 1)
        # TODO: is this needed?
        projections = projections.softmax(dim=0) # (blocks*heads, 1)

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, 1)
        z = weighted_attnmap * predmap.T # (tokens-1, 1)
        ##########################################
        return z.T  # (1, tokens-1)    


    def predmap29(self, x, idx, return_extras=False, temperature=1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)
 
        # all_v # (B, layers, heads, tokens, embed_dim)
        all_v = torch.stack(
            [
                blk.attn.get_v() # (B, heads, tokens, embed_dim)
                for blk in self.blocks],
            dim=1)
        


        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)

        # Select predicted class
        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()


        x = x[..., idx] # (B, layers, tokens)
        
        # calc Entropy
        n_tokens = x.shape[-1]
        ent = x.softmax(dim=-1) # (B, layers, tokens)
        ent = ent * ent.log() # (B, layers, tokens)
        ent = -ent.sum(dim=-1, keepdims=True) # (B, layers, 1)
        ent = ent / torch.tensor(n_tokens, dtype=torch.float32).log() # (B, layers, 1)
        focus = 1 - ent # (B, layers, 1)
        focus = focus / temperature # (B, layers, 1)
        focus[:,-1,:] = -99999
        focus = focus.softmax(dim=1) # (B, layers, 1)
        y = focus * x # (B, layers, tokens)
        y = y.sum(dim=1) # (B, tokens)
        predmap = y # (B, tokens)
        predmap = predmap[..., 1:] # (1, tokens-1)

        
        ##########################################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)

        # all_attn (B, layers, heads, tokens, tokens)
        # all_v (B, layers, heads, tokens, embed_dim)

        a1 = (all_attn[..., 0, :] - all_attn[..., 0, [0]]).abs() # (B, layers, heads, tokens)
        a1 = a1 / a1.norm(p=1, dim=-1, keepdim=True) # (B, layers, heads, tokens)

        a2 = (all_v[..., [0], :] - all_v).abs() # (B, layers, heads, tokens, embed_dim)
        a2 = a2.norm(p=1, dim=-1) # (B, layers, heads, tokens)
        a2 = a2 / a2.norm(p=1, dim=-1, keepdim=True) # (B, layers, heads, tokens)

        alpha = (a1 * a2).sum(dim=-1) # (B, layers, heads)
        alpha = alpha.flatten(1,2) # (B, layers*heads)
        alpha = alpha.T # (layers*heads, 1)
        alpha = alpha.softmax(dim=0) # (layers*heads, 1)



        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (blocks, heads, tokens-1)
        all_attnmap = all_attnmap.flatten(0,1) # (blocks*heads, tokens-1)

        # (blocks*heads, tokens-1) * (1, tokens-1) -> (blocks*heads, classes)
        # projections = all_attnmap * predmap # (blocks*heads, tokens-1)
        # projections = projections.sum(dim=-1, keepdims=True) # (blocks*heads, 1)
        # TODO: is this needed?
        # projections = projections.softmax(dim=0) # (blocks*heads, 1)

        projections = alpha

        weighted_attnmap = all_attnmap.T @ projections # (tokens-1, 1)
        z = weighted_attnmap * predmap.T # (tokens-1, 1)
        ##########################################
        return z.T  # (1, tokens-1)
    
    
    def predmap30(self, x, idx, return_extras=False, temperature=1):
        B = x.shape[0]
        x = self.patch_embed(x) # (B, tokens-1, embed_dim)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1) # (B, tokens, embed_dim)
        x = self.add([x, self.pos_embed]) # (B, tokens, embed_dim)

        # x.register_hook(self.save_inp_grad)

        layers_out = []
        for blk in self.blocks:
            x = blk(x) # (B, tokens, embed_dim)
            layers_out.append(x)
        layers_out = torch.stack(layers_out, dim=1) # (B, layers, tokens, embed_dim)
        x = layers_out.flatten(0,1) # (B*layers, tokens, embed_dim)
        x = self.norm(x) # (B*layers, tokens, embed_dim)
        x = self.head(x) # (B*layers, tokens, classes)

        n_layers = len(self.blocks)
        x = x.unflatten(0, (B, n_layers)) # (B, layers, tokens, classes)

        # Select predicted class
        if idx is None:
            idx = torch.argmax(x[:,-1,0], dim=-1).item()


        x = x[..., idx] # (B, layers, tokens)

        predmaps = x[..., 1:] # (B, layers, tokens-1)
        predmaps = predmaps[0] # (layers, tokens-1)
               
        
        ##########################################
        # get all attention from all blocks
        all_attn = torch.stack(
            [
                blk.attn.get_attn() # (B, heads, tokens, tokens)
                for blk in self.blocks],
            dim=1) # (B, blocks, heads, tokens, tokens)
        all_attn = all_attn[0].clamp(min=0) # (blocks, heads, tokens, tokens)
        all_attnmap = all_attn[:, :, 0, 1:] # (layers, heads, tokens-1)

        # predmaps[None, None, :] # (1,1,layers, tokens-1)
        # all_attnmap   # (layers, heads, 1, tokens-1)
        hadamard_prod = (predmaps[None, None, ...] * all_attnmap[...,None,:]) # (layers, heads, layers, tokens-1)
        alphas = hadamard_prod.sum(dim=-1) # (layers, heads, layers)
        shape_ = alphas.shape
        alphas = alphas.flatten()
        # alphas = alphas / 0.0001
        alphas = alphas.softmax(dim=-1).reshape(shape_)



        z = (alphas.unsqueeze(-1) * hadamard_prod).flatten(0,2).sum(dim=0) # (tokens-1, )
        z = z.unsqueeze(0) # (1, tokens-1)
        return z
    

    def predmap31_batched_layer(self, x, idx, return_extras=False, layer=-1, 
                                apply_softmax_classes=False, apply_softmax_tokens=False):
        layer = len(self.blocks) + layer if layer < 0 else layer
        assert 0<=layer<len(self.blocks), f"Layer index out of range: {layer}"
        extras = {}
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = self.add([x, self.pos_embed])

        # x.register_hook(self.save_inp_grad)

        for l, blk in enumerate(self.blocks):
            x = blk(x)
            if l == layer:
                # calc predmap
                x_layer = x
                x_layer = self.norm(x_layer) # (B, tokens, embed_dim)
                x_layer = self.head(x_layer) # (B, tokens, classes)
                predmap = x_layer[:, 1:, :] # (B, tokens-1, classes)


        x = self.norm(x) # (B, tokens, embed_dim)
        x = self.head(x) # (B, tokens, classes)
        if idx is None:
            idx = torch.argmax(x[:,0], dim=-1) # (B, )
        
        if return_extras:
            extras["xpredmap"] = x # (B, tokens, classes)


        x = predmap # (B, tokens-1, classes)
        if apply_softmax_classes:
            x = x.softmax(dim=-1) # (B, tokens-1, classes)
        x = x[range(B), ..., idx] # (B, tokens-1, )
        if apply_softmax_tokens:
            x = x.softmax(dim=-1) # (B, tokens-1)
        if return_extras:
            # extras["all_attnmap"] = all_attnmap.unflatten(1, (len(self.blocks), self.blocks[0].attn.num_heads)) # (B, blocks, heads, tokens-1)
            # extras["projections"] = projections.unflatten(1, (len(self.blocks), self.blocks[0].attn.num_heads)) # (B, blocks, heads, classes)
            # extras["weighted_attnmap"] = weighted_attnmap # (B, tokens-1, classes)
            extras["predmap"] = predmap # (B, tokens-1, classes)
            # extras["CLS_self_attend"] = all_attn[..., 0, 0] # (B, blocks, heads)
            return x, extras
        return x # (B, tokens-1)
    
    def predmap31_softmax_classes_batched_layer(self, x, idx, return_extras=False, layer=-1):
        return self.predmap31_batched_layer(x, idx, return_extras, layer=layer, apply_softmax_classes=True)
    
    def predmap31_softmax_classes_tokens_batched_layer(self, x, idx, return_extras=False, layer=-1):
        return self.predmap31_batched_layer(x, idx, return_extras, layer=layer, apply_softmax_classes=True, apply_softmax_tokens=True)
    
    def predmap31_softmax_tokens_batched_layer(self, x, idx, return_extras=False, layer=-1):
        return self.predmap31_batched_layer(x, idx, return_extras, layer=layer, apply_softmax_tokens=True)
    
    

    def relprop(self, cam=None,method="transformer_attribution", is_ablation=False, start_layer=0, **kwargs):
        # print(kwargs)
        # print("conservation 1", cam.sum())
        # cam (1, classes)
        cam = self.head.relprop(cam, **kwargs) # (1, embed_dim)
        cam = cam.unsqueeze(1) # (1, 1, embed_dim)
        cam = self.pool.relprop(cam, **kwargs) # (1, tokens-1, embed_dim)
        cam = self.norm.relprop(cam, **kwargs) # (1, tokens-1, embed_dim)
        for blk in reversed(self.blocks):
            cam = blk.relprop(cam, **kwargs) # (1, tokens-1, embed_dim)

        # print("conservation 2", cam.sum())
        # print("min", cam.min())

        if method == "full":
            (cam, _) = self.add.relprop(cam, **kwargs)
            cam = cam[:, 1:]
            cam = self.patch_embed.relprop(cam, **kwargs)
            # sum on channels
            cam = cam.sum(dim=1)
            return cam

        elif method == "rollout":
            # cam rollout
            attn_cams = []
            for blk in self.blocks:
                attn_heads = blk.attn.get_attn_cam().clamp(min=0)
                avg_heads = (attn_heads.sum(dim=1) / attn_heads.shape[1]).detach()
                attn_cams.append(avg_heads)
            cam = compute_rollout_attention(attn_cams, start_layer=start_layer)
            cam = cam[:, 0, 1:]
            return cam
        
        # our method, method name grad is legacy
        elif method == "transformer_attribution" or method == "grad":
            cams = []
            for blk in self.blocks:
                grad = blk.attn.get_attn_gradients() # (1, heads, tokens, tokens)
                cam = blk.attn.get_attn_cam() # (1, heads, tokens, tokens)
                cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1]) # (heads, tokens, tokens)
                grad = grad[0].reshape(-1, grad.shape[-1], grad.shape[-1]) # (heads, tokens, tokens)
                cam = grad * cam # (heads, tokens, tokens)
                cam = cam.clamp(min=0).mean(dim=0) # (tokens, tokens)
                cams.append(cam.unsqueeze(0)) # (1, tokens, tokens)
            # cams List[layers x (1, tokens, tokens)]
            rollout = compute_rollout_attention(cams, start_layer=start_layer) # (1, tokens, tokens)
            cam = rollout[:, 0, 1:] # (1, tokens-1)
            return cam
            
        elif method == "last_layer":
            cam = self.blocks[-1].attn.get_attn_cam()
            cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])
            if is_ablation:
                grad = self.blocks[-1].attn.get_attn_gradients()
                grad = grad[0].reshape(-1, grad.shape[-1], grad.shape[-1])
                cam = grad * cam
            cam = cam.clamp(min=0).mean(dim=0)
            cam = cam[0, 1:]
            return cam

        elif method == "last_layer_attn":
            cam = self.blocks[-1].attn.get_attn()
            cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])
            cam = cam.clamp(min=0).mean(dim=0)
            cam = cam[0, 1:]
            return cam

        elif method == "second_layer":
            cam = self.blocks[1].attn.get_attn_cam()
            cam = cam[0].reshape(-1, cam.shape[-1], cam.shape[-1])
            if is_ablation:
                grad = self.blocks[1].attn.get_attn_gradients()
                grad = grad[0].reshape(-1, grad.shape[-1], grad.shape[-1])
                cam = grad * cam
            cam = cam.clamp(min=0).mean(dim=0)
            cam = cam[0, 1:]
            return cam


def _conv_filter(state_dict, patch_size=16):
    """ convert patch embedding weight from manual patchify + linear proj to conv"""
    out_dict = {}
    for k, v in state_dict.items():
        if 'patch_embed.proj.weight' in k:
            v = v.reshape((v.shape[0], 3, patch_size, patch_size))
        out_dict[k] = v
    return out_dict

def vit_base_patch16_224(pretrained=False, **kwargs):
    model = VisionTransformer(
        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, **kwargs)
    model.default_cfg = default_cfgs['vit_base_patch16_224']
    if pretrained:
        load_pretrained(
            model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3), filter_fn=_conv_filter)
    return model

def vit_large_patch16_224(pretrained=False, **kwargs):
    model = VisionTransformer(
        patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True, **kwargs)
    model.default_cfg = default_cfgs['vit_large_patch16_224']
    if pretrained:
        load_pretrained(model, num_classes=model.num_classes, in_chans=kwargs.get('in_chans', 3))
    return model

def deit_base_patch16_224(pretrained=False, **kwargs):
    model = VisionTransformer(
        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, **kwargs)
    model.default_cfg = _cfg()
    if pretrained:
        checkpoint = torch.hub.load_state_dict_from_url(
            url="https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth",
            map_location="cpu", check_hash=True
        )
        model.load_state_dict(checkpoint["model"])
    return model