# 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.model.vilbert.vilbert import BertConfig, BertModel

class GuesserModel(nn.Module):
    def __init__(
        self,
        vilbert_config,
        ans_embed_size,
        num_cats,
        cat_embed_size,
        cat_pad_id,
        dropout_prob,
        state_alpha,
        use_category=True,
        **kwargs
    ):
        super(GuesserModel, self).__init__()

        self.use_category = use_category
        # Vilbert part
        vilbert_config = BertConfig.from_dict(vilbert_config)
        self.bert = BertModel(vilbert_config)
        # self.fusion_method = fusion_method
        self.bi_hidden_size = vilbert_config.bi_hidden_size
        # Yes / No / NA
        self.ans_embed = nn.Embedding(3, ans_embed_size)
        self.ans_proj = nn.Linear(ans_embed_size, self.bi_hidden_size)
        self.dropout = nn.Dropout(dropout_prob)
        fused_proj_out = self.bi_hidden_size//2 if use_category else self.bi_hidden_size
        self.fused_proj = nn.Linear(self.bi_hidden_size, fused_proj_out)
        self.state_alpha = state_alpha
        self.state_proj = nn.Sequential(
            nn.Linear(self.bi_hidden_size, self.bi_hidden_size//2),
            nn.ReLU(),
            nn.Linear(self.bi_hidden_size//2, 1)
        )
        self.softmax = nn.Softmax(dim=-1)
        if use_category:
            self.cat_embed = nn.Embedding(
                num_cats, self.bi_hidden_size//2, padding_idx=cat_pad_id)

    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 init_state(self, batch_size, num_bboxs, device):
        return torch.ones(batch_size, num_bboxs).to(device) / num_bboxs

    def compute_next_state(
        self, curr_state, seq_out_vis, sent_feat, ans_emb, cats_emb, bboxs_mask=None):
        # curr_state: (batch_size, num_bboxs)
        # seq_out_vis:  (batch_size, num_bboxs, dim of vis_feat)
        # sent_feat:  (batch_size, dim of sent_feat)
        
        vis_feats = self.bert.v_pooler(seq_out_vis, idx=None)
        fused_feats = self.dropout(sent_feat.unsqueeze(1) * vis_feats)
        fused_feats = self.fused_proj(fused_feats)
        if self.use_category:
            fused_feats = torch.cat([fused_feats, cats_emb], dim=-1)

        fused_feats = fused_feats * curr_state.unsqueeze(-1)
        fuesd_feats = fused_feats + self.ans_proj(ans_emb).unsqueeze(1)
        # batch_size, num_bboxs
        logits = self.state_proj(fuesd_feats).squeeze(-1)
        
        if bboxs_mask is not None:
            logits[~bboxs_mask] = -1e10
        stat = self.softmax(logits)
        stat = (1-self.state_alpha) * curr_state + self.state_alpha * stat
        return stat, logits

    def forward_turn(
        self, 
        q, 
        ans, 
        cats, 
        img_feats,
        bboxs, 
        curr_state=None,
        bboxs_mask=None,
        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,
        update_vilbert=True,
        ):
        if update_vilbert:
            seq_out_txt, seq_out_vis, pooled_out_txt, pooled_out_vis, _ = self.bert(
                q,
                img_feats,
                bboxs,
                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:
            with torch.no_grad():
                # output of each  token, output of each object, output of cls token, output of global vis
                seq_out_txt, seq_out_vis, pooled_out_txt, pooled_out_vis, _ = self.bert(
                    q,
                    img_feats,
                    bboxs,
                    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,
                )
        ans = self.ans_embed(ans)
        cats = self.cat_embed(cats) if self.use_category else None
        if curr_state is None:
            curr_state = self.init_state(
                img_feats.size(0), img_feats.size(1), img_feats.device)
        stat, logits = self.compute_next_state(
            curr_state, seq_out_vis, pooled_out_txt, ans, cats, bboxs_mask)
        return stat, logits

    def forward_session(
        self, 
        qs, 
        ans, 
        end_turn,
        cats, 
        img_feats,
        bboxs, 
        bboxs_mask=None,
        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,
        update_vilbert=True,
        return_state_history=False,
        ):

        stat = self.init_state(
            img_feats.size(0), img_feats.size(1), img_feats.device)
        batch_size = qs.size(0)
        max_turns = qs.size(1)
        
        stat_his = []
        final_logits = torch.zeros_like(stat)
        for t in range(max_turns):
            q_t = qs[:, t]
            ans_t = ans[:, t]
            next_stat, logits = self.forward_turn(
                q_t, ans_t, cats, img_feats, bboxs, 
                curr_state=stat,
                bboxs_mask=bboxs_mask,
                token_type_ids=token_type_ids,
                attention_mask=attention_mask[:, t],
                image_attention_mask=image_attention_mask,
                co_attention_mask=co_attention_mask,
                decode_mask=decode_mask,
                task_ids=task_ids,
                output_all_encoded_layers=output_all_encoded_layers,
                output_all_attention_masks=output_all_attention_masks,
                update_vilbert=update_vilbert,
                )
            end = end_turn == t
            final_logits[end] = logits[end]
            stat = next_stat
            stat_his.append(stat)
        
        stat_his = torch.stack(stat_his).transpose(0, 1)
        if return_state_history:
            return final_logits[:, 1:], stat_his
        else:
            # First one is for global feat.
            return final_logits[:, 1:]
        
