"""Abstract decoder definition for Transducer models."""

from abc import ABC, abstractmethod
from typing import Any, List, Optional, Tuple

import torch


class AbsDecoder(torch.nn.Module, ABC):
    """Abstract decoder module."""

    @abstractmethod
    def forward(self, labels: torch.Tensor) -> torch.Tensor:
        """Encode source label sequences.

        Args:
            labels: Label ID sequences. (B, L)

        Returns:
            dec_out: Decoder output sequences. (B, T, D_dec)

        """
        raise NotImplementedError

    @abstractmethod
    def score(
        self,
        label: torch.Tensor,
        label_sequence: List[int],
        dec_state: Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]],
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]]]:
        """One-step forward hypothesis.

        Args:
            label: Previous label. (1, 1)
            label_sequence: Current label sequence.
            dec_state: Previous decoder hidden states.
                         ((N, 1, D_dec), (N, 1, D_dec) or None) or None

        Returns:
            dec_out: Decoder output sequence. (1, D_dec) or (1, D_emb)
            dec_state: Decoder hidden states.
                         ((N, 1, D_dec), (N, 1, D_dec) or None) or None

        """
        raise NotImplementedError

    @abstractmethod
    def batch_score(
        self,
        hyps: List[Any],
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]]]:
        """One-step forward hypotheses.

        Args:
            hyps: Hypotheses.

        Returns:
            dec_out: Decoder output sequences. (B, D_dec) or (B, D_emb)
            states: Decoder hidden states.
                      ((N, B, D_dec), (N, B, D_dec) or None) or None

        """
        raise NotImplementedError

    @abstractmethod
    def set_device(self, device: torch.Tensor) -> None:
        """Set GPU device to use.

        Args:
            device: Device ID.

        """
        raise NotImplementedError

    @abstractmethod
    def init_state(
        self, batch_size: int
    ) -> Optional[Tuple[torch.Tensor, Optional[torch.tensor]]]:
        """Initialize decoder states.

        Args:
            batch_size: Batch size.

        Returns:
            : Initial decoder hidden states.
                ((N, B, D_dec), (N, B, D_dec) or None) or None

        """
        raise NotImplementedError

    @abstractmethod
    def select_state(
        self,
        states: Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]] = None,
        idx: int = 0,
    ) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor]]]:
        """Get specified ID state from batch of states, if provided.

        Args:
            states: Decoder hidden states.
                      ((N, B, D_dec), (N, B, D_dec) or None) or None
            idx: State ID to extract.

        Returns:
            : Decoder hidden state for given ID.
                ((N, 1, D_dec), (N, 1, D_dec) or None) or None

        """
        raise NotImplementedError