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

from mmseg.registry import MODELS
import os
_DINOV2_MMSEG_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
_DINOV2_TORCHHUB_DIR = os.path.join(_DINOV2_MMSEG_ROOT, 'torchhub', 'facebookresearch_dinov2_main')


@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(_DINOV2_TORCHHUB_DIR, 'dinov2_vitl14', source='local', pretrained=False)
        else:
            raise NotImplementedError

        if load_from is not None:
            if load_from.split('/')[-1] == 'depth_anything_vitl14.pth':
                print(load_from)
                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)
            else:
                print(load_from)
                all_d = torch.load(load_from, map_location='cpu')
                d = all_d['state_dict']
                new_d = {}
                for key, value in d.items():
                    if 'backbone.dinov2' in key:
                        new_d[key.replace('backbone.dinov2.', '')] = 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