# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the Apache License, Version 2.0
# found in the LICENSE file in the root directory of this source tree.

# References:
#   https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
#   https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py

from functools import partial
import math
import logging
from typing import Sequence, Tuple, Union, Callable
from collections import OrderedDict

import torch
import torch.nn as nn
import torch.utils.checkpoint
from torch.nn.init import trunc_normal_

from .dinov2.hub.backbones import dinov2_vitb14

class FrozenDinoV2ImageEmbedder(nn.Module):
    """
        Uses the dinov2 image encoder with camera modulation.
        Not actually frozen... If you want that set cond_stage_trainable=False in cfg
        """
    def __init__(
            self,
            version='dinov2_vitb14',
            ckpt_path=None,
            lrm_mode='plain_lrm',
        ):
        super().__init__()
        self.lrm_mode = lrm_mode
        assert version in ['dinov2_vitb14', 'dinov2_vits14', 'dinov2_vitl14', 'dinov2_vitg14']

    
        self.model = dinov2_vitb14(pretrained=False)

        if ckpt_path is not None:
            self.load_pretrained(ckpt_path)
        else:
            print('None pretrained model for dinov2 encoder ...')


    def load_pretrained(self, ckpt_path):
        print('Loading dinov2 encoder ...')
        orig_state_dict = torch.load(ckpt_path, map_location='cpu')
        try:
            ret = self.model.load_state_dict(orig_state_dict, strict=False)
            print(ret)
            print('Successfully loaded orig state dict')
        except:
            new_state_dict = OrderedDict()
            for k, v in orig_state_dict['state_dict'].items():
                if 'img_encoder' in k:
                    new_state_dict[k.replace('img_encoder.model.', '')] = v
            ret = self.model.load_state_dict(new_state_dict, strict=False)
            print(ret)
            print('Successfully loaded new state dict')
            

    def forward(self, x, *args, **kwargs):
        ret = self.model.forward_features_with_camera(x, *args, **kwargs)
        output = torch.cat([ret['x_norm_clstoken'].unsqueeze(1), ret['x_norm_patchtokens']], dim=1)
        return output