# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import math
import torch
import torch.nn as nn
from torch.nn.utils.rnn import pad_sequence
# from src.model.guesser_vilbert import GuesserModel
from src.model.vilbert.vilbert import BertConfig, BertModel


class QGenModel(nn.Module):
    def __init__(
        self,
        num_wrds,
        wrd_pad_id,
        wrd_embed_size,
        obj_feat_size,
        lstm_hidden_size,
        num_bboxs,
        num_glimpses,
        state_handler_config,
        use_osda_glimpse=True,
        see_one_region_per_q=False,
        **kwargs
    ):
        super(QGenModel, self).__init__()
        self.state_handler = GuesserModel(**state_handler_config)

        self.lstm_hidden_size = lstm_hidden_size
        # Attention glimpse in the original paper
        self.num_glimpses = num_glimpses
        self.num_bboxs = num_bboxs
        self.see_one_region_per_q = see_one_region_per_q
        self.use_osda_glimpse = use_osda_glimpse

        self.img_mlp = nn.Sequential(
            nn.Linear(obj_feat_size, wrd_embed_size),
            nn.ReLU() # TODO: swish activation 
        )

        if see_one_region_per_q:
            self.onehot = nn.Embedding.from_pretrained(torch.eye(num_bboxs), freeze=True)
            self.post_mlp = nn.Sequential(
                nn.Linear(2*wrd_embed_size, wrd_embed_size),
                nn.ReLU()
            )
            vis_repr_out_dim = wrd_embed_size
        elif use_osda_glimpse:
            self.attn_glimpse = nn.Sequential(
                nn.Linear(num_bboxs * wrd_embed_size, num_glimpses),
                nn.Softmax(dim=1)
            )
            vis_repr_out_dim = num_glimpses * wrd_embed_size
        else:
            vis_repr_out_dim = wrd_embed_size

        self.wrd_embed = nn.Embedding(
            num_wrds, wrd_embed_size, padding_idx=wrd_pad_id)
        self.lstm = nn.LSTM(
            wrd_embed_size + vis_repr_out_dim, 
            lstm_hidden_size, batch_first=True, bidirectional=False)
        if wrd_embed_size != lstm_hidden_size:
            self.ans_proj = nn.Linear(wrd_embed_size, lstm_hidden_size)
        self.proj = nn.Linear(lstm_hidden_size, num_wrds)
        self.softmax = nn.Softmax(dim=-1)

    def self_diff_attention(self, obj_repr, pi):
        num_bboxs = obj_repr.size(1)
        obj_feat_dim = obj_repr.size(-1)
        # Object-self Difference Attention (OsDA)
        # (batch_size, num_bboxs, num_bboxs, obj_feat_dim)
        diff = obj_repr.view(-1, num_bboxs, 1, obj_feat_dim) -\
                obj_repr.view(-1, 1, num_bboxs, obj_feat_dim)
        diff = obj_repr.unsqueeze(2) * diff
        diff = diff.view(-1, num_bboxs, num_bboxs * obj_feat_dim)

        if self.see_one_region_per_q:
            picked_idx = pi.argmax(dim=-1)
            # (1, 0, 0, ...) onehot
            pi_hard = pi + (self.onehot(picked_idx) - pi).detach()
            # obj_repr:      (batch_size, num_bboxs, obj_feat_dim)
            # obj_repr_hard: (batch_size, obj_feat_dim)
            obj_repr_hard = (obj_repr * pi_hard.unsqueeze(-1)).sum(dim=1)

            # diff:      (batch_size, num_bboxs, num_bboxs * obj_feat_dim)
            # diff_hard: (batch_size, num_bboxs * obj_feat_dim)
            diff_hard = (diff * pi_hard.unsqueeze(-1)).sum(dim=1).reshape(-1, num_bboxs, obj_feat_dim)
            vis_repr = diff_hard.sum(dim=1)
            vis_repr = self.post_mlp(torch.cat([obj_repr_hard, vis_repr], dim=-1))
            weight = pi_hard
        elif self.use_osda_glimpse:
            # (batch_size, num_glimpses, num_bboxs)
            weight = self.attn_glimpse(diff).transpose(1, 2)
            vis_repr = torch.bmm(weight, obj_repr)
            vis_repr = vis_repr.view(-1, self.num_glimpses * obj_feat_dim)
        else:
            vis_repr = obj_repr.sum(dim=1)
            # vis_repr = self.vis_proj(vis_repr)
            weight = pi
        return vis_repr, weight

    # Forward w/ teacher forcing 
    def forward_sentence(self, wrd_emb, wrd_seq_len, vis_repr, last_state):
        """
        wrd_emb: (batch_size, len of input question (padded), wrd_embed_size)
        wrd_seq_len: (batch_size)
        vis_repr: (batch_size, num_glimpses * wrd_embed_size)
        """
        # wrd_embed = self.wrd_embed(wrd_seq)
        lstm_input = torch.cat(
            [wrd_emb, vis_repr.unsqueeze(1).repeat(1, wrd_emb.size(1), 1)], 
            dim=-1)
        lstm_input = nn.utils.rnn.pack_padded_sequence(
                lstm_input, wrd_seq_len, batch_first=True, enforce_sorted=False)
        lstm_hidden, state = self.lstm(lstm_input, tuple(last_state))
        # (batch_size, seq_len, lstm_hidden_size)
        lstm_hidden, _ = nn.utils.rnn.pad_packed_sequence(
            lstm_hidden, batch_first=True, total_length=wrd_emb.size(1))
        state = torch.stack(state)
        return lstm_hidden, state

    # Forward w/ teacher forcing 
    def forward(
        self, qs, qs_tf_in, question_len, answers, img_feats, bboxs, txt_attn_mask, end_turn, pi=None, update_vilbert=False):
        # qs (an example in batch):       [[<sos>, is, it, a, dog, ?], ...]
        # q_len (an example in batch):    [5 (no <sos>), ...]
        # qs_tf_in (an example in batch): [[<sos>, is, it, a, dog], ...]
        batch_size = img_feats.size(0)
        num_bboxs = img_feats.size(1)
        device = img_feats.device
        max_num_turns = qs.size(1)
        assert num_bboxs == self.num_bboxs

        # qs_tf_in = qs[:, :-1]

        obj_feats = torch.cat([img_feats, bboxs], dim=-1)
        init_obj_repr = self.img_mlp(obj_feats)
        # obj_feat_dim = init_obj_repr.size(-1)
        q_emb = self.wrd_embed(qs_tf_in)
        # a_emb = self.wrd_embed(answers)

        # finished = torch.zeros(batch_size).bool().to(device)
        
        result_logits = []
        result_pi = []
        last_state = torch.zeros(2, 1, batch_size, self.lstm_hidden_size).to(device)
        if pi is None:
            pi = self.state_handler.init_state(batch_size, num_bboxs, device)
        final_guess_logits = torch.zeros_like(pi)

        for t in range(max_num_turns):
            # Update object representations
            # (batch_size, num_bboxs, obj_feat_dim)
            obj_repr = init_obj_repr * pi.unsqueeze(-1)
            vis_repr, weight = self.self_diff_attention(obj_repr, pi)
            # HACK: RNN can not forward w/ len == 0
            fake_len = question_len[:, t].clone()
            fake_len[fake_len==0] = 1
            # LSTM
            lstm_hidden, state = self.forward_sentence(
                q_emb[:, t],
                fake_len,
                vis_repr,
                last_state
            )
            # (batch_size, seq_len, num_wrds)
            logits = self.proj(lstm_hidden)
            # Not update finished states
            not_finished = question_len[:, t] != 0
            last_state = last_state.clone()
            last_state[:, :, not_finished] = state[:, :, not_finished]
            # last_state = state
            result_logits.append(logits)
            # update pi
            if t != max_num_turns-1:
                pi, guess_logits = self.state_handler.forward_turn(
                    qs[:, t], 
                    answers[:, t], 
                    None, # cats 
                    img_feats,
                    bboxs, 
                    curr_state=pi,
                    attention_mask=txt_attn_mask[:, t],
                    update_vilbert=update_vilbert,
                    )
                # pi = self.refresh_pi(pi, a_emb[:, t], last_state[0,0], obj_repr)
                result_pi.append(pi)
                end = end_turn == t
                final_guess_logits[end] = guess_logits[end]

        result_logits = torch.stack(result_logits).transpose(0, 1)
        result_pi = torch.stack(result_pi).transpose(0, 1)
        return result_logits, result_pi, final_guess_logits

    # w/o teacher forcing
    def generate_word(self, wrd, vis_repr, state):
        batch_size = wrd.size(0)
        wrd_embed = self.wrd_embed(wrd)
        # (batch_size, 1, (num_glimpses+1) * wrd_embed_size)
        lstm_input = torch.cat([wrd_embed, vis_repr], dim=-1).unsqueeze(1)
        lstm_hidden, state = self.lstm(lstm_input, tuple(state))        
        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

    # w/o teacher forcing
    def generate_sentence(
        self, last_wrd, img_feats, bboxs, eoq_token, eod_token, end_of_dialog, 
        max_q_len, pi=None, last_state=None, greedy=True):
        last_wrd = last_wrd.clone()
        batch_size = last_wrd.size(0)
        device = last_wrd.device
        num_bboxs = img_feats.size(1)
        # 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)
        if pi is None:
            # pi = (torch.ones(batch_size, self.num_bboxs) / self.num_bboxs).to(device)
            pi = self.state_handler.init_state(batch_size, num_bboxs, device)

        # UoDR + OsDA
        obj_feats = torch.cat([img_feats, bboxs], dim=-1)
        init_obj_repr = self.img_mlp(obj_feats)
        obj_repr = init_obj_repr * pi.unsqueeze(-1)
        vis_repr, weight = self.self_diff_attention(obj_repr, pi)
        
        _q_tokens = []

        for t in range(max_q_len):
            logit, state = self.generate_word(last_wrd, vis_repr, last_state)
            if greedy:
                q_t = logit.argmax(dim=-1)
            else:
                q_t = torch.multinomial(self.softmax(logit), 1).view(-1)
            
            # Only update those not finished
            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
            #print('end_of_dialog')
            #print(end_of_dialog)
            if finished.sum() == batch_size:
                # all finished
                break

        _q_tokens = torch.stack(_q_tokens).transpose(0, 1) # padded 1 question
        q_tokens = [] # not padded
        for _q_tok, act_len in zip(_q_tokens, actual_length):
            q_tokens.append(_q_tok[:act_len])
        return q_tokens, actual_length, last_state, end_of_dialog


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,
        ans_fuse_method='add',
        **kwargs
    ):
        super(GuesserModel, self).__init__()

        self.use_category = use_category
        self.ans_fuse_method = ans_fuse_method
        assert self.ans_fuse_method in ('add', 'mul'), "Only support these two options."
        # 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)
        if self.ans_fuse_method == 'add':
            fuesd_feats = fused_feats + self.ans_proj(ans_emb).unsqueeze(1)
        elif self.ans_fuse_method == 'mul':
            fuesd_feats = fused_feats * self.ans_proj(ans_emb).unsqueeze(1)
        else:
            raise NotImplementedError
        # 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():
                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:]
        
