"""HuggingFace modeling file for the distilled CLIP-mini student.

Architecture identical to the training bed's Student class: ViT with
d=240, depth 12, heads 4, patch 16, 160px input, CLS readout, linear
head to the 512-d projection space. If the checkpoint carries a
`rotation` buffer (the consensus-distilled champion does), forward
applies it by default, returning embeddings in the CLIP-B/16 (LAION-2B)
deployment frame — compatible with that teacher's text tower.

Inputs are standard `pixel_values`: images resized so the shorter edge
is 182 (bicubic), center-cropped to 160, rescaled to [0,1], normalized
with the CLIP mean/std (see preprocessor_config.json). The published
evaluations used the torchvision transform pipeline with exactly those
constants.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPooling

from .configuration_clip_mini import ClipMiniConfig


class _Block(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.n1 = nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.n2 = nn.LayerNorm(d)
        self.fc1 = nn.Linear(d, 4 * d)
        self.fc2 = nn.Linear(4 * d, d)
        self.heads = heads

    def forward(self, x):
        B, N, C = x.shape
        q, k, v = (self.qkv(self.n1(x))
                   .reshape(B, N, 3, self.heads, C // self.heads)
                   .permute(2, 0, 3, 1, 4))
        a = F.scaled_dot_product_attention(q, k, v)
        x = x + self.proj(a.transpose(1, 2).reshape(B, N, C))
        return x + self.fc2(F.gelu(self.fc1(self.n2(x))))


class ClipMiniModel(PreTrainedModel):
    config_class = ClipMiniConfig
    main_input_name = "pixel_values"

    def __init__(self, config):
        super().__init__(config)
        d = config.hidden_size
        n_tok = (config.image_size // config.patch_size) ** 2 + 1
        self.patch = nn.Conv2d(3, d, config.patch_size, config.patch_size)
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.pos = nn.Parameter(torch.zeros(1, n_tok, d))
        self.blocks = nn.ModuleList(
            _Block(d, config.num_attention_heads)
            for _ in range(config.num_hidden_layers))
        self.norm = nn.LayerNorm(d)
        self.head = nn.Linear(d, config.projection_dim)
        if config.has_rotation:
            self.register_buffer(
                "rotation", torch.eye(config.projection_dim), persistent=True)
        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            nn.init.trunc_normal_(module.weight, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward_features(self, pixel_values):
        x = self.patch(pixel_values).flatten(2).transpose(1, 2)
        x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], 1) + self.pos
        for b in self.blocks:
            x = b(x)
        return self.norm(x)[:, 0]

    def get_image_features(self, pixel_values, apply_rotation=None):
        """L2-normalized (B, 512) image embeddings. With the rotation
        applied (default when present), outputs live in the CLIP-B/16
        LAION-2B frame."""
        z = F.normalize(self.head(self.forward_features(pixel_values)),
                        dim=-1)
        rot = (self.config.apply_rotation
               if apply_rotation is None else apply_rotation)
        if rot and self.config.has_rotation:
            z = F.normalize(z.double() @ self.rotation.double(),
                            dim=-1).to(z.dtype)
        return z

    def forward(self, pixel_values, apply_rotation=None, **kwargs):
        feats = self.get_image_features(pixel_values, apply_rotation)
        return BaseModelOutputWithPooling(
            last_hidden_state=feats.unsqueeze(1), pooler_output=feats)