"""Pointer-generator model classes."""

from typing import Callable, Optional, Tuple, Union

import torch
from torch import nn

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


class PointerGeneratorModel(base.BaseModel):
    """Base class for pointer-generator models."""

    def _get_loss_func(
        self,
    ) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]:
        """Returns the actual function used to compute loss.

        This overrides the loss function behavior in
        models.base.BaseModel because we need to use NLLLoss in
        order to factor the addition of two separate probability
        distributions. An NLLLoss-compatible implementation of label smoothing
        is also provided when label smoothing is requested.

        Returns:
            Callable[[torch.Tensor, torch.Tensor], torch.Tensor]: configured
                loss function.
        """
        if not self.label_smoothing:
            return nn.NLLLoss(ignore_index=special.PAD_IDX)
        else:
            return self._smooth_nllloss

    def _smooth_nllloss(
        self, predictions: torch.Tensor, target: torch.Tensor
    ) -> torch.Tensor:
        """Computes the NLLLoss with a smoothing factor such that some
        proportion of the output distribution is replaced with a
        uniform distribution.

        After:
            https://github.com/NVIDIA/DeepLearningExamples/blob/
            8d8b21a933fff3defb692e0527fca15532da5dc6/PyTorch/Classification/
            ConvNets/image_classification/smoothing.py#L18

        Args:
            predictions (torch.Tensor): tensor of prediction
                distribution of shape B x target_vocab_size x seq_len.
            target (torch.Tensor): tensor of golds of shape
                B x seq_len.

        Returns:
            torch.Tensor: loss.
        """
        # -> (B * seq_len) x target_vocab_size.
        predictions = predictions.transpose(1, 2).reshape(
            -1, self.target_vocab_size
        )
        # -> (B * seq_len) x 1.
        target = target.view(-1, 1)
        non_pad_mask = target.ne(special.PAD_IDX)
        # Gets the ordinary loss.
        nll_loss = -predictions.gather(dim=-1, index=target)[
            non_pad_mask
        ].mean()
        # Gets the smoothed loss.
        smooth_loss = -predictions.sum(dim=-1, keepdim=True)[
            non_pad_mask
        ].mean()
        smooth_loss = smooth_loss / self.target_vocab_size
        # Combines both according to label smoothing weight.
        loss = (1.0 - self.label_smoothing) * nll_loss
        loss.add_(self.label_smoothing * smooth_loss)
        return loss


class PointerGeneratorRNNModel(PointerGeneratorModel):
    """Abstract base class for pointer-generator models with RNN backends.

    If features are provided, a separate features attention module computes
    the feature encodings. Because of this, the source and features encoders
    can have differently-sized hidden layers. However, because of the way
    they are combined, they must have the same number of hidden layers.

    Whereas See et al. use a two-layer MLP  for the vocabulary distribution,
    this uses a single linear layer.

    This supports optional student forcing during training/validation.

    After:
        See, A., Liu, P. J., and Manning, C. D. 2017. Get to the point:
        summarization with pointer-generator networks. In _Proceedings of the
        55th Annual Meeting of the Association for Computational Linguistics
        (Volume 1: Long Papers)_, pages 1073-1083.

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

    Raises:
        base.ConfigurationError: Number of encoder layers and decoder layers
            must match.
    """

    classifier: nn.Linear
    generation_probability: modules.GenerationProbability
    teacher_forcing: bool

    def __init__(
        self,
        *args,
        teacher_forcing: bool = defaults.TEACHER_FORCING,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        if self.has_features_encoder:
            self.features_attention = modules.Attention(
                self.features_encoder.output_size, self.decoder_hidden_size
            )
        self.decoder = self.get_decoder()
        self.generation_probability = modules.GenerationProbability(
            self.embedding_size,
            self.decoder_hidden_size,
            self.decoder_input_size,
        )
        self.classifier = nn.Linear(
            self.decoder_hidden_size + self.decoder_input_size,
            self.target_vocab_size,
        )
        # Compatibility check.
        if self.source_encoder.layers != self.decoder_layers:
            raise base.ConfigurationError(
                f"Number of encoder layers ({self.source_encoder.layers}) and "
                f"decoder layers ({self.decoder_layers}) must match"
            )
        self.teacher_forcing = teacher_forcing
        self._log_model()
        self.save_hyperparameters(
            ignore=[
                "classifier",
                "decoder",
                "embeddings",
                "generation_probability",
                "features_encoder",
                "source_encoder",
            ]
        )

    def beam_decode(
        self,
        source: torch.Tensor,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        *,
        features_encoded: Optional[torch.Tensor] = None,
        features_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Decodes with beam search.

        Implementationally this is almost identical to the method of the same
        name in RNNModel.

        Args:
            source (torch.Tensor): source symbols, used to compute pointer
                weights.
            source_encoded (torch.Tensor): encoded source symbols.
            source_mask (torch.Tensor): mask for the source.
            features_encoded (torch.Tensor, optional): encoded feaure symbols.
            features_mask (torch.Tensor, optional): mask for the features.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: predictions of shape
                B x beam_width x seq_len and log-likelihoods of shape
                B x beam_width.

        Raises:
            NotImplementedError: Beam search is not implemented for
                batch_size > 1.
        """
        # TODO: modify to work with batches larger than 1.
        batch_size = source_mask.size(0)
        if batch_size != 1:
            raise NotImplementedError(
                "Beam search is not implemented for batch_size > 1"
            )
        state = self.decoder.initial_state(batch_size)
        # The start symbol is not needed here because the beam puts that in
        # automatically.
        beam = beam_search.Beam(self.beam_width, state)
        for _ in range(self.max_target_length):
            for cell in beam.cells:
                if cell.final:
                    beam.push(cell)
                else:
                    symbol = torch.tensor([[cell.symbol]], device=self.device)
                    prediction, state = self.decode_step(
                        source,
                        source_encoded,
                        source_mask,
                        symbol,
                        cell.state,
                        features_encoded,
                        features_mask,
                    )
                    logits = nn.functional.log_softmax(
                        prediction.squeeze(), dim=0
                    )
                    for new_cell in cell.extensions(state, logits):
                        beam.push(new_cell)
            beam.update()
            if beam.final:
                break
        return beam.predictions(self.device), beam.logits(self.device)

    def decode_step(
        self,
        source: torch.Tensor,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        symbol: torch.Tensor,
        state: modules.RNNState,
        *,
        features_encoded: Optional[torch.Tensor] = None,
        features_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Single decoder step.

        This predicts a distribution for one symbol.

        Args:
            source (torch.Tensor): source symbols, used to compute pointer
                weights.
            source_encoded (torch.Tensor): encoded source symbols.
            source_mask (torch.Tensor): mask for the source.
            symbol (torch.Tensor): next symbol.
            state (modules.RNNState): RNN state.
            features_encoded (torch.Tensor, optional): encoded features
                symbols.
            features_mask (torch.Tensor, optional): mask for the features.

        Returns:
            Tuple[torch.Tensor, modules.RNNState]: predictions for that state
                and the RNN state.
        """
        # TODO: there are a few Law of Demeter violations here. Is there an
        # obvious refactoring?
        embedded = self.decoder.embed(symbol, self.embeddings)
        context, attention_weights = self.decoder.attention(
            source_encoded,
            state.hidden.transpose(0, 1),
            source_mask,
        )
        if self.has_features_encoder:
            features_context, _ = self.features_attention(
                features_encoded,
                state.hidden.transpose(0, 1),
                features_mask,
            )
            context = torch.cat((context, features_context), dim=2)
        _, state = self.decoder.module(
            torch.cat((embedded, context), dim=2), state
        )
        hidden = state.hidden[-1, :, :].unsqueeze(1)
        output_dist = nn.functional.softmax(
            self.classifier(torch.cat((hidden, context), dim=2)),
            dim=2,
        )
        # -> B x 1 x target_vocab_size.
        pointer_dist = torch.zeros(
            (symbol.size(0), 1, self.target_vocab_size),
            device=self.device,
            dtype=attention_weights.dtype,
        )
        # Gets the attentions to the source in terms of the output generations.
        # These are the "pointer" distribution.
        pointer_dist.scatter_add_(2, source.unsqueeze(1), attention_weights)
        # Probability of generating from output_dist.
        gen_probs = self.generation_probability(context, hidden, embedded)
        scaled_output_dist = output_dist * gen_probs
        scaled_pointer_dist = pointer_dist * (1 - gen_probs)
        # First argument is log-probs.
        return torch.log(scaled_output_dist + scaled_pointer_dist), state

    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.normal_embedding(num_embeddings, embedding_size)

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

        Args:
            batch (data.Batch).

        Returns:
            Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]: beam search
                returns a tuple with a tensor of predictions of shape
                B x beam_width x seq_len and a tensor of B x shape beam_width
                with the likelihood (the unnormalized sum of sequence
                log-probabilities) for each prediction; greedy search returns
                a tensor of predictions of shape
                B x target_vocab_size x seq_len.

        Raises:
            base.ConfigurationError: Features encoder specified but no feature
                column specified.
            base.ConfigurationError: Features column specified but no feature
                encoder specified.
        """
        source_encoded = self.source_encoder(
            batch.source, self.embeddings, is_source=True
        )
        if self.has_features_encoder:
            if not batch.has_features:
                raise base.ConfigurationError(
                    "Features encoder specified but "
                    "no feature column specified"
                )
            features_encoded = self.features_encoder(
                batch.features,
                self.embeddings,
                is_source=False,
            )
            if self.beam_width > 1:
                return self.beam_decode(
                    batch.source.tensor,
                    source_encoded,
                    batch.source.mask,
                    features_encoded=features_encoded,
                    features_mask=batch.features.mask,
                )
            elif self.training or self.validating:
                # This version supports teacher forcing.
                return self.greedy_decode_train_validate(
                    batch.source.tensor,
                    source_encoded,
                    batch.source.mask,
                    target=(
                        batch.target.tensor if self.teacher_forcing else None
                    ),
                    features_encoded=features_encoded,
                    features_mask=batch.features.mask,
                )
            else:
                return self.greedy_decode_predict_test(
                    batch.source.tensor,
                    source_encoded,
                    batch.source.mask,
                    features_encoded=features_encoded,
                    features_mask=batch.features.mask,
                )
        elif batch.has_features:
            raise base.ConfigurationError(
                "Features column specified but no feature encoder specified"
            )
        elif self.beam_width > 1:
            return self.beam_decode(
                batch.source.tensor, source_encoded, batch.source.mask
            )
        elif self.training or self.validating:
            # This version supports teacher forcing.
            return self.greedy_decode_train_validate(
                batch.source.tensor,
                source_encoded,
                batch.source.mask,
                target=batch.target.tensor if self.teacher_forcing else None,
            )
        else:
            return self.greedy_decode_predict_test(
                batch.source.tensor,
                source_encoded,
                batch.source.mask,
            )

    def greedy_decode_train_validate(
        self,
        source: torch.Tensor,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        *,
        target: Optional[torch.Tensor] = None,
        features_encoded: Optional[torch.Tensor] = None,
        features_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Decodes greedily during training and validation.

        Provide the target for teacher forcing.

        If the target is provided, decoding halts at target length. If the
        target is not provided, decoding will halt once each sequence in the
        batch generates END or the maximum target length is reached,
        whichever comes first.

        Args:
            source (torch.Tensor): source symbols, needed to compute pointer
                weights.
            source_encoded (torch.Tensor).
            source_mask (torch.Tensor).
            features_encoded (torch.Tensor, optional).
            features_mask (torch.Tensor, optional).
            target (torch.Tensor, optional): target symbols; if provided,
                these are used for teacher forcing.

        Returns:
            torch.Tensor: predictions of B x target_vocab_size x seq_len.
        """
        batch_size = source_mask.size(0)
        symbol = self.start_symbol(batch_size)
        state = self.decoder.initial_state(batch_size)
        predictions = []
        if target is None:
            target_length = self.max_target_length
            final = torch.zeros(batch_size, device=self.device, dtype=bool)
        else:
            target_length = target.size(1)
        for t in range(target_length):
            log_probs, state = self.decode_step(
                source,
                source_encoded,
                source_mask,
                symbol,
                state,
                features_encoded=features_encoded,
                features_mask=features_mask,
            )
            predictions.append(log_probs.squeeze(1))
            if target is None:
                # Student forcing.
                symbol = log_probs.argmax(dim=2)
                final = torch.logical_or(final, symbol == special.END_IDX)
                if final.all():
                    break
            else:
                # Teacher forcing.
                symbol = target[:, t].unsqueeze(1)
        predictions = torch.stack(predictions, dim=2)
        return predictions

    def greedy_decode_predict_test(
        self,
        source: torch.Tensor,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        *,
        features_encoded: Optional[torch.Tensor] = None,
        features_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Decodes greedily during prediction and testing.

        These are different because teacher forcing is not supported, but
        decoding halts once each sequence in the batch generates END or the
        maximum target length is reached, whichever comes first.

        Args:
            source (torch.Tensor): source symbols, needed to compute pointer
                weights.
            source_encoded (torch.Tensor).
            source_mask (torch.Tensor).
            features_encoded (torch.Tensor, optional).
            features_mask (torch.Tensor, optional).
            target (torch.Tensor, optional): target symbols; if provided,
                these are used for teacher forcing.
            target_length (int, optional): maximum target length during
                decoding. If not specified, max_target_length is used.

        Returns:
            torch.Tensor: predictions of B x target_vocab_size x seq_len.
        """
        batch_size = source_mask.size(0)
        symbol = self.start_symbol(batch_size)
        state = self.decoder.initial_state(batch_size)
        predictions = []
        final = torch.zeros(batch_size, device=self.device, dtype=bool)
        for _ in range(self.max_target_length):
            log_probs, state = self.decode_step(
                source,
                source_encoded,
                source_mask,
                symbol,
                state,
                features_encoded=features_encoded,
                features_mask=features_mask,
            )
            predictions.append(log_probs.squeeze(1))
            symbol = log_probs.argmax(dim=2)
            final = torch.logical_or(final, symbol == special.END_IDX)
            if final.all():
                break
        predictions = torch.stack(predictions, dim=2)
        return predictions

    @property
    def decoder_input_size(self) -> int:
        # We concatenate along the encoding dimension.
        if self.has_features_encoder:
            return (
                self.source_encoder.output_size
                + self.features_encoder.output_size
            )
        else:
            return self.source_encoder.output_size


class PointerGeneratorGRUModel(PointerGeneratorRNNModel):
    """Pointer-generator model with an GRU backend."""

    def get_decoder(self) -> modules.SoftAttentionGRUDecoder:
        return modules.SoftAttentionGRUDecoder(
            attention_input_size=self.source_encoder.output_size,
            decoder_input_size=self.decoder_input_size,
            dropout=self.decoder_dropout,
            embedding_size=self.embedding_size,
            hidden_size=self.decoder_hidden_size,
            layers=self.decoder_layers,
            num_embeddings=self.num_embeddings,
        )

    @property
    def name(self) -> str:
        return "pointer-generator GRU"


class PointerGeneratorLSTMModel(PointerGeneratorRNNModel):
    """Pointer-generator model with an LSTM backend."""

    def get_decoder(self) -> modules.SoftAttentionLSTMDecoder:
        return modules.SoftAttentionLSTMDecoder(
            attention_input_size=self.source_encoder.output_size,
            decoder_input_size=self.decoder_input_size,
            dropout=self.decoder_dropout,
            embedding_size=self.embedding_size,
            hidden_size=self.decoder_hidden_size,
            layers=self.decoder_layers,
            num_embeddings=self.num_embeddings,
        )

    @property
    def name(self) -> str:
        return "pointer-generator LSTM"


class PointerGeneratorTransformerModel(PointerGeneratorModel):
    """Pointer-generator model with a transformer backend.

    After:
        Singer, A., and Kann, K. 2020. The NYU-CUBoulder Systems for
        SIGMORPHON 2020 Task 0 and Task 2. In _Proceedings of the 17th
        SIGMORPHON Workshop on Computational Research in Phonetics, Phonology,
        and Morphology_, pages 90–98.

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

    attention_heads: int
    classifier: nn.Linear
    decoder_positional_encoding: modules.BasePositionalEncoding
    generation_probability: modules.GenerationProbability
    teacher_forcing: bool

    def __init__(
        self,
        *args,
        attention_heads: int = defaults.ATTENTION_HEADS,
        decoder_positional_encoding: Optional[
            modules.BasePositionalEncoding
        ] = None,
        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
        )
        self.decoder_positional_encoding = decoder_positional_encoding
        self.decoder = self.get_decoder()
        if self.has_features_encoder:
            self.generation_probability = modules.GenerationProbability(
                self.embedding_size,
                self.embedding_size,
                self.embedding_size,
            )
        else:
            self.generation_probability = modules.GenerationProbability(
                self.embedding_size,
                self.embedding_size,
                self.source_encoder.output_size,
            )
        self.teacher_forcing = teacher_forcing
        self._log_model()
        self.save_hyperparameters(
            ignore=[
                "classifier",
                "decoder",
                # This lives in the decoder.
                "decoder_positional_encoding",
                "embeddings",
                "features_encoder",
                "generation_probability",
                "source_encoder",
            ],
        )

    def decode_step(
        self,
        source: torch.Tensor,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        target: torch.Tensor,
        target_mask: torch.Tensor,
        *,
        features_encoded: Optional[torch.Tensor] = None,
        features_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Single decoder step.

        This will work on any sequence length, and returns output
        probabilities for all targets, meaning we can use this method for
        student forcing, in which only a single new token is decoded at a time,
        or for teacher forcing, in which all tokens are be decoded in
        parallel using a diagonal mask.

        Args:
            source (torch.Tensor): source symbols, used to compute pointer
                weights.
            source_encoded (torch.Tensor): encoded source_symbols.
            source_mask (torch.Tensor): mask for the source.
            target (torch.Tensor): tensor of predictions thus far.
            target_mask (torch.Tensor): mask for the target.
            features_encoded (torch.Tensor, optional): encoded features
                symbols.
            features_mask (torch.Tensor, optional): mask for the features.

        Returns:
            torch.Tensor: predictions for that state.
        """
        self.decoder.attention_output.clear()
        if self.has_features_encoder:
            decoded, target_embedded = self.decoder(
                source_encoded,
                source_mask,
                target,
                target_mask,
                self.embeddings,
                features_encoded=features_encoded,
                features_mask=features_mask,
            )
        else:
            decoded, target_embedded = self.decoder(
                source_encoded,
                source_mask,
                target,
                target_mask,
                self.embeddings,
            )
        # Outputs from the multi-headed attention from each decoder step to
        # the encoded source; values have been averaged over each attention
        # head.
        # -> B x target_seq_len x source_seq_len.
        mha_outputs = self.decoder.attention_output[0]
        logits = self.classifier(decoded)
        output_dist = nn.functional.softmax(logits, dim=2)
        # -> B x target-seq_len x target_vocab_size.
        pointer_dist = torch.zeros(
            mha_outputs.size(0),
            mha_outputs.size(1),
            self.target_vocab_size,
            device=self.device,
            dtype=mha_outputs.dtype,
        )
        # Repeats the source indices for each target.
        # -> B x target_seq_len x source_seq_len.
        repeated_source = source.unsqueeze(1).repeat(1, mha_outputs.size(1), 1)
        # Scatters the attention weights onto the pointer_dist at their vocab
        # indices in order to get outputs that match the indexing of the
        # generation probability.
        pointer_dist.scatter_add_(2, repeated_source, mha_outputs)
        # A matrix of context vectors from applying attention to the encoder
        # representations w.r.t. each decoder step.
        context = torch.bmm(mha_outputs, source_encoded)
        # Probability of generating from output_dist.
        gen_probs = self.generation_probability(
            context, decoded, target_embedded
        )
        scaled_pointer_dist = pointer_dist * (1 - gen_probs)
        scaled_output_dist = output_dist * gen_probs
        return torch.log(scaled_output_dist + scaled_pointer_dist)

    def get_decoder(self) -> modules.PointerGeneratorTransformerDecoder:
        return modules.PointerGeneratorTransformerDecoder(
            attention_heads=self.attention_heads,
            decoder_input_size=self.source_encoder.output_size,
            dropout=self.decoder_dropout,
            embedding_size=self.embedding_size,
            has_features_encoder=self.has_features_encoder,
            hidden_size=self.decoder_hidden_size,
            layers=self.decoder_layers,
            max_length=self.max_decoder_length,
            num_embeddings=self.num_embeddings,
            positional_encoding=self.decoder_positional_encoding,
        )

    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)

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

        Args:
            batch (data.Batch).

        Returns:
            torch.Tensor.

        Raises:
            base.ConfigurationError: Features encoder specified but no feature
                column specified.
            base.ConfigurationError: Features column specified but no feature
                encoder specified.
        """
        source_encoded = self.source_encoder(
            batch.source, self.embeddings, is_source=True
        )
        if self.has_features_encoder:
            if not batch.has_features:
                raise base.ConfigurationError(
                    "Features encoder specified but "
                    "no feature column specified"
                )
            features_encoded = self.features_encoder(
                batch.features,
                self.embeddings,
                is_source=False,
            )
            if (self.training or self.validating) and self.teacher_forcing:
                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,
                )
                log_probs = self.decode_step(
                    batch.source.tensor,
                    source_encoded,
                    batch.source.mask,
                    target,
                    target_mask,
                    features_encoded=features_encoded,
                    features_mask=batch.features.mask,
                )
                # Truncates the prediction generated by the END_IDX token,
                # which corresponds to nothing in the target tensor.
                return log_probs[:, :-1, :].transpose(1, 2)
            else:
                return self.greedy_decode(
                    batch.source.tensor,
                    source_encoded,
                    batch.source.mask,
                    features_encoded=features_encoded,
                    features_mask=batch.features.mask,
                )
        elif batch.has_features:
            raise base.ConfigurationError(
                "Feature column specified but no feature encoder specified"
            )
        if (self.training or self.validating) and self.teacher_forcing:
            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,
            )
            log_probs = self.decode_step(
                batch.source.tensor,
                source_encoded,
                batch.source.mask,
                target,
                target_mask,
            )
            # Truncates the prediction generated by the END_IDX token, which
            # corresponds to nothing in the target tensor.
            return log_probs[:, :-1, :].transpose(1, 2)
        else:
            return self.greedy_decode(
                batch.source.tensor,
                source_encoded,
                batch.source.mask,
            )

    def greedy_decode(
        self,
        source: torch.Tensor,
        source_encoded: torch.Tensor,
        source_mask: torch.Tensor,
        *,
        features_encoded: Optional[torch.Tensor] = None,
        features_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Decodes greedily during prediction and testing.

        This performs student forcing.

        Decoding halts once each sequence in the batch generates END or the
        maximum target length is reached, whichever comes first.

        Args:
            source (torch.Tensor): source symbols, used to compute pointer
                weights.
            source_encoded (torch.Tensor): encoded source symbols.
            source_mask (torch.Tensor): mask for the source.
            features_encoded (torch.Tensor, optional): encoded feaure symbols.
            features_mask (torch.Tensor, optional): mask for the features.

        Returns:
            torch.Tensor: predictions of B x target_vocab_size x seq_len.
        """
        batch_size = source_encoded.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):
            if self.has_features_encoder:
                assert features_encoded is not None
                assert features_mask is not None
                log_probs = self.decode_step(
                    source,
                    source_encoded,
                    source_mask,
                    torch.stack(predictions, dim=1),
                    None,
                    features_encoded=features_encoded,
                    features_mask=features_mask,
                )
            else:
                log_probs = self.decode_step(
                    source,
                    source_encoded,
                    source_mask,
                    torch.stack(predictions, dim=1),
                    None,
                )
            log_probs = log_probs[:, -1, :]  # From last symbol.
            outputs.append(log_probs)
            symbol = log_probs.argmax(dim=1)
            final = torch.logical_or(final, symbol == special.END_IDX)
            if final.all():
                break
            predictions.append(symbol)
        outputs = torch.stack(outputs, dim=2)
        return outputs

    @property
    def max_decoder_length(self) -> int:
        return self.max_target_length + 1  # Including the START symbol.

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


class RotaryPointerGeneratorTransformerModel(PointerGeneratorTransformerModel):
    """Pointer-generator transformer model using rotary positional encoding.

    RoPE is applied inside attention layers rather than to the token
    embeddings, so this model does not accept a decoder_positional_encoding
    argument. Passing one is a configuration error.

    All encoders used with this model should also use rotary positional
    encoding. Specifically, use RotaryTransformerEncoder or
    RotaryFeatureInvariantTransformerEncoder as the source_encoder and
    features_encoder. Mixing rotary decoders with non-rotary encoders
    is possible but is not recommended.

    Args:
        *args: passed to superclass.
        **kwargs: passed to superclass.
    """

    def __init__(self, *args, **kwargs):
        if kwargs.get("decoder_positional_encoding") is not None:
            raise base.ConfigurationError(
                f"{self.__class__.__name__} does not accept "
                "decoder_positional_encoding"
            )
        super().__init__(*args, **kwargs)

    def get_decoder(
        self,
    ) -> modules.RotaryPointerGeneratorTransformerDecoder:
        return modules.RotaryPointerGeneratorTransformerDecoder(
            attention_heads=self.attention_heads,
            decoder_input_size=self.source_encoder.output_size,
            dropout=self.decoder_dropout,
            embedding_size=self.embedding_size,
            has_features_encoder=self.has_features_encoder,
            hidden_size=self.decoder_hidden_size,
            layers=self.decoder_layers,
            max_length=self.max_decoder_length,
            num_embeddings=self.num_embeddings,
        )

    @property
    def name(self) -> str:
        return f"rotary {super().name}"
