import logging
from typing import Any, Dict, List, Optional

import torch
from torch.nn.functional import nll_loss

from allennlp.common.checks import check_dimensions_match
from allennlp.data import Vocabulary
from allennlp.models.model import Model
from allennlp.modules import Highway
from allennlp.modules import Seq2SeqEncoder, TimeDistributed, TextFieldEmbedder
from allennlp.modules.matrix_attention import MatrixAttention
from allennlp.nn import util, InitializerApplicator, RegularizerApplicator
from allennlp.training.metrics import BooleanAccuracy, CategoricalAccuracy

from vimarc.metrics import SquadEmAndF1
from vimarc.models.utils import (
    get_best_span,
    replace_masked_values_with_big_negative_number,
)

logger = logging.getLogger(__name__)


@Model.register("bidaf")
class BidirectionalAttentionFlow(Model):
    def __init__(
        self,
        vocab: Vocabulary,
        text_field_embedder: TextFieldEmbedder,
        num_highway_layers: int,
        phrase_layer: Seq2SeqEncoder,
        matrix_attention: MatrixAttention,
        modeling_layer: Seq2SeqEncoder,
        span_end_encoder: Seq2SeqEncoder,
        dropout: float = 0.2,
        mask_lstms: bool = True,
        initializer: InitializerApplicator = InitializerApplicator(),
        regularizer: Optional[RegularizerApplicator] = None,
    ) -> None:
        super().__init__(vocab, regularizer)

        self._text_field_embedder = text_field_embedder
        self._highway_layer = TimeDistributed(
            Highway(text_field_embedder.get_output_dim(), num_highway_layers)
        )
        self._phrase_layer = phrase_layer
        self._matrix_attention = matrix_attention
        self._modeling_layer = modeling_layer
        self._span_end_encoder = span_end_encoder

        encoding_dim = phrase_layer.get_output_dim()
        modeling_dim = modeling_layer.get_output_dim()
        span_start_input_dim = encoding_dim * 4 + modeling_dim
        self._span_start_predictor = TimeDistributed(torch.nn.Linear(span_start_input_dim, 1))

        span_end_encoding_dim = span_end_encoder.get_output_dim()
        span_end_input_dim = encoding_dim * 4 + span_end_encoding_dim
        self._span_end_predictor = TimeDistributed(torch.nn.Linear(span_end_input_dim, 1))

        # Bidaf has lots of layer dimensions which need to match up - these aren't necessarily
        # obvious from the configuration files, so we check here.
        check_dimensions_match(
            modeling_layer.get_input_dim(),
            4 * encoding_dim,
            "modeling layer input dim",
            "4 * encoding dim",
        )
        check_dimensions_match(
            text_field_embedder.get_output_dim(),
            phrase_layer.get_input_dim(),
            "text field embedder output dim",
            "phrase layer input dim",
        )
        check_dimensions_match(
            span_end_encoder.get_input_dim(),
            4 * encoding_dim + 3 * modeling_dim,
            "span end encoder input dim",
            "4 * encoding dim + 3 * modeling dim",
        )

        self._span_start_accuracy = CategoricalAccuracy()
        self._span_end_accuracy = CategoricalAccuracy()
        self._span_accuracy = BooleanAccuracy()
        self._squad_metrics = SquadEmAndF1()
        if dropout > 0:
            self._dropout = torch.nn.Dropout(p=dropout)
        else:
            self._dropout = lambda x: x
        self._mask_lstms = mask_lstms

        initializer(self)

    def forward(  # type: ignore
        self,
        question: Dict[str, torch.LongTensor],
        context: Dict[str, torch.LongTensor],
        span_start: torch.IntTensor = None,
        span_end: torch.IntTensor = None,
        metadata: List[Dict[str, Any]] = None,
    ) -> Dict[str, torch.Tensor]:

        """
        # Parameters

        question : `Dict[str, torch.LongTensor]`
            From a ``TextField``.
        context : `Dict[str, torch.LongTensor]`
            From a ``TextField``.  The model assumes that this context contains the answer to the
            question, and predicts the beginning and ending positions of the answer within the
            context.
        span_start : `torch.IntTensor`, optional
            From an ``IndexField``.  This is one of the things we are trying to predict - the
            beginning position of the answer with the context.  This is an `inclusive` token index.
            If this is given, we will compute a loss that gets included in the output dictionary.
        span_end : `torch.IntTensor`, optional
            From an ``IndexField``.  This is one of the things we are trying to predict - the
            ending position of the answer with the context.  This is an `inclusive` token index.
            If this is given, we will compute a loss that gets included in the output dictionary.
        metadata : `List[Dict[str, Any]]`, optional
            metadata : `List[Dict[str, Any]]`, optional
            If present, this should contain the question tokens, context tokens, original context
            text, and token offsets into the context for each instance in the batch.  The length
            of this list should be the batch size, and each dictionary should have the keys
            ``question_tokens``, ``context_tokens``, ``original_context``, and ``token_offsets``.

        Returns
        -------
        An output dictionary consisting of:

        span_start_logits : `torch.FloatTensor`
            A tensor of shape ``(batch_size, context_length)`` representing unnormalized log
            probabilities of the span start position.
        span_start_probs : `torch.FloatTensor`
            The result of ``softmax(span_start_logits)``.
        span_end_logits : `torch.FloatTensor`
            A tensor of shape ``(batch_size, context_length)`` representing unnormalized log
            probabilities of the span end position (inclusive).
        span_end_probs : `torch.FloatTensor`
            The result of ``softmax(span_end_logits)``.
        best_span : `torch.IntTensor`
            The result of a constrained inference over ``span_start_logits`` and
            ``span_end_logits`` to find the most probable span.  Shape is ``(batch_size, 2)``
            and each offset is a token index.
        loss : `torch.FloatTensor`, optional
            A scalar loss to be optimised.
        best_span_str : `List[str]`
            If sufficient metadata was provided for the instances in the batch, we also return the
            string from the original context that the model thinks is the best answer to the
            question.
        """
        embedded_question = self._highway_layer(self._text_field_embedder(question))
        embedded_context = self._highway_layer(self._text_field_embedder(context))
        batch_size = embedded_question.size(0)
        context_length = embedded_context.size(1)
        question_mask = util.get_text_field_mask(question)
        context_mask = util.get_text_field_mask(context)
        question_lstm_mask = question_mask if self._mask_lstms else None
        context_lstm_mask = context_mask if self._mask_lstms else None

        encoded_question = self._dropout(self._phrase_layer(embedded_question, question_lstm_mask))
        encoded_context = self._dropout(self._phrase_layer(embedded_context, context_lstm_mask))
        encoding_dim = encoded_question.size(-1)

        # Shape: (batch_size, context_length, question_length)
        context_question_similarity = self._matrix_attention(encoded_context, encoded_question)
        # Shape: (batch_size, context_length, question_length)
        context_question_attention = util.masked_softmax(context_question_similarity, question_mask)
        # Shape: (batch_size, context_length, encoding_dim)
        context_question_vectors = util.weighted_sum(encoded_question, context_question_attention)

        # We replace masked values with something really negative here, so they don't affect the
        # max below.
        masked_similarity = replace_masked_values_with_big_negative_number(
            context_question_similarity, question_mask.unsqueeze(1)
        )
        # Shape: (batch_size, context_length)
        question_context_similarity = masked_similarity.max(dim=-1)[0].squeeze(-1)
        # Shape: (batch_size, context_length)
        question_context_attention = util.masked_softmax(question_context_similarity, context_mask)
        # Shape: (batch_size, encoding_dim)
        question_context_vector = util.weighted_sum(encoded_context, question_context_attention)
        # Shape: (batch_size, context_length, encoding_dim)
        tiled_question_context_vector = question_context_vector.unsqueeze(1).expand(
            batch_size, context_length, encoding_dim
        )

        # Shape: (batch_size, context_length, encoding_dim * 4)
        final_merged_context = torch.cat(
            [
                encoded_context,
                context_question_vectors,
                encoded_context * context_question_vectors,
                encoded_context * tiled_question_context_vector,
            ],
            dim=-1,
        )

        modeled_context = self._dropout(
            self._modeling_layer(final_merged_context, context_lstm_mask)
        )
        modeling_dim = modeled_context.size(-1)

        # Shape: (batch_size, context_length, encoding_dim * 4 + modeling_dim))
        span_start_input = self._dropout(torch.cat([final_merged_context, modeled_context], dim=-1))
        # Shape: (batch_size, context_length)
        span_start_logits = self._span_start_predictor(span_start_input).squeeze(-1)
        # Shape: (batch_size, context_length)
        span_start_probs = util.masked_softmax(span_start_logits, context_mask)

        # Shape: (batch_size, modeling_dim)
        span_start_representation = util.weighted_sum(modeled_context, span_start_probs)
        # Shape: (batch_size, context_length, modeling_dim)
        tiled_start_representation = span_start_representation.unsqueeze(1).expand(
            batch_size, context_length, modeling_dim
        )

        # Shape: (batch_size, context_length, encoding_dim * 4 + modeling_dim * 3)
        span_end_representation = torch.cat(
            [
                final_merged_context,
                modeled_context,
                tiled_start_representation,
                modeled_context * tiled_start_representation,
            ],
            dim=-1,
        )
        # Shape: (batch_size, context_length, encoding_dim)
        encoded_span_end = self._dropout(
            self._span_end_encoder(span_end_representation, context_lstm_mask)
        )
        # Shape: (batch_size, context_length, encoding_dim * 4 + span_end_encoding_dim)
        span_end_input = self._dropout(torch.cat([final_merged_context, encoded_span_end], dim=-1))
        span_end_logits = self._span_end_predictor(span_end_input).squeeze(-1)
        span_end_probs = util.masked_softmax(span_end_logits, context_mask)

        # Replace the masked values with a very negative constant.
        span_start_logits = replace_masked_values_with_big_negative_number(
            span_start_logits, context_mask
        )
        span_end_logits = replace_masked_values_with_big_negative_number(
            span_end_logits, context_mask
        )
        best_span = get_best_span(span_start_logits, span_end_logits)

        output_dict = {
            "context_question_attention": context_question_attention,
            "span_start_logits": span_start_logits,
            "span_start_probs": span_start_probs,
            "span_end_logits": span_end_logits,
            "span_end_probs": span_end_probs,
            "best_span": best_span,
        }

        # Compute the loss for training.
        if span_start is not None:
            loss = nll_loss(
                util.masked_log_softmax(span_start_logits, context_mask), span_start.squeeze(-1)
            )
            self._span_start_accuracy(span_start_logits, span_start.squeeze(-1))
            loss += nll_loss(
                util.masked_log_softmax(span_end_logits, context_mask), span_end.squeeze(-1)
            )
            self._span_end_accuracy(span_end_logits, span_end.squeeze(-1))
            self._span_accuracy(best_span, torch.cat([span_start, span_end], -1))
            output_dict["loss"] = loss

        # Compute the EM and F1 on SQuAD and add the tokenized input to the output.
        if metadata is not None:
            output_dict["best_span_str"] = []
            question_tokens = []
            context_tokens = []
            token_offsets = []
            for i in range(batch_size):
                question_tokens.append(metadata[i]["question_tokens"])
                context_tokens.append(metadata[i]["context_tokens"])
                token_offsets.append(metadata[i]["token_offsets"])
                context_str = metadata[i]["original_context"]
                offsets = metadata[i]["token_offsets"]
                predicted_span = tuple(best_span[i].detach().cpu().numpy())
                start_offset = offsets[predicted_span[0]][0]
                end_offset = offsets[predicted_span[1]][1]
                best_span_string = context_str[start_offset:end_offset]
                output_dict["best_span_str"].append(best_span_string)
                answer_texts = metadata[i].get("answer_texts", [])
                if answer_texts:
                    self._squad_metrics(best_span_string, answer_texts)
            output_dict["question_tokens"] = question_tokens
            output_dict["context_tokens"] = context_tokens
            output_dict["token_offsets"] = token_offsets
        return output_dict

    def get_metrics(self, reset: bool = False) -> Dict[str, float]:
        exact_match, f1_score = self._squad_metrics.get_metric(reset)
        return {
            "start_acc": self._span_start_accuracy.get_metric(reset),
            "end_acc": self._span_end_accuracy.get_metric(reset),
            "span_acc": self._span_accuracy.get_metric(reset),
            "em": exact_match,
            "f1": f1_score,
        }

    @staticmethod
    def get_best_span(
        span_start_logits: torch.Tensor, span_end_logits: torch.Tensor
    ) -> torch.Tensor:
        # We call the inputs "logits" - they could either be unnormalized logits or normalized log
        # probabilities.  A log_softmax operation is a constant shifting of the entire logit
        # vector, so taking an argmax over either one gives the same result.
        if span_start_logits.dim() != 2 or span_end_logits.dim() != 2:
            raise ValueError("Input shapes must be (batch_size, context_length)")
        batch_size, context_length = span_start_logits.size()
        device = span_start_logits.device
        # (batch_size, context_length, context_length)
        span_log_probs = span_start_logits.unsqueeze(2) + span_end_logits.unsqueeze(1)
        # Only the upper triangle of the span matrix is valid; the lower triangle has entries where
        # the span ends before it starts.
        span_log_mask = (
            torch.triu(torch.ones((context_length, context_length), device=device))
            .log()
            .unsqueeze(0)
        )
        valid_span_log_probs = span_log_probs + span_log_mask

        # Here we take the span matrix and flatten it, then find the best span using argmax.  We
        # can recover the start and end indices from this flattened list using simple modular
        # arithmetic.
        # (batch_size, context_length * context_length)
        best_spans = valid_span_log_probs.view(batch_size, -1).argmax(-1)
        span_start_indices = best_spans // context_length
        span_end_indices = best_spans % context_length
        return torch.stack([span_start_indices, span_end_indices], dim=-1)

    default_predictor = "reading_comprehension"
