from typing import Tuple

import torch
from torch import nn as nn, Tensor
from torch.nn.functional import softmax

from .. import utils


class Encoder(nn.Module):
    """Encoder module.

    The encoder represents its inputs using fixed, pre-trained GloVe embeddings, which are passed through a
    bidirectional GRU.

    Attributes
    ----------
    units: int
        Number of hidden units in the GRU.
    embed_dim: int
        Embedding dimensionality.
    embedding: Module
        Embedding module.
    gru: Module
        Bidirectional GRU module.
    fc: Module
        Fully-connected layer. Used to perform dimensionality reduction on final (bidirectional) hidden state, which
        will be used to initialize the (unidirectional) hidden state of the decoder.
    dropout: Module
        Dropout module for the fully-connected layer.
    """

    def __init__(self, embed_weights: Tensor, units: int, dropout: float) -> None:
        """Initialize the encoder.

        Constructs encoder sub-modules and initializes their weights using `utils.init_weights`.

        :param embed_weights: Tensor
            Embedding weights of shape (src_vocab_size, embed_dim).
        :param units: int
            Number of GRU hidden units.
        :param dropout: float
            Dropout probability.
        """
        super().__init__()
        self.units = units
        self.embed_dim = embed_weights.shape[1]

        self.embedding = nn.Embedding.from_pretrained(embed_weights)
        self.gru = nn.GRU(self.embed_dim, units, bidirectional=True)

        self.fc = nn.Linear(2 * units, units)
        self.dropout = nn.Dropout(dropout)

        utils.init_weights_(self.gru)
        utils.init_weights_(self.fc)

    def forward(self, source: Tensor) -> Tuple[Tensor, Tensor]:
        """Encoder a source sequence by forward propagation.

        :param source: Tensor
            Source sequence of shape (seq_len, batch_size)
        :return: Tuple[Tensor, Tensor]
            Outputs of shape (seq_len, batch_size, 2 * units) and state of shape (batch_size, units)
        """
        embedded = self.embedding(source)
        embedded = self.dropout(embedded)

        outputs, state = self.gru(embedded)
        outputs = self.dropout(outputs)

        # combine and reshape bidirectional states for compatibility with (unidirectional) decoder
        combined = torch.cat((state[-2, :, :], state[-1, :, :]), dim=1)
        state = torch.tanh(self.fc(combined))

        return outputs, state


class Attention(nn.Module):
    """Attention module.

    Used to compute Bahdanau attention weights over encoder outputs.

    Attributes
    ----------
    attn_in: int
        Dimensionality of attention input.
    fc: Module
        First fully-connected layer
    fc2: Module
        Second fully-connected layer.
    dropout: Module
        Dropout module
    """
    def __init__(self, enc_units: int, dec_units: int, units: int) -> None:
        """Initializes the attention module.

        :param enc_units: int
            Encoder output dimensionality.
        :param dec_units: int
            Decoder state dimensionality.
        :param units: int
            Attention dimensionality.
        """
        super().__init__()
        self.attn_in = 2 * enc_units + dec_units

        self.fc = nn.Linear(self.attn_in, units)
        self.fc2 = nn.Linear(units, 1)

        utils.init_weights_(self.fc)
        utils.init_weights_(self.fc2)

    def forward(self, dec_state: Tensor, enc_outputs: Tensor) -> Tensor:
        """Computes (Bahdanau) attention weights.

        :param dec_state: Tensor
            Previous or initial decoder state of shape (batch_size, dec_units).
        :param enc_outputs: Tensor
            Encoder outputs of shape (source_len, batch_size, 2 * enc_units).
        :return:
            Attention weights of shape (batch_size, source_len).
        """
        source_len = enc_outputs.shape[0]

        # vectorize computation of Bahdanau attention scores for all encoder outputs
        dec_state = dec_state.unsqueeze(1).repeat(1, source_len, 1)
        enc_outputs = enc_outputs.permute(1, 0, 2)

        combined = torch.cat((dec_state, enc_outputs), dim=2)
        scores = torch.tanh(self.fc(combined))                                          # (len, batch, units)
        scores = self.fc2(scores).squeeze(-1)

        weights = softmax(scores, dim=1)
        return weights


class Decoder(nn.Module):
    """Decoder module.

    Predicts next token probabilities using previous token and state and context vector.

    Attributes
    ----------
    vocab_size: int
        Target vocabulary size.
    embed_dim: int
        Embedding dimension.
    embedding: Module
        Embedding layer.
    rnn: Module
        GRU.
    fc: Module
        Fully-connected layer that outputs scores over target vocabulary.
    dropout: Module
        Dropout module.
    """
    def __init__(self, embed_weights: Tensor, units: int, dropout: float, enc_units: int) -> None:
        """Initializes the decoder module.

        :param embed_weights: Tensor
            Embedding weights of shape (trg_vocab_size, embed_dim).
        :param units: int
            Number of hidden units in the GRU.
        :param dropout: float
            Dropout probability.
        """
        super().__init__()
        self.vocab_size, embed_dim = embed_weights.shape

        self.embedding = nn.Embedding.from_pretrained(embed_weights)
        self.rnn = nn.GRU((units * 2) + embed_dim, units)
        self.fc = nn.Linear(units + 2 * enc_units + embed_dim, self.vocab_size)
        self.dropout = nn.Dropout(dropout)
        self.ln = nn.LayerNorm(units)

        utils.init_weights_(self.rnn)
        utils.init_weights_(self.fc)

    def forward(self, token: Tensor, dec_state: Tensor, context: Tensor) -> Tuple[Tensor, Tensor]:
        """Decodes the next token.

        :param token: Tensor
            Token of shape (batch_size,).
        :param dec_state:
            Decoder state of shape (batch_size, dec_units).
        :param context:
            Context vector of shape (1, batch_size, 2 * enc_units).
        :return: Tuple[Tensor, Tensor]
            Vocabulary scores of shape (batch_size, vocab_size) and decoder state of shape (batch_size, dec_units).
        """
        token = token.unsqueeze(0)
        embedded = self.embedding(token)
        embedded = self.dropout(embedded)

        rnn_input = torch.cat((embedded, context), dim=2)
        output, dec_state = self.rnn(rnn_input, dec_state.unsqueeze(0))

        embedded = embedded.squeeze(0)
        output = output.squeeze(0)
        context = context.squeeze(0)

        # dropout already applied to embedded
        output = self.dropout(output)
        context = self.dropout(context)

        # compute output using gru output, context vector, and embedding of previous output
        combined = torch.cat((output, context, embedded), dim=1)
        output = self.fc(combined)
        output = self.dropout(output)

        # return logits (rather than softmax activations) for compatibility with cross-entropy loss
        dec_state = dec_state.squeeze(0)
        dec_state = self.ln(dec_state)
        return output, dec_state
