# 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 torch.nn.utils.rnn import pad_sequence


class QGenModel(nn.Module):
    def __init__(
        self,
        num_wrds,
        wrd_embed_size,
        wrd_pad_id,
        img_feat_size,
        img_hidden_size,
        lstm_hidden_size,
        **kwargs
    ):
        super(QGenModel, self).__init__()
        self.lstm_hidden_size = lstm_hidden_size
        self.img_mlp = nn.Linear(
            img_feat_size, img_hidden_size)
        self.wrd_embed = nn.Embedding(
            num_wrds, wrd_embed_size, padding_idx=wrd_pad_id)
        self.lstm = nn.LSTM(
            wrd_embed_size + img_hidden_size, 
            lstm_hidden_size, batch_first=True, bidirectional=False)
        self.proj = nn.Linear(lstm_hidden_size, num_wrds)
        self.softmax = nn.Softmax(dim=-1)


    def forward(self, wrds, wrd_len, img_feat, mask=None, return_last_state=False):
        batch_size = wrds.size(0)
        img_hidden = self.img_mlp(img_feat)
        img_hidden = img_hidden.unsqueeze(1).repeat(1, wrds.size(1), 1)
        wrd_embed = self.wrd_embed(wrds)
        lstm_input = torch.cat([wrd_embed, img_hidden], dim=-1)
        lstm_input = nn.utils.rnn.pack_padded_sequence(
                lstm_input, wrd_len, batch_first=True, enforce_sorted=False)
        lstm_hidden, last_state = self.lstm(lstm_input)
        # (batch_size, seq_len, lstm_hidden_size)
        lstm_hidden, _ = nn.utils.rnn.pad_packed_sequence(
            lstm_hidden, batch_first=True)
        # (batch_size, seq_len, num_wrds)
        logits = self.proj(lstm_hidden)
        if mask is not None:
            # (batch_size, seq_len, num_wrds)
            logits = mask.unsqueeze(-1).repeat(1,1,logits.size(-1)) * logits
        if return_last_state:
            return logits, last_state
        else:
            return logits

    def step(self, wrd, img_feat, state):
        batch_size = wrd.size(0)
        img_hidden = self.img_mlp(img_feat)
        wrd_embed = self.wrd_embed(wrd)
        # (batch_size, 1, lstm_hidden_size)
        lstm_input = torch.cat([wrd_embed, img_hidden], dim=-1).unsqueeze(1)
        lstm_hidden, state = self.lstm(lstm_input, (state[0], state[1]))        
        lstm_hidden = lstm_hidden.view(batch_size, -1)
        # (batch_size, num_wrds)
        logit = self.proj(lstm_hidden)
        # (2, ...)
        state = torch.stack(state)
        return logit, state

    def generate_from_dialog(self, dialog, dialog_len, last_ans, img_feat, eoq_token, eod_token, max_q_len):
        _, last_state = self.forward(dialog, dialog_len, img_feat, return_last_state=True)
        q_token, _, _, _ = self.generate(
            last_ans, 
            img_feat, 
            eoq_token, 
            eod_token, 
            torch.zeros(dialog.size(0)).to(dialog.device).bool(), 
            max_q_len, 
            last_state=torch.stack(last_state)
            )
        return q_token

    def generate(self, last_wrd, img_feat, eoq_token, eod_token, end_of_dialog, 
                 max_q_len, last_state=None, greedy=True):
        last_wrd = last_wrd.clone()
        batch_size = last_wrd.size(0)
        device = last_wrd.device
        # Have generated <eoq> or <eod> 
        # Shape: (batch_size)
        finished = end_of_dialog.clone()
        # Shape: (batch_size)
        actual_length = torch.zeros_like(last_wrd).long()
        if last_state is None:
            # (h, c)
            last_state = torch.zeros(2, 1, batch_size, self.lstm_hidden_size).to(device)
                          
        _q_tokens = []
        for t in range(max_q_len):
            logit, state = self.step(last_wrd, img_feat, last_state)
            if greedy:
                q_t = logit.argmax(dim=-1)
            else:
                q_t = torch.multinomial(self.softmax(logit), 1).view(-1)
            
            updated_indices = torch.logical_not(finished).nonzero().view(-1)
            actual_length[updated_indices] = actual_length[updated_indices] + 1
            last_state[:, :, updated_indices] = state[:, :, updated_indices]
            last_wrd[updated_indices] = q_t[updated_indices]
            # _q_tokens.append(last_wrd) # Bad because this "last_wrd" will be altered in-place
            _q_tokens.append(q_t)
            # Update finished flags
            new_end_of_question = (q_t == eoq_token)
            new_end_of_dialog = (q_t == eod_token)
            new_finished_indices = (new_end_of_question | new_end_of_dialog).nonzero().view(-1)
            finished[new_finished_indices] = 1
            end_of_dialog = end_of_dialog | new_end_of_dialog
            if finished.sum() == batch_size:
                # all finished
                break

        _q_tokens = torch.stack(_q_tokens).transpose(0, 1)
        q_tokens = []
        for _q_tok, act_len in zip(_q_tokens, actual_length):
            q_tokens.append(_q_tok[:act_len])

        return q_tokens, last_state, actual_length, end_of_dialog
