# 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


class GuesserModel(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,
        **kwargs
    ):
        super(GuesserModel, self).__init__()

        self.wrd_embed = nn.Embedding(
            num_wrds, wrd_embed_size, padding_idx=wrd_pad_id)

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

        self.cat_embed = nn.Embedding(
            num_cats, cat_embed_size, padding_idx=cat_pad_id)

        self.bbox_mlp = nn.Sequential(
            nn.Linear(cat_embed_size + spatial_size, wrd_hidden_size),
            nn.ReLU(),
            nn.Linear(wrd_hidden_size, wrd_hidden_size),
            nn.ReLU()
        )

    def forward(self, wrds, wrd_len, cats, bboxs, bboxs_mask=None):
        batch_size = wrds.size(0)
        wrd_embed = self.wrd_embed(wrds)
        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_hidden, _) = self.lstm(wrd_embed)
        # (batch_size, wrd_hidden_size)
        wrd_hidden = wrd_hidden.view(batch_size, -1)
        # (batch_size, num_bboxs, cat_embed_size)
        cats = self.cat_embed(cats)
        obj_embed = torch.cat([cats, bboxs], dim=-1)
        # (batch_size, num_bboxs, wrd_hidden_size)
        obj_hidden = self.bbox_mlp(obj_embed)
        # inner product: (batch_size, num_bboxs)
        logits = torch.bmm(obj_hidden, wrd_hidden.unsqueeze(-1)).squeeze(-1)
        if bboxs_mask is not None:
            # (batch_size, num_bboxs)
            # logits = bboxs_mask * logits
            logits[~bboxs_mask] = -1e10
        return logits
