from typing import Optional, Tuple
from einops import rearrange
import torch
import torch.nn.functional as F
import timm
from PIL import Image
from torch import nn
import numpy as np
import os
import time

import gradio as gr

MODEL_DICT = {}


def transform_image(image, resolution=(1024, 1024), use_cuda=False):
    image = image.convert('RGB').resize(resolution, Image.Resampling.NEAREST)
    # Convert to torch tensor
    image = torch.tensor(np.array(image).transpose(2, 0, 1)).float()
    if use_cuda:
        image = image.cuda()
    image = image / 255
    # Normalize
    image = (image - 0.5) / 0.5
    return image

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_outputs, mlp_outputs, block_outputs


MODEL_DICT["MobileSAM"] = MobileSAM()


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_outputs, mlp_outputs, block_outputs


MODEL_DICT["SAM(sam_vit_b)"] = SAM()


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

        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] - 5).astype(int)
            self.attn_output = rearrange(
                attn_output.clone()[:, 5:], "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()[:, 5:], "b (h w) c -> b h w c", h=hw
            )
            x = x + mlp_output
            block_output = x
            self.block_output = rearrange(
                block_output.clone()[:, 5:], "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_outputs, mlp_outputs, block_outputs


MODEL_DICT["DiNO(dinov2_vitb14_reg)"] = DiNOv2()

class CLIP(torch.nn.Module):
    def __init__(self):
        super().__init__()

        from transformers import CLIPProcessor, CLIPModel

        model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
        
        # resample the patch embeddings to 64x64, take 1024x1024 input
        embeddings = model.vision_model.embeddings.position_embedding.weight
        cls_embeddings = embeddings[0]
        patch_embeddings = embeddings[1:]  # [14*14, 768]
        patch_embeddings = rearrange(patch_embeddings, "(h w) c -> c h w", h=14)
        patch_embeddings = F.interpolate(patch_embeddings.unsqueeze(0), size=(64, 64), mode="bilinear", align_corners=False).squeeze(0)
        patch_embeddings = rearrange(patch_embeddings, "c h w -> (h w) c")
        embeddings = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
        model.vision_model.embeddings.position_embedding.weight = nn.Parameter(embeddings)
        model.vision_model.embeddings.position_ids = torch.arange(0, 1+64*64)
        
        # 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)"] = CLIP()


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 64x64, take 1024x1024 input
        pos_embed = self.pos_embed[0]
        cls_embeddings = pos_embed[0]
        patch_embeddings = pos_embed[1:]  # [14*14, 768]
        patch_embeddings = rearrange(patch_embeddings, "(h w) c -> c h w", h=14)
        patch_embeddings = F.interpolate(patch_embeddings.unsqueeze(0), size=(64, 64), mode="bilinear", align_corners=False).squeeze(0)
        patch_embeddings = rearrange(patch_embeddings, "c h w -> (h w) c")
        pos_embed = torch.cat([cls_embeddings.unsqueeze(0), patch_embeddings], dim=0)
        self.pos_embed = nn.Parameter(pos_embed.unsqueeze(0))
        self.img_size = (1024, 1024)
        self.patch_embed.img_size = (1024, 1024)

        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_nodes = [remove_cls_and_reshape(block.saved_attn_node) for block in self.blocks]
        mlp_nodes = [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_nodes, mlp_nodes, block_outputs


MODEL_DICT["MAE(vit_base)"] = MAE()


def extract_features(images, model_name, node_type, layer):
    use_cuda = torch.cuda.is_available()
    
    resolution = (1024, 1024)
    resolution_dict = {
        "DiNO(dinov2_vitb14_reg)": (896, 896),
    }
    if model_name in resolution_dict:
        resolution = resolution_dict[model_name]

    model = MODEL_DICT[model_name]

    if use_cuda:
        model = model.cuda()

    outputs = []
    for i in range(len(images)):
        image = transform_image(images[i], resolution=resolution, use_cuda=use_cuda)
        inp = image.unsqueeze(0)
        attn_output, mlp_output, block_output = model(inp)
        out_dict = {
            "attn": attn_output,
            "mlp": mlp_output,
            "block": block_output,
        }
        out = out_dict[node_type]
        out = out[layer]
        outputs.append(out)
    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)