import torch
from transformers import PreTrainedModel

from .configuration_xmag import XMagConfig


class XMagModel(PreTrainedModel):
    config_class = XMagConfig
    base_model_prefix = "model"

    def __init__(self, config):
        super().__init__(config)
        self.model = torch.hub.load("facebookresearch/dinov2", config.backbone_name)

    def forward(self, pixel_values):
        if pixel_values.ndim != 4:
            raise ValueError("pixel_values must have shape (B, 3, H, W).")

        _, channels, height, width = pixel_values.shape
        if channels != 3:
            raise ValueError("pixel_values must have 3 RGB channels.")
        if height != self.config.image_size or width != self.config.image_size:
            raise ValueError(f"pixel_values must be {self.config.image_size}x{self.config.image_size}.")

        outputs = self.model.forward_features(pixel_values)

        return {
            "cls_embedding": outputs["x_norm_clstoken"],
            "patch_embeddings": outputs["x_norm_patchtokens"],
        }