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

import gradio as gr

MODEL_DICT = {}


def transform_images(images, resolution=(1024, 1024)):
    images = [image.convert("RGB").resize(resolution) for image in images]
    # Convert to torch tensor
    images = [
        torch.tensor(np.array(image).transpose(2, 0, 1)).float() / 255
        for image in images
    ]
    # Normalize
    images = [(image - 0.5) / 0.5 for image in images]
    images = torch.stack(images)
    return images


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")
        # 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()


def extract_features(images, model_name, node_type, layer):
    resolution_dict = {
        "MobileSAM": (1024, 1024),
        "SAM(sam_vit_b)": (1024, 1024),
        "DiNO(dinov2_vitb14_reg)": (448, 448),
        "CLIP(openai/clip-vit-base-patch16)": (224, 224),
    }
    images = transform_images(images, resolution=resolution_dict[model_name])

    model = MODEL_DICT[model_name]

    use_cuda = torch.cuda.is_available()
    if use_cuda:
        model = model.cuda()

    outputs = []
    for i in range(images.shape[0]):
        inp = images[i].unsqueeze(0)
        if use_cuda:
            inp = inp.cuda()
        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