"""Hugging Face model implementation for the Mettle tile encoder."""

from dataclasses import dataclass
from typing import Optional

import timm
import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

try:
    from .configuration_mettle import MettleConfig
except ImportError:
    from configuration_mettle import MettleConfig


class MettleRefineHead(nn.Module):
    """Content-conditional refinement of the pooled CLS embedding."""

    def __init__(self, dim, num_atoms, rank, hidden_size):
        super().__init__()
        self.P = nn.Parameter(torch.zeros(num_atoms, dim, rank))
        self.Q = nn.Parameter(torch.zeros(num_atoms, dim, rank))
        self.mix = nn.Sequential(
            nn.Linear(dim, hidden_size),
            nn.GELU(),
            nn.Linear(hidden_size, num_atoms),
        )
        self.shift = nn.Sequential(
            nn.Linear(dim, hidden_size),
            nn.GELU(),
            nn.Linear(hidden_size, dim),
        )

    def forward(self, embedding):
        mixture = self.mix(embedding)
        shift = self.shift(embedding)
        projected = torch.einsum("bd,kdr->bkr", embedding, self.P)
        delta = torch.einsum("bkr,kdr->bkd", projected, self.Q)
        return embedding + (mixture.unsqueeze(-1) * delta).sum(1) + shift


@dataclass
class MettleModelOutput(ModelOutput):
    """Outputs from :class:`MettleModel`.

    ``pooler_output`` follows the requested feature view. ``cls_embedding`` is
    the headed 1536-dimensional CLS representation. ``cls_mean_embedding`` is
    the concatenation of headed CLS and the unheaded mean spatial patch token.
    """

    last_hidden_state: Optional[torch.Tensor] = None
    pooler_output: Optional[torch.Tensor] = None
    cls_embedding: Optional[torch.Tensor] = None
    mean_patch_embedding: Optional[torch.Tensor] = None
    cls_mean_embedding: Optional[torch.Tensor] = None


class MettleModel(PreTrainedModel):
    """Mettle scanner-resilient pathology tile encoder."""

    config_class = MettleConfig
    base_model_prefix = "backbone"
    main_input_name = "pixel_values"

    def __init__(self, config):
        super().__init__(config)
        self.backbone = timm.create_model(
            config.backbone_name,
            pretrained=False,
            num_classes=0,
            init_values=1e-5,
            dynamic_img_size=False,
            img_size=config.image_size,
        )
        if int(self.backbone.num_features) != config.hidden_size:
            raise ValueError(
                f"{config.backbone_name} has width {self.backbone.num_features}, "
                f"but the config declares {config.hidden_size}"
            )
        self.head = None
        if config.head_enabled:
            self.head = MettleRefineHead(
                dim=config.hidden_size,
                num_atoms=config.head_num_atoms,
                rank=config.head_rank,
                hidden_size=config.head_hidden_size,
            )
        self.post_init()

    def _embeddings(self, pixel_values):
        tokens = self.backbone.forward_features(pixel_values)
        expected_prefix = 1 + self.config.num_register_tokens
        actual_prefix = int(self.backbone.num_prefix_tokens)
        if actual_prefix != expected_prefix:
            raise RuntimeError(
                f"expected {expected_prefix} prefix tokens, found {actual_prefix}"
            )
        patches = tokens[:, actual_prefix:]
        if patches.shape[1] == 0:
            raise RuntimeError("the backbone returned no spatial patch tokens")

        cls_embedding = tokens[:, 0].to(torch.float32)
        if self.head is not None:
            cls_embedding = self.head(cls_embedding)
        mean_patch_embedding = patches.to(torch.float32).mean(dim=1)
        cls_mean_embedding = torch.cat(
            (cls_embedding, mean_patch_embedding),
            dim=-1,
        )
        return tokens, cls_embedding, mean_patch_embedding, cls_mean_embedding

    def forward(
        self,
        pixel_values,
        feature_view=None,
        return_dict=None,
    ):
        """Encode normalized 224-pixel RGB tiles.

        Args:
            pixel_values: Tensor with shape ``(batch, 3, 224, 224)`` after the
                preprocessing specified by ``preprocessor_config.json``.
            feature_view: ``"cls"`` for a 1536-dimensional embedding or
                ``"cls_mean"`` for headed CLS concatenated with mean spatial
                patch pooling (3072 dimensions).
            return_dict: Follow the standard Transformers return convention.
        """
        feature_view = feature_view or self.config.default_feature_view
        if feature_view not in {"cls", "cls_mean"}:
            raise ValueError(
                "feature_view must be 'cls' or 'cls_mean', got "
                f"{feature_view!r}"
            )
        return_dict = (
            self.config.use_return_dict if return_dict is None else return_dict
        )
        tokens, cls_embedding, mean_patch, cls_mean = self._embeddings(
            pixel_values
        )
        selected = cls_embedding if feature_view == "cls" else cls_mean
        if not return_dict:
            return tokens, selected, cls_embedding, mean_patch, cls_mean
        return MettleModelOutput(
            last_hidden_state=tokens,
            pooler_output=selected,
            cls_embedding=cls_embedding,
            mean_patch_embedding=mean_patch,
            cls_mean_embedding=cls_mean,
        )

    def encode(self, pixel_values, feature_view=None):
        """Return only the selected tile embedding."""
        return self(
            pixel_values=pixel_values,
            feature_view=feature_view,
            return_dict=True,
        ).pooler_output

    def get_input_embeddings(self):
        """Return the ViT patch projection for Transformers integrations."""
        return self.backbone.patch_embed.proj