import torch
from torch import nn
from PIL import Image
from einops import rearrange
from torchvision.transforms.v2 import (
    Compose,
    Resize,
    InterpolationMode,
    ToImage,
    ToDtype,
    Normalize,
)
import timm


class VisualHolder(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.visual = model

    def forward(self, x):
        return self.visual(x)


class ModelHolder(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, x):
        return self.model(x)


class LinearPatchEmbedding(nn.Module):
    def __init__(self, conv):
        super().__init__()
        self.linear = nn.Linear(588, 1152)
        self.linear.weight.data = conv.weight.data.view(1152, -1)
        if conv.bias is not None:
            self.linear.bias.data = conv.bias.data

    def forward(self, x):
        return self.linear(x)


class MLP(nn.Module):
    def __init__(
        self,
        in_features: int,
        hidden_features: int = None,
        out_features: int = None,
        act_layer: nn.Module = nn.GELU,
    ) -> None:
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)

        torch.nn.init.kaiming_normal_(
            self.fc1.weight, mode="fan_in", nonlinearity="relu"
        )
        torch.nn.init.kaiming_normal_(
            self.fc2.weight, mode="fan_in", nonlinearity="relu"
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x


class VisionProjection(nn.Module):
    def __init__(self):
        super().__init__()

        image_embedding_dim = 1152
        model_dim = 2048
        hidden_dim = model_dim * 4

        self.mlp1 = MLP(image_embedding_dim, hidden_dim, model_dim)
        self.mlp2 = MLP(model_dim, hidden_dim, model_dim)
        self.ln = nn.LayerNorm(model_dim)

    @property
    def device(self):
        return self.mlp1.fc1.weight.device

    def forward(self, x):
        x = self.mlp1(x)
        x = self.ln(x)
        x = x + self.mlp2(x)
        return x


class VisionTower(nn.Module):
    def __init__(self):
        super().__init__()

        self.encoder = ModelHolder(
            VisualHolder(timm.create_model("vit_so400m_patch14_siglip_384"))
        )
        self.encoder.model.visual.patch_embed = LinearPatchEmbedding(
            self.encoder.model.visual.patch_embed.proj
        )
        self.encoder.model.visual.attn_pool = nn.Identity()

        self.projection = VisionProjection()

    def forward(self, x):
        x = self.encoder(x)
        x = self.projection(x)
        return x


class VisionEncoder(nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.model = VisionTower()
        self.preprocess = Compose(
            [
                Resize(size=(378, 378), interpolation=InterpolationMode.BICUBIC),
                ToImage(),
                ToDtype(torch.float32, scale=True),
                Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
            ]
        )

    @property
    def device(self):
        return self.model.projection.mlp1.fc1.weight.device

    @property
    def dtype(self):
        return self.model.projection.mlp1.fc1.weight.dtype

    def __call__(self, image: Image) -> torch.Tensor:
        with torch.no_grad():
            image_vec = (
                self.preprocess(image.convert("RGB"))
                .unsqueeze(0)
                .to(self.device, dtype=self.dtype)
            )
            image_vec = rearrange(
                image_vec, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=14, p2=14
            )
            return self.model(image_vec)