import torch
from mmengine.model import BaseModule
from torch import nn

from mmseg.registry import MODELS


@MODELS.register_module()
class DINOv2(nn.Module):
    """Use DINOv2 pre-trained models"""

    def __init__(self, version="large", freeze=False, load_from=None):
        super().__init__()

        if version == "large":
            self.dinov2 = torch.hub.load(
                "torchhub/facebookresearch_dinov2_main",
                "dinov2_vitl14",
                source="local",
                pretrained=False,
            )
        else:
            raise NotImplementedError

        if load_from is not None:
            d = torch.load(load_from, map_location="cpu")
            new_d = {}
            for key, value in d.items():
                if "pretrained" in key:
                    new_d[key.replace("pretrained.", "")] = value
            self.dinov2.load_state_dict(new_d)

        self.freeze = freeze

    def forward(self, inputs):
        B, _, h, w = inputs.shape

        if self.freeze:
            with torch.no_grad():
                features = self.dinov2.get_intermediate_layers(inputs, 4)
        else:
            features = self.dinov2.get_intermediate_layers(inputs, 4)

        outs = []
        for feature in features:
            C = feature.shape[-1]
            feature = (
                feature.permute(0, 2, 1).reshape(B, C, h // 14, w // 14).contiguous()
            )
            outs.append(feature)

        return outs
