"""DINOv3 feature extraction via HuggingFace, handling CLS and register tokens."""

import os
import torch
import torch.nn as nn
from transformers import DINOv3ViTConfig, DINOv3ViTModel, DINOv3ViTImageProcessorFast


class DINOv3HFExtractor(nn.Module):
    """
    Extracts intermediate features from DINOv3 via HuggingFace transformers.

    Builds the model from config (no gated download required); weights are loaded
    from the MMDiff checkpoint which bundles the DINOv3 backbone.

    Returns 4 feature maps of shape [B, C_dino, H//16, W//16] from selected layers.
    Input images must be [B, 3, H, W] in [0, 1] range.
    """

    def __init__(self, repo_id="facebook/dinov3-vitb16-pretrain-lvd1689m",
                 take_last=None, take_indices=None, trainable=False,
                 hidden_size=768, num_hidden_layers=12, num_attention_heads=12,
                 intermediate_size=3072, patch_size=16, image_size=512,
                 num_register_tokens=4):
        super().__init__()

        # Build image processor from default config (no gated download needed)
        self.proc = DINOv3ViTImageProcessorFast()

        # Disable resizing/cropping so native resolution maps to patches
        for k in ("do_resize", "do_center_crop"):
            if hasattr(self.proc, k):
                setattr(self.proc, k, False)

        # Build model from config (random weights; real weights loaded from checkpoint)
        config = DINOv3ViTConfig(
            hidden_size=hidden_size,
            num_hidden_layers=num_hidden_layers,
            num_attention_heads=num_attention_heads,
            intermediate_size=intermediate_size,
            patch_size=patch_size,
            image_size=image_size,
            num_register_tokens=num_register_tokens,
            hidden_act="gelu",
        )
        self.model = DINOv3ViTModel(config)
        self.model.config.output_hidden_states = True

        self._frozen = not trainable
        if self._frozen:
            self.model.eval()
            for p in self.model.parameters():
                p.requires_grad = False
        else:
            self.model.train()
            for p in self.model.parameters():
                p.requires_grad = True

        # ImageNet normalization stats
        mean = torch.tensor(self.proc.image_mean).view(1, 3, 1, 1)
        std = torch.tensor(self.proc.image_std).view(1, 3, 1, 1)
        self.register_buffer("mean", mean, persistent=False)
        self.register_buffer("std", std, persistent=False)

        if take_indices is not None:
            self.take_indices = take_indices
            self.take_last = None
        else:
            self.take_last = take_last if take_last is not None else 4
            self.take_indices = None

        self.patch_size = getattr(self.model.config, "patch_size", 16)
        self.num_register_tokens = getattr(self.model.config, "num_register_tokens", 0)

        hidden_size = getattr(self.model.config, "hidden_size", 768)

        layers = self.take_indices if self.take_indices is not None else f"last {self.take_last}"
        trainable_str = "trainable" if not self._frozen else "frozen"
        print(f"[DINOv3] Built from config: dim={hidden_size}, patch={self.patch_size}, "
              f"layers={layers} ({trainable_str})")

    def train(self, mode: bool = True):
        self.training = mode
        if self._frozen:
            self.model.eval()
        else:
            self.model.train(mode)
        return self

    def forward(self, images_512: torch.Tensor):
        with torch.set_grad_enabled(not self._frozen):
            return self._forward(images_512)

    def _forward(self, images_512: torch.Tensor):
        x = (images_512 - self.mean) / self.std

        out = self.model(pixel_values=x, output_hidden_states=True)
        hidden_states = out.hidden_states

        B, _, H, W = images_512.shape
        H_patches = H // self.patch_size
        W_patches = W // self.patch_size
        P = H_patches * W_patches
        R = self.num_register_tokens

        maps = []
        if self.take_indices is not None:
            for idx in self.take_indices:
                hidden = hidden_states[idx]
                spatial = hidden[:, 1:1+P, :]
                C = spatial.shape[-1]
                spatial_map = spatial.transpose(1, 2).reshape(B, C, H_patches, W_patches).contiguous()
                maps.append(spatial_map)
        else:
            for hidden in hidden_states[-self.take_last:]:
                spatial = hidden[:, 1:1+P, :]
                C = spatial.shape[-1]
                spatial_map = spatial.transpose(1, 2).reshape(B, C, H_patches, W_patches).contiguous()
                maps.append(spatial_map)

        return maps


def create_dinov3_hf_extractor(repo_id="facebook/dinov3-vitb16-pretrain-lvd1689m", take_last=None, take_indices=None, trainable=False):
    """
    Factory for DINOv3HFExtractor (frozen in eval mode unless trainable=True).
    Builds from config — no gated download required.
    """
    return DINOv3HFExtractor(
        repo_id=repo_id, take_last=take_last, take_indices=take_indices, trainable=trainable,
    )