# Author: Huzheng Yang
# %%
from typing import Optional, Tuple
from einops import rearrange
import requests
import torch
import torch.nn.functional as F
import timm
from torch import nn
import numpy as np
import os
from functools import partial

MODEL_DICT = {}
LAYER_DICT = {}
RES_DICT = {}

class SAM2(nn.Module):

    def __init__(self, model_cfg='sam2_hiera_b+',):
        super().__init__()

        try:
            from sam2.build_sam import build_sam2
        except ImportError:
            print("Please install segment_anything_2 from https://github.com/facebookresearch/segment-anything-2.git")
            return
        
        config_dict = {
            'sam2_hiera_l': ("sam2_hiera_large.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt"),
            'sam2_hiera_b+': ("sam2_hiera_base_plus.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt"),
            'sam2_hiera_s': ("sam2_hiera_small.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt"),
            'sam2_hiera_t': ("sam2_hiera_tiny.pt", "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt"),
        }
        filename, url = config_dict[model_cfg]
        if not os.path.exists(filename):
            print(f"Downloading {url}")
            r = requests.get(url)
            with open(filename, 'wb') as f:
                f.write(r.content)
        sam2_checkpoint = filename
        
        device = 'cpu'
        sam2_model = build_sam2(model_cfg, sam2_checkpoint, device=device)

        image_encoder = sam2_model.image_encoder
        image_encoder.eval()
        
        from sam2.modeling.backbones.hieradet import do_pool
        from sam2.modeling.backbones.utils import window_partition, window_unpartition
        def new_forward(self, x: torch.Tensor) -> torch.Tensor:
            shortcut = x  # B, H, W, C
            x = self.norm1(x)

            # Skip connection
            if self.dim != self.dim_out:
                shortcut = do_pool(self.proj(x), self.pool)

            # Window partition
            window_size = self.window_size
            if window_size > 0:
                H, W = x.shape[1], x.shape[2]
                x, pad_hw = window_partition(x, window_size)

            # Window Attention + Q Pooling (if stage change)
            x = self.attn(x)
            if self.q_stride:
                # Shapes have changed due to Q pooling
                window_size = self.window_size // self.q_stride[0]
                H, W = shortcut.shape[1:3]

                pad_h = (window_size - H % window_size) % window_size
                pad_w = (window_size - W % window_size) % window_size
                pad_hw = (H + pad_h, W + pad_w)

            # Reverse window partition
            if self.window_size > 0:
                x = window_unpartition(x, window_size, pad_hw, (H, W))

            self.attn_output = x.clone()
            
            x = shortcut + self.drop_path(x)
            # MLP
            mlp_out = self.mlp(self.norm2(x))
            self.mlp_output = mlp_out.clone()
            x = x + self.drop_path(mlp_out)
            self.block_output = x.clone()
            return x        
        
        setattr(image_encoder.trunk.blocks[0].__class__, 'forward', new_forward)
        
        self.image_encoder = image_encoder
        
        
        
    @torch.no_grad()
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        output = self.image_encoder(x)
        attn_outputs, mlp_outputs, block_outputs = [], [], []
        for block in self.image_encoder.trunk.blocks:
            attn_outputs.append(block.attn_output)
            mlp_outputs.append(block.mlp_output)
            block_outputs.append(block.block_output)
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }

MODEL_DICT["SAM2(sam2_hiera_t)"] = partial(SAM2, model_cfg='sam2_hiera_t')
LAYER_DICT["SAM2(sam2_hiera_t)"] = 12
RES_DICT["SAM2(sam2_hiera_t)"] = (1024, 1024)
MODEL_DICT["SAM2(sam2_hiera_s)"] = partial(SAM2, model_cfg='sam2_hiera_s')
LAYER_DICT["SAM2(sam2_hiera_s)"] = 16
RES_DICT["SAM2(sam2_hiera_s)"] = (1024, 1024)
MODEL_DICT["SAM2(sam2_hiera_b+)"] = partial(SAM2, model_cfg='sam2_hiera_b+')
LAYER_DICT["SAM2(sam2_hiera_b+)"] = 24
RES_DICT["SAM2(sam2_hiera_b+)"] = (1024, 1024)
MODEL_DICT["SAM2(sam2_hiera_l)"] = partial(SAM2, model_cfg='sam2_hiera_l')
LAYER_DICT["SAM2(sam2_hiera_l)"] = 48
RES_DICT["SAM2(sam2_hiera_l)"] = (1024, 1024)


class SAM(torch.nn.Module):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        from segment_anything import sam_model_registry, SamPredictor
        from segment_anything.modeling.sam import Sam

        checkpoint = "sam_vit_b_01ec64.pth"
        if not os.path.exists(checkpoint):
            checkpoint_url = (
                "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
            )
            import requests

            r = requests.get(checkpoint_url)
            with open(checkpoint, "wb") as f:
                f.write(r.content)

        sam: Sam = sam_model_registry["vit_b"](checkpoint=checkpoint)

        from segment_anything.modeling.image_encoder import (
            window_partition,
            window_unpartition,
        )

        def new_block_forward(self, x: torch.Tensor) -> torch.Tensor:
            shortcut = x
            x = self.norm1(x)
            # Window partition
            if self.window_size > 0:
                H, W = x.shape[1], x.shape[2]
                x, pad_hw = window_partition(x, self.window_size)

            x = self.attn(x)
            # Reverse window partition
            if self.window_size > 0:
                x = window_unpartition(x, self.window_size, pad_hw, (H, W))
            self.attn_output = x.clone()

            x = shortcut + x
            mlp_outout = self.mlp(self.norm2(x))
            self.mlp_output = mlp_outout.clone()
            x = x + mlp_outout
            self.block_output = x.clone()

            return x

        setattr(sam.image_encoder.blocks[0].__class__, "forward", new_block_forward)

        self.image_encoder = sam.image_encoder
        self.image_encoder.eval()

    @torch.no_grad()
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            x = torch.nn.functional.interpolate(x, size=(1024, 1024), mode="bilinear")
        out = self.image_encoder(x)

        attn_outputs, mlp_outputs, block_outputs = [], [], []
        for i, blk in enumerate(self.image_encoder.blocks):
            attn_outputs.append(blk.attn_output)
            mlp_outputs.append(blk.mlp_output)
            block_outputs.append(blk.block_output)
        attn_outputs = torch.stack(attn_outputs)
        mlp_outputs = torch.stack(mlp_outputs)
        block_outputs = torch.stack(block_outputs)
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }

MODEL_DICT["SAM(sam_vit_b)"] = partial(SAM)
LAYER_DICT["SAM(sam_vit_b)"] = 12
RES_DICT["SAM(sam_vit_b)"] = (1024, 1024)


class MobileSAM(nn.Module):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        from mobile_sam import sam_model_registry

        url = "https://raw.githubusercontent.com/ChaoningZhang/MobileSAM/master/weights/mobile_sam.pt"
        model_type = "vit_t"
        sam_checkpoint = "mobile_sam.pt"
        if not os.path.exists(sam_checkpoint):
            import requests

            r = requests.get(url)
            with open(sam_checkpoint, "wb") as f:
                f.write(r.content)

        mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)

        def new_forward_fn(self, x):
            shortcut = x

            x = self.conv1(x)
            x = self.act1(x)

            x = self.conv2(x)
            x = self.act2(x)

            self.attn_output = rearrange(x.clone(), "b c h w -> b h w c")

            x = self.conv3(x)

            self.mlp_output = rearrange(x.clone(), "b c h w -> b h w c")

            x = self.drop_path(x)

            x += shortcut
            x = self.act3(x)

            self.block_output = rearrange(x.clone(), "b c h w -> b h w c")

            return x

        setattr(
            mobile_sam.image_encoder.layers[0].blocks[0].__class__,
            "forward",
            new_forward_fn,
        )

        def new_forward_fn2(self, x):
            H, W = self.input_resolution
            B, L, C = x.shape
            assert L == H * W, "input feature has wrong size"
            res_x = x
            if H == self.window_size and W == self.window_size:
                x = self.attn(x)
            else:
                x = x.view(B, H, W, C)
                pad_b = (self.window_size - H % self.window_size) % self.window_size
                pad_r = (self.window_size - W % self.window_size) % self.window_size
                padding = pad_b > 0 or pad_r > 0

                if padding:
                    x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))

                pH, pW = H + pad_b, W + pad_r
                nH = pH // self.window_size
                nW = pW // self.window_size
                # window partition
                x = (
                    x.view(B, nH, self.window_size, nW, self.window_size, C)
                    .transpose(2, 3)
                    .reshape(B * nH * nW, self.window_size * self.window_size, C)
                )
                x = self.attn(x)
                # window reverse
                x = (
                    x.view(B, nH, nW, self.window_size, self.window_size, C)
                    .transpose(2, 3)
                    .reshape(B, pH, pW, C)
                )

                if padding:
                    x = x[:, :H, :W].contiguous()

                x = x.view(B, L, C)

            hw = np.sqrt(x.shape[1]).astype(int)
            self.attn_output = rearrange(x.clone(), "b (h w) c -> b h w c", h=hw)

            x = res_x + self.drop_path(x)

            x = x.transpose(1, 2).reshape(B, C, H, W)
            x = self.local_conv(x)
            x = x.view(B, C, L).transpose(1, 2)

            mlp_output = self.mlp(x)
            self.mlp_output = rearrange(
                mlp_output.clone(), "b (h w) c -> b h w c", h=hw
            )

            x = x + self.drop_path(mlp_output)
            self.block_output = rearrange(x.clone(), "b (h w) c -> b h w c", h=hw)
            return x

        setattr(
            mobile_sam.image_encoder.layers[1].blocks[0].__class__,
            "forward",
            new_forward_fn2,
        )

        mobile_sam.eval()
        self.image_encoder = mobile_sam.image_encoder

    @torch.no_grad()
    def forward(self, x):
        with torch.no_grad():
            x = torch.nn.functional.interpolate(x, size=(1024, 1024), mode="bilinear")
        out = self.image_encoder(x)

        attn_outputs, mlp_outputs, block_outputs = [], [], []
        for i_layer in range(len(self.image_encoder.layers)):
            for i_block in range(len(self.image_encoder.layers[i_layer].blocks)):
                blk = self.image_encoder.layers[i_layer].blocks[i_block]
                attn_outputs.append(blk.attn_output)
                mlp_outputs.append(blk.mlp_output)
                block_outputs.append(blk.block_output)
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }

MODEL_DICT["MobileSAM"] = partial(MobileSAM)
LAYER_DICT["MobileSAM"] = 12
RES_DICT["MobileSAM"] = (1024, 1024)


class DiNOv2(torch.nn.Module):
    def __init__(self, ver="dinov2_vitb14_reg", num_reg=5):
        super().__init__()
        self.dinov2 = torch.hub.load("facebookresearch/dinov2", ver)
        self.dinov2.requires_grad_(False)
        self.dinov2.eval()
        self.num_reg = num_reg

        def new_block_forward(self, x: torch.Tensor) -> torch.Tensor:
            def attn_residual_func(x):
                return self.ls1(self.attn(self.norm1(x)))

            def ffn_residual_func(x):
                return self.ls2(self.mlp(self.norm2(x)))

            attn_output = attn_residual_func(x)

            hw = np.sqrt(attn_output.shape[1] - num_reg).astype(int)
            self.attn_output = rearrange(
                attn_output.clone()[:, num_reg:], "b (h w) c -> b h w c", h=hw
            )

            x = x + attn_output
            mlp_output = ffn_residual_func(x)
            self.mlp_output = rearrange(
                mlp_output.clone()[:, num_reg:], "b (h w) c -> b h w c", h=hw
            )
            x = x + mlp_output
            block_output = x
            self.block_output = rearrange(
                block_output.clone()[:, num_reg:], "b (h w) c -> b h w c", h=hw
            )
            return x

        setattr(self.dinov2.blocks[0].__class__, "forward", new_block_forward)

    @torch.no_grad()
    def forward(self, x):

        out = self.dinov2(x)

        attn_outputs, mlp_outputs, block_outputs = [], [], []
        for i, blk in enumerate(self.dinov2.blocks):
            attn_outputs.append(blk.attn_output)
            mlp_outputs.append(blk.mlp_output)
            block_outputs.append(blk.block_output)

        attn_outputs = torch.stack(attn_outputs)
        mlp_outputs = torch.stack(mlp_outputs)
        block_outputs = torch.stack(block_outputs)
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }

MODEL_DICT["DiNOv2reg(dinov2_vitb14_reg)"] = partial(DiNOv2, ver="dinov2_vitb14_reg", num_reg=5)
LAYER_DICT["DiNOv2reg(dinov2_vitb14_reg)"] = 12
RES_DICT["DiNOv2reg(dinov2_vitb14_reg)"] = (672, 672)
MODEL_DICT["DiNOv2(dinov2_vitb14)"] = partial(DiNOv2, ver="dinov2_vitb14", num_reg=1)
LAYER_DICT["DiNOv2(dinov2_vitb14)"] = 12
RES_DICT["DiNOv2(dinov2_vitb14)"] = (672, 672)

class DiNO(nn.Module):
    def __init__(self, ver="dino_vitb8"):
        super().__init__()
        model = torch.hub.load('facebookresearch/dino:main', ver)
        model = model.eval()

        def remove_cls_and_reshape(x):
            x = x.clone()
            x = x[:, 1:]
            hw = np.sqrt(x.shape[1]).astype(int)
            x = rearrange(x, "b (h w) c -> b h w c", h=hw)
            return x

        def new_forward(self, x, return_attention=False):
            y, attn = self.attn(self.norm1(x))
            self.attn_output = remove_cls_and_reshape(y.clone())
            if return_attention:
                return attn
            x = x + self.drop_path(y)
            mlp_output = self.mlp(self.norm2(x))
            self.mlp_output = remove_cls_and_reshape(mlp_output.clone())
            x = x + self.drop_path(mlp_output)
            self.block_output = remove_cls_and_reshape(x.clone())
            return x

        setattr(model.blocks[0].__class__, "forward", new_forward)
        
        self.model = model
        self.model.eval()
        self.model.requires_grad_(False)
        
    def forward(self, x):
        out = self.model(x)
        attn_outputs = [block.attn_output for block in self.model.blocks]
        mlp_outputs = [block.mlp_output for block in self.model.blocks]
        block_outputs = [block.block_output for block in self.model.blocks]
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }
            
MODEL_DICT["DiNO(dino_vitb8)"] = partial(DiNO)
LAYER_DICT["DiNO(dino_vitb8)"] = 12
RES_DICT["DiNO(dino_vitb8)"] = (448, 448)

def resample_position_embeddings(embeddings, h, w):
    cls_embeddings = embeddings[0]
    patch_embeddings = embeddings[1:]  # [14*14, 768]
    hw = np.sqrt(patch_embeddings.shape[0]).astype(int)
    patch_embeddings = rearrange(patch_embeddings, "(h w) c -> c h w", h=hw)
    patch_embeddings = F.interpolate(patch_embeddings.unsqueeze(0), size=(h, w), mode="nearest").squeeze(0)
    patch_embeddings = rearrange(patch_embeddings, "c h w -> (h w) c")
    embeddings = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
    return embeddings

# class CLIP(torch.nn.Module):
#     def __init__(self, ver="openai/clip-vit-base-patch16"):
#         super().__init__()

#         from transformers import CLIPProcessor, CLIPModel

#         model = CLIPModel.from_pretrained(ver)
                
#         # resample the patch embeddings to 56x56, take 896x896 input
#         embeddings = model.vision_model.embeddings.position_embedding.weight
#         embeddings = resample_position_embeddings(embeddings, 42, 42)
#         model.vision_model.embeddings.position_embedding.weight = nn.Parameter(embeddings)
#         model.vision_model.embeddings.position_ids = torch.arange(0, 1+56*56)
        
#         # processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
#         self.model = model.eval()

#         def new_forward(
#             self,
#             hidden_states: torch.Tensor,
#             attention_mask: torch.Tensor,
#             causal_attention_mask: torch.Tensor,
#             output_attentions: Optional[bool] = False,
#         ) -> Tuple[torch.FloatTensor]:

#             residual = hidden_states

#             hidden_states = self.layer_norm1(hidden_states)
#             hidden_states, attn_weights = self.self_attn(
#                 hidden_states=hidden_states,
#                 attention_mask=attention_mask,
#                 causal_attention_mask=causal_attention_mask,
#                 output_attentions=output_attentions,
#             )
#             hw = np.sqrt(hidden_states.shape[1] - 1).astype(int)
#             self.attn_output = rearrange(
#                 hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
#             )
#             hidden_states = residual + hidden_states

#             residual = hidden_states
#             hidden_states = self.layer_norm2(hidden_states)
#             hidden_states = self.mlp(hidden_states)
#             self.mlp_output = rearrange(
#                 hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
#             )

#             hidden_states = residual + hidden_states

#             outputs = (hidden_states,)

#             if output_attentions:
#                 outputs += (attn_weights,)

#             self.block_output = rearrange(
#                 hidden_states.clone()[:, 1:], "b (h w) c -> b h w c", h=hw
#             )
#             return outputs

#         setattr(
#             self.model.vision_model.encoder.layers[0].__class__, "forward", new_forward
#         )

#     @torch.no_grad()
#     def forward(self, x):

#         out = self.model.vision_model(x)

#         attn_outputs, mlp_outputs, block_outputs = [], [], []
#         for i, blk in enumerate(self.model.vision_model.encoder.layers):
#             attn_outputs.append(blk.attn_output)
#             mlp_outputs.append(blk.mlp_output)
#             block_outputs.append(blk.block_output)

#         attn_outputs = torch.stack(attn_outputs)
#         mlp_outputs = torch.stack(mlp_outputs)
#         block_outputs = torch.stack(block_outputs)
#         return attn_outputs, mlp_outputs, block_outputs


# MODEL_DICT["CLIP(openai/clip-vit-base-patch16)"] = partial(CLIP, ver="openai/clip-vit-base-patch16")
# LAYER_DICT["CLIP(openai/clip-vit-base-patch16)"] = 12
# RES_DICT["CLIP(openai/clip-vit-base-patch16)"] = (896, 896)


class OpenCLIPViT(nn.Module):
    def __init__(self, version='ViT-B-16', pretrained='laion2b_s34b_b88k'):
        super().__init__()
        try:
            import open_clip
        except ImportError:
            print("Please install open_clip to use this class.")
            return
        
        model, _, _ = open_clip.create_model_and_transforms(version, pretrained=pretrained)
        
        positional_embedding = resample_position_embeddings(model.visual.positional_embedding, 42, 42)
        model.visual.positional_embedding = nn.Parameter(positional_embedding)
        
        def new_forward(
                self,
                q_x: torch.Tensor,
                k_x: Optional[torch.Tensor] = None,
                v_x: Optional[torch.Tensor] = None,
                attn_mask: Optional[torch.Tensor] = None,
        ):
            def remove_cls_and_reshape(x):
                x = x.clone()
                x = x[1:]
                hw = np.sqrt(x.shape[0]).astype(int)
                x = rearrange(x, "(h w) b c -> b h w c", h=hw)
                return x
            
            k_x = self.ln_1_kv(k_x) if hasattr(self, "ln_1_kv") and k_x is not None else None
            v_x = self.ln_1_kv(v_x) if hasattr(self, "ln_1_kv") and v_x is not None else None

            attn_output = self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, attn_mask=attn_mask)
            self.attn_output = remove_cls_and_reshape(attn_output.clone())
            x = q_x + self.ls_1(attn_output)
            mlp_output = self.mlp(self.ln_2(x))
            self.mlp_output = remove_cls_and_reshape(mlp_output.clone())
            x = x + self.ls_2(mlp_output)    
            self.block_output = remove_cls_and_reshape(x.clone())
            return x


        setattr(model.visual.transformer.resblocks[0].__class__, "forward", new_forward)
        
        self.model = model
        self.model.eval()
    
    def forward(self, x):
        out = self.model(x)
        attn_outputs, mlp_outputs, block_outputs = [], [], []
        for block in self.model.visual.transformer.resblocks:
            attn_outputs.append(block.attn_output)
            mlp_outputs.append(block.mlp_output)
            block_outputs.append(block.block_output)
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }
        
MODEL_DICT["CLIP(ViT-B-16/openai)"] = partial(OpenCLIPViT, version='ViT-B-16', pretrained='openai')
LAYER_DICT["CLIP(ViT-B-16/openai)"] = 12
RES_DICT["CLIP(ViT-B-16/openai)"] = (672, 672)
MODEL_DICT["CLIP(ViT-B-16/laion2b_s34b_b88k)"] = partial(OpenCLIPViT, version='ViT-B-16', pretrained='laion2b_s34b_b88k')
LAYER_DICT["CLIP(ViT-B-16/laion2b_s34b_b88k)"] = 12
RES_DICT["CLIP(ViT-B-16/laion2b_s34b_b88k)"] = (672, 672)

class EVA02(nn.Module):
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
        model = timm.create_model(
            'eva02_base_patch14_448.mim_in22k_ft_in1k',
            pretrained=True,
            num_classes=0,  # remove classifier nn.Linear
        )
        model = model.eval()

        def new_forward(self, x, rope: Optional[torch.Tensor] = None, attn_mask: Optional[torch.Tensor] = None):
        
            def remove_cls_and_reshape(x):
                x = x.clone()
                x = x[:, 1:]
                hw = np.sqrt(x.shape[1]).astype(int)
                x = rearrange(x, "b (h w) c -> b h w c", h=hw)
                return x
            
            if self.gamma_1 is None:
                attn_output = self.attn(self.norm1(x), rope=rope, attn_mask=attn_mask)
                self.attn_output = remove_cls_and_reshape(attn_output.clone())
                x = x + self.drop_path1(attn_output)
                mlp_output = self.mlp(self.norm2(x))
                self.mlp_output = remove_cls_and_reshape(mlp_output.clone())
                x = x + self.drop_path2(mlp_output)
            else:
                attn_output = self.attn(self.norm1(x), rope=rope, attn_mask=attn_mask)
                self.attn_output = remove_cls_and_reshape(attn_output.clone())
                x = x + self.drop_path1(self.gamma_1 * attn_output)
                mlp_output = self.mlp(self.norm2(x))
                self.mlp_output = remove_cls_and_reshape(mlp_output.clone())
                x = x + self.drop_path2(self.gamma_2 * mlp_output)
            self.block_output = remove_cls_and_reshape(x.clone())
            return x

        setattr(model.blocks[0].__class__, "forward", new_forward)
        
        self.model = model
        
    def forward(self, x):
        out = self.model(x)
        attn_outputs = [block.attn_output for block in self.model.blocks]
        mlp_outputs = [block.mlp_output for block in self.model.blocks]
        block_outputs = [block.block_output for block in self.model.blocks]
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }
        
MODEL_DICT["CLIP(eva02_base_patch14_448.mim_in22k_ft_in1k)"] = partial(EVA02)
LAYER_DICT["CLIP(eva02_base_patch14_448.mim_in22k_ft_in1k)"] = 12
RES_DICT["CLIP(eva02_base_patch14_448.mim_in22k_ft_in1k)"] = (448, 448)

class CLIPConvnext(nn.Module):
    def __init__(self):
        super().__init__()
        try:
            import open_clip
        except ImportError:
            print("Please install open_clip to use this class.")
            return
    
        model, _, _ = open_clip.create_model_and_transforms('convnext_base_w_320', pretrained='laion_aesthetic_s13b_b82k')
        
        def new_forward(self, x):
            shortcut = x
            x = self.conv_dw(x)
            if self.use_conv_mlp:
                x = self.norm(x)
                x = self.mlp(x)
            else:
                x = x.permute(0, 2, 3, 1)
                x = self.norm(x)
                x = self.mlp(x)
                x = x.permute(0, 3, 1, 2)
            if self.gamma is not None:
                x = x.mul(self.gamma.reshape(1, -1, 1, 1))

            x = self.drop_path(x) + self.shortcut(shortcut)
            self.block_output = rearrange(x.clone(), "b c h w -> b h w c")
            return x

        setattr(model.visual.trunk.stages[0].blocks[0].__class__, "forward", new_forward)
        
        self.model = model
        self.model.eval()
    
    def forward(self, x):
        out = self.model(x)
        block_outputs = []
        for stage in self.model.visual.trunk.stages:
            for block in stage.blocks:
                block_outputs.append(block.block_output)
        return {
            'attn': None,
            'mlp': None,
            'block': block_outputs
        }
        
    
MODEL_DICT["CLIP(convnext_base_w_320/laion_aesthetic_s13b_b82k)"] = partial(CLIPConvnext)
LAYER_DICT["CLIP(convnext_base_w_320/laion_aesthetic_s13b_b82k)"] = 36
RES_DICT["CLIP(convnext_base_w_320/laion_aesthetic_s13b_b82k)"] = (960, 960)


class MAE(timm.models.vision_transformer.VisionTransformer):
    def __init__(self, **kwargs):
        super(MAE, self).__init__(**kwargs)

        sd = torch.hub.load_state_dict_from_url(
            "https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth"
        )

        checkpoint_model = sd["model"]
        state_dict = self.state_dict()
        for k in ["head.weight", "head.bias"]:
            if (
                k in checkpoint_model
                and checkpoint_model[k].shape != state_dict[k].shape
            ):
                print(f"Removing key {k} from pretrained checkpoint")
                del checkpoint_model[k]

        # load pre-trained model
        msg = self.load_state_dict(checkpoint_model, strict=False)
        print(msg)
        
        # resample the patch embeddings to 56x56, take 896x896 input
        pos_embed = self.pos_embed[0]
        pos_embed = resample_position_embeddings(pos_embed, 42, 42)
        self.pos_embed = nn.Parameter(pos_embed.unsqueeze(0))
        self.img_size = (672, 672)
        self.patch_embed.img_size = (672, 672)

        self.requires_grad_(False)
        self.eval()
        
        def forward(self, x):
            self.saved_attn_node = self.ls1(self.attn(self.norm1(x)))
            x = x + self.saved_attn_node.clone()
            self.saved_mlp_node = self.ls2(self.mlp(self.norm2(x)))
            x = x + self.saved_mlp_node.clone()
            self.saved_block_output = x.clone()
            return x
        
        setattr(self.blocks[0].__class__, "forward", forward)
        
    def forward(self, x):
        out = super().forward(x)
        def remove_cls_and_reshape(x):
            x = x.clone()
            x = x[:, 1:]
            hw = np.sqrt(x.shape[1]).astype(int)
            x = rearrange(x, "b (h w) c -> b h w c", h=hw)
            return x
        
        attn_outputs = [remove_cls_and_reshape(block.saved_attn_node) for block in self.blocks]
        mlp_outputs = [remove_cls_and_reshape(block.saved_mlp_node) for block in self.blocks]
        block_outputs = [remove_cls_and_reshape(block.saved_block_output) for block in self.blocks]
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }
        

MODEL_DICT["MAE(vit_base)"] = partial(MAE)
LAYER_DICT["MAE(vit_base)"] = 12
RES_DICT["MAE(vit_base)"] = (672, 672)

class ImageNet(nn.Module):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        model = timm.create_model(
            'vit_base_patch16_224.augreg2_in21k_ft_in1k',
            pretrained=True,
            num_classes=0,  # remove classifier nn.Linear
        )
        
        # resample the patch embeddings to 56x56, take 896x896 input
        pos_embed = model.pos_embed[0]
        pos_embed = resample_position_embeddings(pos_embed, 42, 42)
        model.pos_embed = nn.Parameter(pos_embed.unsqueeze(0))
        model.img_size = (672, 672)
        model.patch_embed.img_size = (672, 672)

        model.requires_grad_(False)
        model.eval()
        
        def forward(self, x):
            self.saved_attn_node = self.ls1(self.attn(self.norm1(x)))
            x = x + self.saved_attn_node.clone()
            self.saved_mlp_node = self.ls2(self.mlp(self.norm2(x)))
            x = x + self.saved_mlp_node.clone()
            self.saved_block_output = x.clone()
            return x
        
        setattr(model.blocks[0].__class__, "forward", forward)
        
        self.model = model
        
    def forward(self, x):
        out = self.model(x)
        def remove_cls_and_reshape(x):
            x = x.clone()
            x = x[:, 1:]
            hw = np.sqrt(x.shape[1]).astype(int)
            x = rearrange(x, "b (h w) c -> b h w c", h=hw)
            return x
        
        attn_outputs = [remove_cls_and_reshape(block.saved_attn_node) for block in self.model.blocks]
        mlp_outputs = [remove_cls_and_reshape(block.saved_mlp_node) for block in self.model.blocks]
        block_outputs = [remove_cls_and_reshape(block.saved_block_output) for block in self.model.blocks]
        return {
            'attn': attn_outputs,
            'mlp': mlp_outputs,
            'block': block_outputs
        }
            
MODEL_DICT["ImageNet(vit_base)"] = partial(ImageNet)
LAYER_DICT["ImageNet(vit_base)"] = 12
RES_DICT["ImageNet(vit_base)"] = (672, 672)

def download_all_models():
    for model_name in MODEL_DICT:
        print(f"Downloading {model_name}")
        try:
            model = MODEL_DICT[model_name]()
        except Exception as e:
            print(f"Error downloading {model_name}: {e}")
            continue
        
def get_all_model_names():
    return list(MODEL_DICT.keys())
        
def get_model(model_name):
    return MODEL_DICT[model_name]()

@torch.no_grad()
def extract_features(images, model, model_name, node_type, layer, batch_size=8):
    use_cuda = torch.cuda.is_available()
    
    if use_cuda:
        model = model.cuda()

    chunked_idxs = torch.split(torch.arange(images.shape[0]), batch_size)
    
    outputs = []
    for idxs in chunked_idxs:
        inp = images[idxs]
        if use_cuda:
            inp = inp.cuda()
        out = model(inp)  # {'attn': [B, H, W, C], 'mlp': [B, H, W, C], 'block': [B, H, W, C]}
        out = out[node_type]
        if out is None:
            raise ValueError(f"Node type {node_type} not found in model {model_name}")
        out = out[layer]
        # normalize
        out = F.normalize(out, dim=-1)
        outputs.append(out.cpu().float())
    outputs = torch.cat(outputs, dim=0)

    return outputs

    
if __name__ == '__main__':
    inp = torch.rand(1, 3, 1024, 1024)
    model = MAE()
    out = model(inp)
    print(out[0][0].shape, out[0][1].shape, out[0][2].shape)