import torch.nn.functional as F
from mmengine.model import BaseModel

from mmdet.registry import MODELS
from mmdet.utils import ConfigType, OptConfigType, OptMultiConfig


@MODELS.register_module()
class SAMSegmentor(BaseModel):
    MASK_THRESHOLD = 0.5

    def __init__(
            self,
            backbone: ConfigType,
            neck: ConfigType,
            prompt_encoder: ConfigType,
            mask_decoder: ConfigType,
            data_preprocessor: OptConfigType = None,
            fpn_neck: OptConfigType = None,
            init_cfg: OptMultiConfig = None,
            use_clip_feat: bool = False,
            use_head_feat: bool = False,
            use_gt_prompt: bool = False,
            use_point: bool = False,
            enable_backbone: bool = False,
    ) -> None:
        super().__init__(data_preprocessor=data_preprocessor, init_cfg=init_cfg)

        self.backbone = MODELS.build(backbone)
        self.neck = MODELS.build(neck)
        self.pe = MODELS.build(prompt_encoder)
        self.mask_decoder = MODELS.build(mask_decoder)
        if fpn_neck is not None:
            self.fpn_neck = MODELS.build(fpn_neck)
        else:
            self.fpn_neck = None

        self.use_clip_feat = use_clip_feat
        self.use_head_feat = use_head_feat
        self.use_gt_prompt = use_gt_prompt
        self.use_point = use_point

        self.enable_backbone = enable_backbone

    def extract_feat(self, inputs):
        backbone_feat = self.backbone(inputs)
        neck_feat = self.neck(backbone_feat)
        if self.fpn_neck is not None:
            fpn_feat = self.fpn_neck(backbone_feat)
        else:
            fpn_feat = None

        return dict(
            backbone_feat=backbone_feat,
            neck_feat=neck_feat,
            fpn_feat=fpn_feat
        )

    def extract_masks(self, feat_cache, prompts):
        sparse_embed, dense_embed = self.pe(
            prompts,
            image_size=(1024, 1024),
            with_points='point_coords' in prompts,
            with_bboxes='bboxes' in prompts,
        )

        kwargs = dict()
        if self.enable_backbone:
            kwargs['backbone_feats'] = feat_cache['backbone_feat']
            kwargs['backbone'] = self.backbone
            kwargs['fpn_feats'] = feat_cache['fpn_feat']
        low_res_masks, iou_predictions, cls_pred = self.mask_decoder(
            image_embeddings=feat_cache['neck_feat'],
            image_pe=self.pe.get_dense_pe(),
            sparse_prompt_embeddings=sparse_embed,
            dense_prompt_embeddings=dense_embed,
            multi_mask_output=False,
            **kwargs
        )
        masks = F.interpolate(
            low_res_masks,
            scale_factor=4.,
            mode='bilinear',
            align_corners=False,
        )

        masks = masks.sigmoid()
        cls_pred = cls_pred.softmax(-1)[..., :-1]
        return masks.detach().cpu().numpy(), cls_pred.detach().cpu()

    def forward(self, inputs):
        return inputs