from functools import partial
from typing import Literal

import torch
import torch.nn as nn
from mmdet.registry import MODELS

from mmengine.model import BaseModule
from mmengine.logging import MMLogger

from ext.sam import ImageEncoderViT
from ext.meta.sam_meta import meta_dict, checkpoint_dict
from utils.load_checkpoint import load_checkpoint_with_prefix


@MODELS.register_module()
class SAMBackbone(BaseModule):

    def __init__(
            self,
            model_name: Literal['vit_h', 'vit_l', 'vit_b'] = 'vit_h',
            fix: bool = True,
            init_cfg=None,
    ):
        assert init_cfg is not None and init_cfg['type'] in \
               ['sam_pretrain', 'Pretrained'], f"{init_cfg['type']} is not supported."
        pretrained = init_cfg['checkpoint']
        super().__init__(init_cfg=None)
        self.init_cfg = init_cfg
        self.logger = MMLogger.get_current_instance()

        backbone_meta = meta_dict[model_name]

        backbone = ImageEncoderViT(
            depth=backbone_meta['encoder_depth'],
            embed_dim=backbone_meta['encoder_embed_dim'],
            num_heads=backbone_meta['encoder_num_heads'],
            patch_size=backbone_meta['vit_patch_size'],
            img_size=backbone_meta['image_size'],
            global_attn_indexes=backbone_meta['encoder_global_attn_indexes'],
            out_chans=backbone_meta['prompt_embed_dim'],
            norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
            qkv_bias=True,
            use_rel_pos=True,
            mlp_ratio=4,
            window_size=14,
        )
        if self.init_cfg['type'] == 'sam_pretrain':
            checkpoint_path = checkpoint_dict[pretrained]
            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix='image_encoder')
            backbone.load_state_dict(state_dict, strict=True)

        self.stem = backbone.patch_embed
        self.pos_embed = backbone.pos_embed

        self.res_layers = []
        last_pos = 0
        for idx, cur_pos in enumerate(backbone_meta['encoder_global_attn_indexes']):
            blocks = backbone.blocks[last_pos:cur_pos + 1]
            layer_name = f'layer{idx + 1}'
            self.add_module(layer_name, nn.Sequential(*blocks))
            self.res_layers.append(layer_name)
            last_pos = cur_pos + 1

        self.out_proj = backbone.neck

        if self.init_cfg['type'] == 'Pretrained':
            checkpoint_path = pretrained
            state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix=self.init_cfg['prefix'])
            self.load_state_dict(state_dict, strict=True)

        self.model_name = model_name
        self.fix = fix
        self.model_type = 'vit'
        self.output_channels = None
        self.out_indices = (0, 1, 2, 3)
        if self.fix:
            self.train(mode=False)
            for name, param in self.named_parameters():
                param.requires_grad = False

    def init_weights(self):
        self.logger.info(f"Init Config for {self.model_name}")
        self.logger.info(self.init_cfg)

    def train(self: torch.nn.Module, mode: bool = True) -> torch.nn.Module:
        if not isinstance(mode, bool):
            raise ValueError("training mode is expected to be boolean")
        if self.fix:
            super().train(mode=False)
        else:
            super().train(mode=mode)
        return self

    def forward_func(self, x):
        x = self.stem(x)
        x = x + self.pos_embed
        outs = []
        for i, layer_name in enumerate(self.res_layers):
            res_layer = getattr(self, layer_name)
            x = res_layer(x)
            if i in self.out_indices:
                outs.append(x.permute(0, 3, 1, 2).contiguous())
        outs[-1] = self.out_proj(outs[-1])
        return tuple(outs)

    def forward(self, x):
        if self.fix:
            with torch.no_grad():
                outs = self.forward_func(x)
        else:
            outs = self.forward_func(x)
        return outs