# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import torch
import torch.nn as nn
from src.tools.tokenizer import load_bert_model, load_glove_embed
from src.model.vilbert.vilbert import BertConfig, BertModel

class OracleModel(nn.Module):
    def __init__(
        self,
        vilbert_config,
        dropout_prob,
        fusion_method,
        num_cats,
        cat_embed_size,
        cat_pad_id,
        mlp_hidden,
        num_classes,
        **kwargs
    ):
        super(OracleModel, self).__init__()
        vilbert_config = BertConfig.from_dict(vilbert_config)
        self.bert = BertModel(vilbert_config)
        self.dropout = nn.Dropout(dropout_prob)
        self.fusion_method = fusion_method
        self.cat_embed = nn.Embedding(
            num_cats, cat_embed_size, padding_idx=cat_pad_id)
        self.bi_hidden_size = vilbert_config.bi_hidden_size
        mlp_input_size = 2*self.bi_hidden_size + cat_embed_size
        self.mlp = nn.Sequential(
            nn.Linear(mlp_input_size, mlp_hidden),
            nn.ReLU(),
            nn.Linear(mlp_hidden, mlp_hidden//2),
            nn.ReLU(),
            nn.Linear(mlp_hidden//2, num_classes)
        )
        
    def load_pretrained_vilbert(self, path, map_location):
        ckpt = torch.load(path, map_location=map_location)
        mismatch = self.load_state_dict(ckpt, strict=False)
        return mismatch

    def forward(
        self, 
        wrds, 
        tgt_cat, 
        tgt_bbox, 
        tgt_img_feat, 
        bg_bboxs, 
        bg_img_feats, 
        update_vilbert=True, 
        token_type_ids=None,
        attention_mask=None,
        image_attention_mask=None,
        co_attention_mask=None,
        decode_mask=None,
        task_ids=None,
        output_all_encoded_layers=False,
        output_all_attention_masks=False,
        ):
        if update_vilbert:
            # with torch.no_grad():
            seq_out_txt, seq_out_vis, pooled_out_txt, pooled_out_vis, _ = self.bert(
                wrds,
                torch.cat([bg_img_feats, tgt_img_feat.unsqueeze(1)], dim=1),
                torch.cat([bg_bboxs, tgt_bbox.unsqueeze(1)], dim=1),
                token_type_ids,
                attention_mask,
                image_attention_mask,
                co_attention_mask,
                task_ids,
                output_all_encoded_layers=output_all_encoded_layers,
                output_all_attention_masks=output_all_attention_masks,
            )
        else:
            raise NotImplementedError
        
        tgt_feat = self.bert.v_pooler(seq_out_vis, idx=-1)
        bg_feat = pooled_out_vis
        txt_feat = pooled_out_txt
        if self.fusion_method == "sum":
            fused_tgt = self.dropout(txt_feat + tgt_feat)
            fused_bg = self.dropout(txt_feat + bg_feat)
        elif self.fusion_method == "mul":
            fused_tgt = self.dropout(txt_feat * tgt_feat)
            fused_bg = self.dropout(txt_feat * bg_feat)
        else:
            raise NotImplementedError
                        
        tgt_cat_embed = self.cat_embed(tgt_cat)
        mlp_input = torch.cat(
            [fused_tgt, fused_bg, tgt_cat_embed], dim=-1)
        logit = self.mlp(mlp_input)
        return logit

