"""Transformer model classes."""

import torch
from torch import nn

from .. import data, defaults, special
from . import base, embeddings, modules


class TransformerModel(base.BaseModel):
    """Vanilla transformer model.

    If features are provided, the encodings are fused by concatenation of the
    features encoding with the source encoding on the sequence length
    dimension.

    This supports optional student forcing during training.

    Args:
        *args: passed to superclass.
        attention_heads (int, optional).
        teacher_forcing (bool, optional).
        **kwargs: passed to superclass.
    """

    # Model arguments.
    attention_heads: int
    teacher_forcing: bool
    classifier: nn.Linear

    def __init__(
        self,
        *args,
        attention_heads: int = defaults.ATTENTION_HEADS,
        teacher_forcing: bool = defaults.TEACHER_FORCING,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.attention_heads = attention_heads
        self.classifier = nn.Linear(
            self.embedding_size, self.target_vocab_size
        )
        if self.has_features_encoder and (
            self.source_encoder.output_size
            != self.features_encoder.output_size
        ):
            raise base.ConfigurationError(
                "Cannot concatenate source encoding "
                f"({self.source_encoder.output_size}) and features "
                f"encoding ({self.features_encoder.output_size})"
            )
        self.decoder = self.get_decoder()
        self.teacher_forcing = teacher_forcing
        self._log_model()
        self.save_hyperparameters(
            ignore=[
                "classifier",
                "decoder",
                "embeddings",
                "features_encoder",
                "source_encoder",
            ]
        )

    def decode_step(
        self,
        encoded: torch.Tensor,
        mask: torch.Tensor,
        predictions: torch.Tensor,
    ) -> torch.Tensor:
        """Single decoder step.

        This predicts a distribution for one symbol.

        Args:
            encoded (torch.Tensor).
            mask (torch.Tensor).
            predictions (torch.Tensor): tensor of predictions thus far.

        Returns:
            torch.Tensor: logits.
        """
        decoded, _ = self.decoder(
            encoded, mask, predictions, None, self.embeddings
        )
        # We only need the logits for the most recent time step.
        return self.classifier(decoded[:, -1, :])

    def forward(self, batch: data.Batch) -> torch.Tensor:
        """Forward pass.

        Args:
            batch (data.Batch).

        Returns:
            torch.Tensor.

        Raises:
            base.ConfigurationError: Features column specified but no features
                encoder specified.
            base.ConfigurationError: Features encoder specified but no features
                column specified.
        """
        encoded = self.source_encoder(
            batch.source, self.embeddings, is_source=True
        )
        mask = batch.source.mask
        if batch.has_features and not self.has_features_encoder:
            raise base.ConfigurationError(
                "Features column provided but no features encoder specified"
            )
        if self.has_features_encoder:
            if not batch.has_features:
                raise base.ConfigurationError(
                    "Features encoder specified but "
                    "no features column specified"
                )
            features_encoded = self.features_encoder(
                batch.features,
                self.embeddings,
                is_source=False,
            )
            encoded = torch.cat((encoded, features_encoded), dim=1)
            mask = torch.cat((mask, batch.features.mask), dim=1)
        if self.teacher_forcing and (self.training or self.validating):
            batch_size = len(batch)
            symbol = self.start_symbol(batch_size)
            target = torch.cat((symbol, batch.target.tensor), dim=1)
            target_mask = torch.cat(
                (
                    torch.zeros_like(symbol, dtype=bool),
                    batch.target.mask,
                ),
                dim=1,
            )
            decoded, _ = self.decoder(
                encoded, mask, target, target_mask, self.embeddings
            )
            # -> B x target_vocab_size x seq_len.
            logits = self.classifier(decoded).transpose(1, 2)
            # Truncates the prediction generated by the END_IDX token, which
            # corresponds to nothing in the target tensor.
            return logits[:, :, :-1]
        else:
            return self.greedy_decode(encoded, mask)

    def get_decoder(self) -> modules.TransformerDecoder:
        return modules.TransformerDecoder(
            decoder_input_size=self.source_encoder.output_size,
            dropout=self.decoder_dropout,
            embedding_size=self.embedding_size,
            hidden_size=self.decoder_hidden_size,
            layers=self.decoder_layers,
            max_length=self.max_target_length + 1,
            num_embeddings=self.num_embeddings,
            attention_heads=self.attention_heads,
        )

    def greedy_decode(
        self,
        encoded: torch.Tensor,
        mask: torch.Tensor,
    ) -> torch.Tensor:
        """Decodes the output sequence greedily.

        This performs student forcing.

        Args:
            encoded (torch.Tensor).
            mask (torch.Tensor).

        Returns:
            torch.Tensor: logits from the decoder.
        """
        batch_size = mask.size(0)
        # The output distributions to be returned.
        outputs = []
        # The predicted symbols at each iteration.
        predictions = [
            torch.tensor([special.START_IDX], device=self.device).repeat(
                batch_size
            )
        ]
        final = torch.zeros(batch_size, device=self.device, dtype=bool)
        for _ in range(self.max_target_length):
            logits = self.decode_step(
                encoded,
                mask,
                torch.stack(predictions, dim=1),
            )
            outputs.append(logits)
            symbol = torch.argmax(logits, dim=1)
            predictions.append(symbol)
            final = torch.logical_or(final, symbol == special.END_IDX)
            if final.all():
                break
        # -> B x target_vocab_size x seq_len.
        outputs = torch.stack(outputs, dim=2)
        return outputs

    def init_embeddings(
        self, num_embeddings: int, embedding_size: int
    ) -> nn.Embedding:
        """Initializes the embedding layer.

        Args:
            num_embeddings (int): number of embeddings.
            embedding_size (int): dimension of embeddings.

        Returns:
            nn.Embedding: embedding layer.
        """
        return embeddings.xavier_embedding(num_embeddings, embedding_size)

    @property
    def max_length(self) -> int:
        if self.has_features_encoder:
            return self.max_source_length + self.max_features_length
        else:
            return self.max_source_length

    @property
    def name(self) -> str:
        return "transformer"
