#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
from typing import Dict, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from pytext.data.squad_tensorizer import SquadTensorizer
from pytext.data.tensorizers import LabelTensorizer, Tensorizer
from pytext.data.utils import Vocabulary
from pytext.models.decoders.mlp_decoder import MLPDecoder
from pytext.models.embeddings import WordEmbedding
from pytext.models.model import BaseModel
from pytext.models.module import create_module
from pytext.models.output_layers.squad_output_layer import SquadOutputLayer
from pytext.models.representations.attention import (
    DotProductSelfAttention,
    MultiplicativeAttention,
    SequenceAlignedAttention,
)
from pytext.models.representations.pooling import SelfAttention
from pytext.models.representations.stacked_bidirectional_rnn import (
    StackedBidirectionalRNN,
)
from pytext.utils.usage import log_class_usage


GLOVE_840B_300D = "/mnt/vol/pytext/users/kushall/pretrained/glove.840B.300d.txt"


class DrQAModel(BaseModel):
    class Config(BaseModel.Config):
        class ModelInput(BaseModel.Config.ModelInput):
            squad_input: SquadTensorizer.Config = SquadTensorizer.Config()
            has_answer: LabelTensorizer.Config = LabelTensorizer.Config(
                column="has_answer"
            )

        # Model inputs.
        inputs: ModelInput = ModelInput()

        # Configrable modules for the model.
        dropout: float = 0.4  # Overrides dropout in sub-modules of the model.
        embedding: WordEmbedding.Config = WordEmbedding.Config(
            embed_dim=300,
            pretrained_embeddings_path=GLOVE_840B_300D,
            vocab_from_pretrained_embeddings=True,
        )
        ques_rnn: StackedBidirectionalRNN.Config = StackedBidirectionalRNN.Config(
            dropout=dropout
        )
        doc_rnn: StackedBidirectionalRNN.Config = StackedBidirectionalRNN.Config(
            dropout=dropout
        )

        # Output layer.
        output_layer: SquadOutputLayer.Config = SquadOutputLayer.Config()

        # Is model traning distilling knowledg from a teacher model?
        is_kd: bool = False

    @classmethod
    def from_config(cls, config: Config, tensorizers: Dict[str, Tensorizer]):
        # Although the RNN params are configurable, for DrQA we want to set
        # the following parameters for all cases.
        config.ques_rnn.dropout = config.dropout
        config.doc_rnn.dropout = config.dropout

        embedding = cls.create_embedding(config, tensorizers)
        ques_aligned_doc_attn = SequenceAlignedAttention(embedding.embedding_dim)
        ques_rnn = create_module(config.ques_rnn, input_size=embedding.embedding_dim)
        doc_rnn = create_module(config.doc_rnn, input_size=embedding.embedding_dim * 2)
        ques_self_attn = DotProductSelfAttention(ques_rnn.representation_dim)
        start_attn = MultiplicativeAttention(
            doc_rnn.representation_dim, ques_rnn.representation_dim, normalize=False
        )
        end_attn = MultiplicativeAttention(
            doc_rnn.representation_dim, ques_rnn.representation_dim, normalize=False
        )
        doc_rep_pool = SelfAttention(
            SelfAttention.Config(dropout=config.dropout),
            n_input=doc_rnn.representation_dim,
        )
        has_answer_labels = ["False", "True"]
        tensorizers["has_answer"].vocab = Vocabulary(has_answer_labels)
        has_ans_decoder = MLPDecoder(
            config=MLPDecoder.Config(),
            in_dim=doc_rnn.representation_dim,
            out_dim=len(has_answer_labels),
        )
        output_layer = create_module(
            config.output_layer, labels=has_answer_labels, is_kd=config.is_kd
        )
        return cls(
            dropout=nn.Dropout(config.dropout),
            embedding=embedding,
            ques_rnn=ques_rnn,
            doc_rnn=doc_rnn,
            ques_self_attn=ques_self_attn,
            ques_aligned_doc_attn=ques_aligned_doc_attn,
            start_attn=start_attn,
            end_attn=end_attn,
            doc_rep_pool=doc_rep_pool,
            has_ans_decoder=has_ans_decoder,
            output_layer=output_layer,
            is_kd=config.is_kd,
        )

    @classmethod
    def create_embedding(cls, model_config: Config, tensorizers: Dict[str, Tensorizer]):
        squad_tensorizer = tensorizers["squad_input"]

        # Initialize the embedding module.
        embedding_module = create_module(model_config.embedding, None, squad_tensorizer)

        # Set ques and doc tensorizer vocab to squad_tensorizer.vocab.
        squad_tensorizer.ques_tensorizer.vocab = squad_tensorizer.vocab
        squad_tensorizer.doc_tensorizer.vocab = squad_tensorizer.vocab

        return embedding_module

    def __init__(
        self,
        dropout: nn.Module,
        embedding: nn.Module,
        ques_rnn: nn.Module,
        doc_rnn: nn.Module,
        ques_self_attn: nn.Module,
        ques_aligned_doc_attn: nn.Module,
        start_attn: nn.Module,
        end_attn: nn.Module,
        doc_rep_pool: nn.Module,
        has_ans_decoder: nn.Module,
        output_layer: nn.Module,
        is_kd: bool = Config.is_kd,
    ) -> None:
        super().__init__()
        self.dropout = dropout
        self.embedding = embedding
        self.ques_rnn = ques_rnn
        self.doc_rnn = doc_rnn
        self.ques_self_attn = ques_self_attn
        self.ques_aligned_doc_attn = ques_aligned_doc_attn
        self.start_attn = start_attn
        self.end_attn = end_attn
        self.doc_rep_pool = doc_rep_pool
        self.has_ans_decoder = has_ans_decoder
        self.output_layer = output_layer
        self.ignore_impossible = output_layer.ignore_impossible
        self.module_list = [
            embedding,
            ques_rnn,
            doc_rnn,
            ques_self_attn,
            ques_aligned_doc_attn,
            start_attn,
            end_attn,
            has_ans_decoder,
        ]
        self.is_kd = is_kd
        log_class_usage(__class__)

    def arrange_model_inputs(self, tensor_dict):
        (
            doc_tokens,
            doc_seq_len,
            doc_mask,
            ques_tokens,
            ques_seq_len,
            ques_mask,
            *ignore,
        ) = tensor_dict["squad_input"]
        return (doc_tokens, doc_seq_len, doc_mask, ques_tokens, ques_seq_len, ques_mask)

    def arrange_targets(self, tensor_dict):
        has_answer = tensor_dict["has_answer"]
        if not self.is_kd:
            _, _, _, _, _, _, answer_start_idx, answer_end_idx = tensor_dict[
                "squad_input"
            ]
            return answer_start_idx, answer_end_idx, has_answer
        else:
            (
                _,
                _,
                _,
                _,
                _,
                _,
                answer_start_idx,
                answer_end_idx,
                start_logits,
                end_logits,
                has_answer_logits,
            ) = tensor_dict["squad_input"]
            return (
                answer_start_idx,
                answer_end_idx,
                has_answer,
                start_logits,
                end_logits,
                has_answer_logits,
            )

    def forward(
        self,
        doc_tokens: torch.Tensor,
        doc_seq_len: torch.Tensor,
        doc_mask: torch.Tensor,
        ques_tokens: torch.Tensor,
        ques_seq_len: torch.Tensor,
        ques_mask: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        # Embedding lookups.
        doc_embedded = self.dropout(self.embedding(doc_tokens))
        ques_embedded = self.dropout(self.embedding(ques_tokens))

        # Document encoding.
        doc_attn_weights = self.ques_aligned_doc_attn(
            doc_embedded, ques_embedded, ques_mask
        )
        ques_aligned_doc_embedded = doc_attn_weights.bmm(ques_embedded)
        doc_embedded = torch.cat([doc_embedded, ques_aligned_doc_embedded], dim=2)
        doc_seq_vec = self.doc_rnn(doc_embedded, doc_mask)

        # Question encoding.
        ques_seq_vec = self.ques_rnn(ques_embedded, ques_mask)
        ques_attn_weights = self.ques_self_attn(ques_seq_vec, ques_mask)
        ques_vec = ques_attn_weights.unsqueeze(1).bmm(ques_seq_vec).squeeze(1)

        # Apply dropout to the ques and document representations.
        doc_seq_vec = self.dropout(doc_seq_vec)
        ques_vec = self.dropout(ques_vec)

        # Compute bilinear attention weights for each document token.
        # These attention weights serve as logits for detecting start and end.
        start_logits = self.start_attn(doc_seq_vec, ques_vec, doc_mask)
        end_logits = self.end_attn(doc_seq_vec, ques_vec, doc_mask)

        has_answer_logits = torch.zeros(
            ques_tokens.size(0),  # batch_size
            self.has_ans_decoder.out_dim,  # Number of classes
        )
        if not self.ignore_impossible:  # Compute whether document has an answer.
            doc_vec = self.doc_rep_pool(doc_seq_vec, doc_seq_len)
            has_answer_logits = F.relu(self.has_ans_decoder(doc_vec))

        # start_logits and end_logits: batch_size, max_seq_len
        # has_answer_logit: batch_size, 2
        return start_logits, end_logits, has_answer_logits
