# This file is a modified version of the Vision Transformer - Pytorch implementation
# https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py
from collections.abc import Callable

from einops import rearrange, repeat
from einops.layers.torch import Rearrange
import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from transformers.utils import logging

from .configuration_metom import MetomConfig

logger = logging.get_logger(__name__)


def size_pair(t):
    return t if isinstance(t, tuple) else (t, t)


def metom_eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float | None = None,
    dropout: float = 0.0,
    **kwargs,
) -> tuple[torch.Tensor, torch.Tensor]:
    if scaling is None:
        scaling = query.size(-1) ** -0.5

    attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
    if attention_mask is not None:
        attention_mask = attention_mask[:, :, :, : key.shape[-2]]
        attn_weights = attn_weights + attention_mask

    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value)
    attn_output = attn_output.transpose(1, 2).contiguous()
    return attn_output, attn_weights


class MetomFeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout),
        )

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


class MetomAttention(nn.Module):
    def __init__(self, config: MetomConfig):
        super().__init__()
        inner_dim = config.dim_head * config.heads
        project_out = not (config.heads == 1 and config.dim_head == config.dim)

        self.config = config
        self.heads = config.heads
        self.scale = config.dim_head ** -0.5
        self.norm = nn.LayerNorm(config.dim)
        self.dropout = nn.Dropout(config.dropout)
        self.to_qkv = nn.Linear(config.dim, inner_dim * 3, bias=False)
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, config.dim),
            nn.Dropout(config.dropout),
        ) if project_out else nn.Identity()

    def forward(self, x: torch.Tensor, **kwargs):
        x = self.norm(x)
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), qkv)

        attn_implementation = self.config._attn_implementation or "eager"
        if attn_implementation == "flex_attention":
            if self.training and self.dropout.p > 0:
                logger.warning_once(
                    "`flex_attention` does not support attention dropout during training. Falling back to `eager`."
                )
                attn_implementation = "eager"
        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            attn_implementation,
            metom_eager_attention_forward,
        )
        out, _ = attention_interface(
            self,
            q,
            k,
            v,
            None,
            is_causal=False,
            scaling=self.scale,
            dropout=0.0 if not self.training else self.dropout.p,
            **kwargs,
        )
        out = rearrange(out, "b n h d -> b n (h d)")
        return self.to_out(out)


class MetomTransformer(nn.Module):
    def __init__(self, config: MetomConfig):
        super().__init__()
        self.norm = nn.LayerNorm(config.dim)
        self.layers = nn.ModuleList([])
        for _ in range(config.depth):
            self.layers.append(
                nn.ModuleList(
                    [
                        MetomAttention(config),
                        MetomFeedForward(config.dim, config.mlp_dim, dropout=config.dropout),
                    ]
                )
            )

    def forward(self, x: torch.Tensor, **kwargs):
        for attn, ff in self.layers:
            x = attn(x, **kwargs) + x
            x = ff(x) + x
        return self.norm(x)


class MetomModel(PreTrainedModel):
    config_class = MetomConfig
    main_input_name = "pixel_values"
    _supports_attention_backend = True
    _supports_flash_attn = True
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    _supports_flex_attn = True

    def __init__(self, config: MetomConfig):
        super().__init__(config)
        image_height, image_width = size_pair(config.image_size)
        patch_height, patch_width = size_pair(config.patch_size)
        assert image_height % patch_height == 0 and image_width % patch_width == 0, (
            "Image dimensions must be divisible by the patch size."
        )

        num_patches = (image_height // patch_height) * (image_width // patch_width)
        patch_dim = config.channels * patch_height * patch_width
        assert config.pool in {"cls", "mean"}, "pool type must be either cls (cls token) or mean (mean pooling)"
        assert len(config.labels) > 0, "labels must be composed of at least one label"

        self.to_patch_embedding = nn.Sequential(
            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=patch_height, p2=patch_width),
            nn.LayerNorm(patch_dim),
            nn.Linear(patch_dim, config.dim),
            nn.LayerNorm(config.dim),
        )
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, config.dim))
        self.cls_token = nn.Parameter(torch.randn(1, 1, config.dim))
        self.dropout = nn.Dropout(config.emb_dropout)
        self.transformer = MetomTransformer(config)
        self.pool = config.pool
        self.to_latent = nn.Identity()
        self.mlp_head = nn.Linear(config.dim, len(config.labels))
        self.labels = config.labels
        self.post_init()

    def forward(self, pixel_values: torch.Tensor, **kwargs):
        x = self.to_patch_embedding(pixel_values)
        b, n, _ = x.shape
        cls_tokens = repeat(self.cls_token, "1 1 d -> b 1 d", b=b)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.pos_embedding[:, : (n + 1)]
        x = self.dropout(x)
        x = self.transformer(x, **kwargs)
        x = x.mean(dim=1) if self.pool == "mean" else x[:, 0]
        x = self.to_latent(x)
        return self.mlp_head(x)

    def get_predictions(self, pixel_values: torch.Tensor) -> list[str]:
        logits = self(pixel_values=pixel_values)
        indices = torch.argmax(logits, dim=-1)
        return [self.labels[i] for i in indices]

    def get_topk_labels(
        self, pixel_values: torch.Tensor, k: int = 5, return_probs: bool = False
    ) -> list[list[str]] | list[list[tuple[str, float]]]:
        assert 0 < k <= len(self.labels), "k must be a positive integer less than or equal to the number of labels"
        logits = self(pixel_values=pixel_values)
        probs = torch.softmax(logits, dim=-1)
        topk_probs, topk_indices = torch.topk(probs, k, dim=-1)
        topk_labels = [[self.labels[i] for i in ti] for ti in topk_indices]
        if return_probs:
            return [
                [(label, prob.item()) for label, prob in zip(labels, probs)]
                for labels, probs in zip(topk_labels, topk_probs)
            ]
        return topk_labels