import torch.nn as nn
import torch
from relogic.logickit.base.utils import log
from typing import Tuple
from relogic.logickit.modules.input_variational_dropout import InputVariationalDropout
from relogic.logickit.modules.bilinear_matrix_attention import BilinearMatrixAttention
import copy
import numpy
from relogic.logickit.utils.utils import get_range_vector, get_device_of, masked_log_softmax
from relogic.logickit.modules.chi_liu_edmonds import decode_mst
import torch.nn.functional as F

class BiaffineDepModule(nn.Module):
  def __init__(self, config, task_name, n_classes):
    super(BiaffineDepModule, self).__init__()
    self.config = config
    self.task_name = task_name
    self.n_classes = n_classes

    if hasattr(self.config, "sequence_labeling_use_cls") and self.config.sequence_labeling_use_cls:
      self.mul = 2
      log("Use CLS in dependency parsing")
    else:
      self.mul = 1

    encoder_dim = config.hidden_size
    arc_representation_dim = tag_representation_dim = config.dep_parsing_mlp_dim
    # self.pos_tag_embedding = nn.Embedding()
    self.head_sentinel = torch.nn.Parameter(torch.randn([1, 1, config.hidden_size]))

    # TODO: Need to check the dropout attribute.
    # TODO: How to design task specific parameter configuration
    self.dropout = InputVariationalDropout(config.dropout)
    self.head_arc_feedforward = nn.Sequential(
      nn.Linear(encoder_dim, arc_representation_dim),
      nn.ELU())
    self.child_arc_feedforward = copy.deepcopy(self.head_arc_feedforward)
    self.head_tag_feedforward = nn.Sequential(
      nn.Linear(encoder_dim, tag_representation_dim),
      nn.ELU())
    self.child_tag_feedforward = copy.deepcopy(self.head_tag_feedforward)

    self.arc_attention = BilinearMatrixAttention(
      matrix_1_dim=arc_representation_dim,
      matrix_2_dim=arc_representation_dim,
      use_input_biases=True)

    self.tag_bilinear = nn.modules.Bilinear(
        tag_representation_dim, tag_representation_dim, self.n_classes)




  def forward(self, *input, **kwargs):
    features = kwargs.pop("features")
    mask = (kwargs.pop("input_head") == 1).long()

    head_indices = kwargs.pop("arcs_ids", None)
    head_tags = kwargs.pop("label_ids", None)


    batch_size = features.size(0)
    encoding_dim = features.size(2)
    head_sentinel = self.head_sentinel.expand(batch_size, 1, encoding_dim)
    encoded_text = torch.cat([head_sentinel, features], dim=1)
    mask = torch.cat([mask.new_ones(batch_size, 1), mask], dim=1)
    float_mask = mask.float()

    if head_indices is not None:
      head_indices = torch.cat([head_indices.new_zeros(batch_size, 1), head_indices], 1)
    if head_tags is not None:
      head_tags = torch.cat([head_tags.new_zeros(batch_size, 1), head_tags], 1)

    encoded_text = self.dropout(encoded_text)

    head_arc_representation = self.dropout(self.head_arc_feedforward(encoded_text))
    child_arc_representation = self.dropout(self.child_arc_feedforward(encoded_text))

    head_tag_representation = self.dropout(self.head_tag_feedforward(encoded_text))
    child_tag_representation = self.dropout(self.child_tag_feedforward(encoded_text))

    attended_arcs = self.arc_attention(head_arc_representation, child_arc_representation)
    # shape (batch_size, sequence_length, sequence_length)

    minus_inf = -1e8
    minus_mask = (1 - float_mask) * minus_inf
    attended_arcs = attended_arcs + minus_mask.unsqueeze(2) + minus_mask.unsqueeze(1)

    if self.training:
      predicted_heads, predicted_head_tags = self.greedy_decode(
        head_tag_representation, child_tag_representation, attended_arcs, mask)

    else:
      predicted_heads, predicted_head_tags = self.mst_decode(
        head_tag_representation, child_tag_representation, attended_arcs, mask
      )

    if head_indices is not None and head_tags is not None:
      arc_nll, tag_nll = self.construct_loss(
        head_tag_representation=head_tag_representation,
        child_tag_representation=child_tag_representation,
        attended_arcs=attended_arcs,
        head_indices=head_indices,
        head_tags=head_tags,
        mask=mask,
      )
    else:
      arc_nll, tag_nll = self.construct_loss(
        head_tag_representation=head_tag_representation,
        child_tag_representation=child_tag_representation,
        attended_arcs=attended_arcs,
        head_indices=predicted_heads.long(),
        head_tags=predicted_head_tags.long(),
        mask=mask,
      )

    loss = arc_nll + tag_nll

    # if head_indices is not None and head_tags is not None:
    #   evaluation_mask = self._get_mask_for_eval(mask[:, 1:], pos_tags)
    #   # We calculate attatchment scores for the whole sentence
    #   # but excluding the symbolic ROOT token at the start,
    #   # which is why we start from the second element in the sequence.
    #   self._attachment_scores(
    #     predicted_heads[:, 1:],
    #     predicted_head_tags[:, 1:],
    #     head_indices,
    #     head_tags,
    #     evaluation_mask,
    #   )

    output_dict = {
      "heads": predicted_heads,
      "head_tags": predicted_head_tags,
      "arc_loss": arc_nll,
      "tag_loss": tag_nll,
      "loss": loss,
      "mask": mask}


    return output_dict



  def construct_loss(
          self,
          head_tag_representation: torch.Tensor,
          child_tag_representation: torch.Tensor,
          attended_arcs: torch.Tensor,
          head_indices: torch.Tensor,
          head_tags: torch.Tensor,
          mask: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Computes the arc and tag loss for a sequence given gold head indices and tags.
    Parameters
    ----------
    head_tag_representation : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length, tag_representation_dim),
        which will be used to generate predictions for the dependency tags
        for the given arcs.
    child_tag_representation : ``torch.Tensor``, required
        A tensor of shape (batch_size, sequence_length, tag_representation_dim),
        which will be used to generate predictions for the dependency tags
        for the given arcs.
    attended_arcs : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length, sequence_length) used to generate
        a distribution over attachments of a given word to all other words.
    head_indices : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length).
        The indices of the heads for every word.
    head_tags : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length).
        The dependency labels of the heads for every word.
    mask : ``torch.Tensor``, required.
        A mask of shape (batch_size, sequence_length), denoting unpadded
        elements in the sequence.
    Returns
    -------
    arc_nll : ``torch.Tensor``, required.
        The negative log likelihood from the arc loss.
    tag_nll : ``torch.Tensor``, required.
        The negative log likelihood from the arc tag loss.
    """
    float_mask = mask.float()
    batch_size, sequence_length, _ = attended_arcs.size()
    # shape (batch_size, 1)
    range_vector = get_range_vector(batch_size, get_device_of(attended_arcs)).unsqueeze(1)
    # shape (batch_size, sequence_length, sequence_length)
    normalised_arc_logits = (
          masked_log_softmax(attended_arcs, mask)
          * float_mask.unsqueeze(2)
          * float_mask.unsqueeze(1)
    )

    # shape (batch_size, sequence_length, num_head_tags)
    head_tag_logits = self.get_head_tags(
      head_tag_representation, child_tag_representation, head_indices
    )
    normalised_head_tag_logits = masked_log_softmax(
      head_tag_logits, mask.unsqueeze(-1)
    ) * float_mask.unsqueeze(-1)
    # index matrix with shape (batch, sequence_length)
    timestep_index = get_range_vector(sequence_length, get_device_of(attended_arcs))
    child_index = (
      timestep_index.view(1, sequence_length).expand(batch_size, sequence_length).long()
    )
    # shape (batch_size, sequence_length)
    arc_loss = normalised_arc_logits[range_vector, child_index, head_indices]
    tag_loss = normalised_head_tag_logits[range_vector, child_index, head_tags]
    # We don't care about predictions for the symbolic ROOT token's head,
    # so we remove it from the loss.
    arc_loss = arc_loss[:, 1:]
    tag_loss = tag_loss[:, 1:]

    # The number of valid positions is equal to the number of unmasked elements minus
    # 1 per sequence in the batch, to account for the symbolic HEAD token.
    valid_positions = mask.sum() - batch_size

    arc_nll = -arc_loss.sum() / valid_positions.float()
    tag_nll = -tag_loss.sum() / valid_positions.float()
    return arc_nll, tag_nll

  def greedy_decode(
        self,
        head_tag_representation: torch.Tensor,
        child_tag_representation: torch.Tensor,
        attended_arcs: torch.Tensor,
        mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Decodes the head and head tag predictions by decoding the unlabeled arcs
    independently for each word and then again, predicting the head tags of
    these greedily chosen arcs independently. Note that this method of decoding
    is not guaranteed to produce trees (i.e. there maybe be multiple roots,
    or cycles when children are attached to their parents).
    Parameters
    ----------
    head_tag_representation : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length, tag_representation_dim),
        which will be used to generate predictions for the dependency tags
        for the given arcs.
    child_tag_representation : ``torch.Tensor``, required
        A tensor of shape (batch_size, sequence_length, tag_representation_dim),
        which will be used to generate predictions for the dependency tags
        for the given arcs.
    attended_arcs : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length, sequence_length) used to generate
        a distribution over attachments of a given word to all other words.
    Returns
    -------
    heads : ``torch.Tensor``
        A tensor of shape (batch_size, sequence_length) representing the
        greedily decoded heads of each word.
    head_tags : ``torch.Tensor``
        A tensor of shape (batch_size, sequence_length) representing the
        dependency tags of the greedily decoded heads of each word.
    """
    attended_arcs = attended_arcs + torch.diag(
      attended_arcs.new(mask.size(1)).fill_(-numpy.inf))
    # Compute the heads greedily.
    # shape (batch_size, sequence_length)
    _, heads = attended_arcs.max(dim=2)
    # Given the greedily predicted heads, decode their dependency tags.
    # shape (batch_size, sequence_length, num_head_tags)
    head_tag_logits = self.get_head_tags(
      head_tag_representation, child_tag_representation, heads
    )
    _, head_tags = head_tag_logits.max(dim=2)
    return heads, head_tags

  def get_head_tags(
        self,
        head_tag_representation: torch.Tensor,
        child_tag_representation: torch.Tensor,
        head_indices: torch.Tensor) -> torch.Tensor:
    """
    Decodes the head tags given the head and child tag representations
    and a tensor of head indices to compute tags for. Note that these are
    either gold or predicted heads, depending on whether this function is
    being called to compute the loss, or if it's being called during inference.
    Parameters
    ----------
    head_tag_representation : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length, tag_representation_dim),
        which will be used to generate predictions for the dependency tags
        for the given arcs.
    child_tag_representation : ``torch.Tensor``, required
        A tensor of shape (batch_size, sequence_length, tag_representation_dim),
        which will be used to generate predictions for the dependency tags
        for the given arcs.
    head_indices : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length). The indices of the heads
        for every word.
    Returns
    -------
    head_tag_logits : ``torch.Tensor``
        A tensor of shape (batch_size, sequence_length, num_head_tags),
        representing logits for predicting a distribution over tags
        for each arc.
    """
    batch_size = head_tag_representation.size(0)
    # shape (batch_size,)
    range_vector = get_range_vector(
      batch_size, get_device_of(head_tag_representation)
    ).unsqueeze(1)

    # This next statement is quite a complex piece of indexing, which you really
    # need to read the docs to understand. See here:
    # https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.indexing.html#advanced-indexing
    # In effect, we are selecting the indices corresponding to the heads of each word from the
    # sequence length dimension for each element in the batch.

    # shape (batch_size, sequence_length, tag_representation_dim)
    selected_head_tag_representations = head_tag_representation[range_vector, head_indices]
    selected_head_tag_representations = selected_head_tag_representations.contiguous()
    # shape (batch_size, sequence_length, num_head_tags)
    head_tag_logits = self.tag_bilinear(
      selected_head_tag_representations, child_tag_representation
    )
    return head_tag_logits

  def mst_decode(
        self,
        head_tag_representation: torch.Tensor,
        child_tag_representation: torch.Tensor,
        attended_arcs: torch.Tensor,
        mask: torch.Tensor,
  ) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Decodes the head and head tag predictions using the Edmonds' Algorithm
    for finding minimum spanning trees on directed graphs. Nodes in the
    graph are the words in the sentence, and between each pair of nodes,
    there is an edge in each direction, where the weight of the edge corresponds
    to the most likely dependency label probability for that arc. The MST is
    then generated from this directed graph.
    Parameters
    ----------
    head_tag_representation : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length, tag_representation_dim),
        which will be used to generate predictions for the dependency tags
        for the given arcs.
    child_tag_representation : ``torch.Tensor``, required
        A tensor of shape (batch_size, sequence_length, tag_representation_dim),
        which will be used to generate predictions for the dependency tags
        for the given arcs.
    attended_arcs : ``torch.Tensor``, required.
        A tensor of shape (batch_size, sequence_length, sequence_length) used to generate
        a distribution over attachments of a given word to all other words.
    Returns
    -------
    heads : ``torch.Tensor``
        A tensor of shape (batch_size, sequence_length) representing the
        greedily decoded heads of each word.
    head_tags : ``torch.Tensor``
        A tensor of shape (batch_size, sequence_length) representing the
        dependency tags of the optimally decoded heads of each word.
    """
    batch_size, sequence_length, tag_representation_dim = head_tag_representation.size()

    lengths = mask.data.sum(dim=1).long().cpu().numpy()

    expanded_shape = [batch_size, sequence_length, sequence_length, tag_representation_dim]
    head_tag_representation = head_tag_representation.unsqueeze(2)
    head_tag_representation = head_tag_representation.expand(*expanded_shape).contiguous()
    child_tag_representation = child_tag_representation.unsqueeze(1)
    child_tag_representation = child_tag_representation.expand(*expanded_shape).contiguous()
    # Shape (batch_size, sequence_length, sequence_length, num_head_tags)
    pairwise_head_logits = self.tag_bilinear(head_tag_representation, child_tag_representation)

    # Note that this log_softmax is over the tag dimension, and we don't consider pairs
    # of tags which are invalid (e.g are a pair which includes a padded element) anyway below.
    # Shape (batch, num_labels,sequence_length, sequence_length)
    normalized_pairwise_head_logits = F.log_softmax(pairwise_head_logits, dim=3).permute(
      0, 3, 1, 2
    )

    # Mask padded tokens, because we only want to consider actual words as heads.
    minus_inf = -1e8
    minus_mask = (1 - mask.float()) * minus_inf
    attended_arcs = attended_arcs + minus_mask.unsqueeze(2) + minus_mask.unsqueeze(1)

    # Shape (batch_size, sequence_length, sequence_length)
    normalized_arc_logits = F.log_softmax(attended_arcs, dim=2).transpose(1, 2)

    # Shape (batch_size, num_head_tags, sequence_length, sequence_length)
    # This energy tensor expresses the following relation:
    # energy[i,j] = "Score that i is the head of j". In this
    # case, we have heads pointing to their children.
    batch_energy = torch.exp(
      normalized_arc_logits.unsqueeze(1) + normalized_pairwise_head_logits
    )
    return self._run_mst_decoding(batch_energy, lengths)

  @staticmethod
  def _run_mst_decoding(
        batch_energy: torch.Tensor, lengths: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor]:
    heads = []
    head_tags = []
    for energy, length in zip(batch_energy.detach().cpu(), lengths):
      scores, tag_ids = energy.max(dim=0)
      # Although we need to include the root node so that the MST includes it,
      # we do not want any word to be the parent of the root node.
      # Here, we enforce this by setting the scores for all word -> ROOT edges
      # edges to be 0.
      scores[0, :] = 0
      # Decode the heads. Because we modify the scores to prevent
      # adding in word -> ROOT edges, we need to find the labels ourselves.
      instance_heads, _ = decode_mst(scores.numpy(), length, has_labels=False)

      # Find the labels which correspond to the edges in the max spanning tree.
      instance_head_tags = []
      for child, parent in enumerate(instance_heads):
        instance_head_tags.append(tag_ids[parent, child].item())
      # We don't care what the head or tag is for the root token, but by default it's
      # not necesarily the same in the batched vs unbatched case, which is annoying.
      # Here we'll just set them to zero.
      instance_heads[0] = 0
      instance_head_tags[0] = 0
      heads.append(instance_heads)
      head_tags.append(instance_head_tags)
    return torch.from_numpy(numpy.stack(heads)), torch.from_numpy(numpy.stack(head_tags))

  def _get_mask_for_eval(
        self, mask: torch.LongTensor, pos_tags: torch.LongTensor
  ) -> torch.LongTensor:
    """
    Dependency evaluation excludes words are punctuation.
    Here, we create a new mask to exclude word indices which
    have a "punctuation-like" part of speech tag.
    Parameters
    ----------
    mask : ``torch.LongTensor``, required.
        The original mask.
    pos_tags : ``torch.LongTensor``, required.
        The pos tags for the sequence.
    Returns
    -------
    A new mask, where any indices equal to labels
    we should be ignoring are masked.
    """
    new_mask = mask.detach()
    for label in self._pos_to_ignore:
      label_mask = pos_tags.eq(label).long()
      new_mask = new_mask * (1 - label_mask)
    return new_mask







