# --------------------------------------------------------
# mcan-vqa (Deep Modular Co-Attention Networks)
# Licensed under The MIT License [see LICENSE for details]
# Written by Yuhao Cui https://github.com/cuiyuhao1996
# --------------------------------------------------------

from core.model.net_utils import FC, MLP, LayerNorm
from core.model.mca import MCA_ED

import numpy as np

import torch.nn as nn
import torch.nn.functional as F
import torch
from core.model.tv2d_layer import TV2DFunction
from entmax import sparsemax
from functools import partial


# ------------------------------
# ---- Flatten the sequence ----
# ------------------------------

class AttFlat(nn.Module):
    def __init__(self, __C, att_func=torch.softmax):
        super(AttFlat, self).__init__()
        self.__C = __C

        self.att_func=att_func

        if str(att_func)=='tvmax':
            self.sparsemax = partial(sparsemax, k=512)
            self.tvmax = TV2DFunction.apply

        self.mlp = MLP(
            in_size=__C.HIDDEN_SIZE,
            mid_size=__C.FLAT_MLP_SIZE,
            out_size=__C.FLAT_GLIMPSES,
            dropout_r=__C.DROPOUT_R,
            use_relu=True)

        self.linear_merge = nn.Linear(__C.HIDDEN_SIZE * __C.FLAT_GLIMPSES,__C.FLAT_OUT_SIZE)

    def forward(self, x, x_mask, size, eval_=False):
        att = self.mlp(x)
        att = att.masked_fill(x_mask.squeeze(1).squeeze(1).unsqueeze(2),-1e9)

        if str(self.att_func)=='tvmax':
            att = att.squeeze(-1)
            for i in range(att.size(0)):
                a = att[i][:size[0][i].item()*size[1][i].item()]
                a = a.view(size[0][i].item(),size[1][i].item())
                att[i][:size[0][i].item()*size[1][i].item()] = self.tvmax(a).view(-1)
            att = self.sparsemax(att.view(att.shape[0],-1), dim=-1).unsqueeze(-1)
        else:
            att = self.att_func(att.squeeze(-1), dim=-1).unsqueeze(-1)


        att_list = []
        for i in range(self.__C.FLAT_GLIMPSES):
            att_list.append(torch.sum(att[:, :, i: i + 1] * x, dim=1))

        x_atted = torch.cat(att_list, dim=1)
        x_atted = self.linear_merge(x_atted)
        return x_atted


# -------------------------
# ---- Main MCAN Model ----
# -------------------------

class Net(nn.Module):
    def __init__(self, __C, pretrained_emb, token_size, answer_size, att_func=torch.softmax):
        super(Net, self).__init__()

        self.embedding = nn.Embedding(num_embeddings=token_size, embedding_dim=__C.WORD_EMBED_SIZE)

        # Loading the GloVe embedding weights
        if __C.USE_GLOVE:
            self.embedding.weight.data.copy_(torch.from_numpy(pretrained_emb))

        self.lstm = nn.LSTM(
            input_size=__C.WORD_EMBED_SIZE,
            hidden_size=__C.HIDDEN_SIZE,
            num_layers=1,
            batch_first=True)

        self.img_feat_linear = nn.Linear(
            __C.IMG_FEAT_SIZE,
            __C.HIDDEN_SIZE)

        self.att_func=att_func
        self.backbone = MCA_ED(__C)

        self.attflat_img = AttFlat(__C, self.att_func)
        self.attflat_lang = AttFlat(__C)

        self.proj_norm = LayerNorm(__C.FLAT_OUT_SIZE)
        self.proj = nn.Linear(__C.FLAT_OUT_SIZE, answer_size)


    def forward(self, img_feat, ques_ix, size, eval_=False):

        # Make mask
        lang_feat_mask = self.make_mask(ques_ix.unsqueeze(2))
        img_feat_mask = self.make_mask(img_feat)
        # Pre-process Language Feature
        lang_feat = self.embedding(ques_ix)
        lang_feat, _ = self.lstm(lang_feat)

        # Pre-process Image Feature
        
        img_feat = self.img_feat_linear(img_feat)

        # Backbone Framework
        lang_feat, img_feat = self.backbone(
            lang_feat,
            img_feat,
            lang_feat_mask,
            img_feat_mask)
        
        lang_feat = self.attflat_lang(
            lang_feat,
            lang_feat_mask, size, eval_)

        img_feat = self.attflat_img(
            img_feat,
            img_feat_mask, size, eval_)

        proj_feat = lang_feat + img_feat
        proj_feat = self.proj_norm(proj_feat)
        proj_feat = torch.sigmoid(self.proj(proj_feat))

        return proj_feat


    # Masking
    def make_mask(self, feature):
        return (torch.sum(torch.abs(feature),dim=-1) == 0).unsqueeze(1).unsqueeze(2)
