# 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


class OracleModel(nn.Module):
    def __init__(
        self,
        num_wrds,
        wrd_embed_size,
        wrd_hidden_size,
        wrd_pad_id,
        num_cats,
        cat_embed_size,
        cat_pad_id,
        spatial_size,
        img_feat_size,
        img_hidden_size,
        mlp_hidden,
        num_classes,
        wrd_emb_mode='from_scratch',
        wrd_emb_path=None,
        num_attn_heads=None,
        **kwargs
    ):
        super(OracleModel, self).__init__()
        
        self.wrd_emb_mode = wrd_emb_mode
        if wrd_emb_mode == 'from_scratch':
            print("[INFO] Wrd Emb.: from scratch")
            self.wrd_embed = nn.Embedding(
                num_wrds, wrd_embed_size, padding_idx=wrd_pad_id)
        elif wrd_emb_mode == 'bert':
            print("[INFO] Wrd Emb.: BERT model.")
            self.wrd_embed = load_bert_model()
        elif wrd_emb_mode == 'glove':
            print("[INFO] Wrd Emb.: glove vectors.")
            self.wrd_embed = load_glove_embed(num_wrds, wrd_emb_path, pad_id=wrd_pad_id)
            

        self.lstm = nn.LSTM(
            wrd_embed_size, wrd_hidden_size, batch_first=True, bidirectional=False)

        self.img_mlp = nn.Sequential(
            nn.Linear(img_feat_size, img_hidden_size),
            nn.ReLU(),
            nn.Linear(img_hidden_size, img_hidden_size)
        )
        self.use_attn = num_attn_heads is not None
        if self.use_attn:
            print("[INFO] Use attention for word seq.")
            self.loc_attn = nn.MultiheadAttention(
                embed_dim=spatial_size, 
                num_heads=num_attn_heads, 
                kdim=wrd_hidden_size,
                vdim=wrd_hidden_size,
                )
            self.cat_attn = nn.MultiheadAttention(
                embed_dim=cat_embed_size, 
                num_heads=num_attn_heads, 
                kdim=wrd_hidden_size,
                vdim=wrd_hidden_size,
                )
            self.attn_proj = nn.Linear(cat_embed_size + spatial_size, wrd_hidden_size)

        self.cat_embed = nn.Embedding(
            num_cats, cat_embed_size, padding_idx=cat_pad_id)
        self.mlp = nn.Sequential(
            nn.Linear(wrd_hidden_size + cat_embed_size + spatial_size + img_hidden_size, mlp_hidden),
            nn.ReLU(),
            nn.Linear(mlp_hidden, wrd_embed_size),
            nn.ReLU(),
            nn.Linear(wrd_embed_size, num_classes)
        )

    def forward(self, wrds, tgt_cat, tgt_bbox, tgt_img_feat, wrd_len, update_wrd_embed=False):
        batch_size = wrds.size(0)
        if not update_wrd_embed:
            with torch.no_grad():
                wrd_embed = self.wrd_embed(wrds)
        else:
            wrd_embed = self.wrd_embed(wrds)
        if self.wrd_emb_mode == 'bert':
            # bert out: hidden state, attention
                wrd_embed = wrd_embed[0]

        wrd_embed = nn.utils.rnn.pack_padded_sequence(
                wrd_embed, wrd_len, batch_first=True, enforce_sorted=False)
        # Use the last hidden state of the LSTM as wrd_hidden
        wrd_hiddens, (last_wrd_hidden, _) = self.lstm(wrd_embed)
        wrd_hiddens, _ = nn.utils.rnn.pad_packed_sequence(
            wrd_hiddens, batch_first=True)

        # (batch_size, cat_embed_size)
        cat_embed = self.cat_embed(tgt_cat)
        # cat_embed = torch.zeros_like(cat_embed)
        img_feat = self.img_mlp(tgt_img_feat)
        # img_feat = torch.zeros_like(img_feat)
        if self.use_attn:
            cat_attn_out, _ = self.cat_attn(
                cat_embed.unsqueeze(0),  
                wrd_hiddens.transpose(0, 1), 
                wrd_hiddens.transpose(0, 1)
            )
            loc_attn_out, _ = self.loc_attn(
                tgt_bbox.unsqueeze(0), 
                wrd_hiddens.transpose(0, 1), 
                wrd_hiddens.transpose(0, 1)
            )
            wrd_hidden = self.attn_proj(
                torch.cat([cat_attn_out.squeeze(0), loc_attn_out.squeeze(0)], dim=-1))
            # print(wrd_hiddens.shape)
            # print(wrd_hidden.shape)
        else:
            last_wrd_hidden = last_wrd_hidden.view(batch_size, -1)
            wrd_hidden = last_wrd_hidden
        mlp_input = torch.cat(
            [wrd_hidden, cat_embed, tgt_bbox, img_feat], dim=-1)
        logit = self.mlp(mlp_input)
        return logit

