from __future__ import annotations

import math

from typing import List, NamedTuple, Union

import torch
import torchaudio

torchaudio._extension._load_lib("libctc_prefix_decoder")
import torchaudio.lib.pybind11_prefixctc as cuctc


__all__ = ["CUCTCHypothesis", "CUCTCDecoder", "cuda_ctc_decoder"]


def _get_vocab_list(vocab_file):
    vocab = []
    with open(vocab_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip().split()
            vocab.append(line[0])
    return vocab


class CUCTCHypothesis(NamedTuple):
    r"""Represents hypothesis generated by CUCTC beam search decoder :class:`CUCTCDecoder`."""
    tokens: List[int]
    """Predicted sequence of token IDs. Shape `(L, )`, where `L` is the length of the output sequence"""

    words: List[str]
    """List of predicted tokens. Algin with modeling unit.
    """

    score: float
    """Score corresponding to hypothesis"""


_DEFAULT_SKIP_THREASHOLD = math.log(0.95)


class CUCTCDecoder:
    """CUDA CTC beam search decoder.

    .. devices:: CUDA

    Note:
        To build the decoder, please use the factory function :func:`cuda_ctc_decoder`.
    """

    def __init__(
        self,
        vocab_list: List[str],
        blank_id: int = 0,
        beam_size: int = 10,
        nbest: int = 1,
        blank_skip_threshold: float = _DEFAULT_SKIP_THREASHOLD,
        cuda_stream: torch.cuda.streams.Stream = None,
    ):
        """
        Args:
            blank_id (int): token id corresopnding to blank (Default: 0)
            vocab_list (List[str]): list of vocabulary tokens
            beam_size (int, optional): max number of hypos to hold after each decode step (Default: 10)
            nbest (int): number of best decodings to return
            blank_skip_threshold (float): skip frames if log_prob(blank) > blank_skip_threshold, to speed up decoding.
                (Default: log(0.95)).
            cuda_stream (torch.cuda.streams.Stream): using assigned cuda stream (Default: using default stream)

        """
        if cuda_stream:
            if not isinstance(cuda_stream, torch.cuda.streams.Stream):
                raise AssertionError("cuda_stream must be torch.cuda.streams.Stream")
        cuda_stream_ = cuda_stream.cuda_stream if cuda_stream else torch.cuda.current_stream().cuda_stream
        self.internal_data = cuctc.prefixCTC_alloc(cuda_stream_)
        self.memory = torch.empty(0, dtype=torch.int8, device=torch.device("cuda"))
        self.blank_id = 0  # blank id has to be zero
        self.vocab_list = vocab_list
        self.space_id = 0
        self.nbest = nbest
        self.blank_skip_threshold = blank_skip_threshold
        self.beam_size = min(beam_size, len(vocab_list))  # beam size must be smaller than vocab size

    def __del__(self):
        if cuctc is not None:
            cuctc.prefixCTC_free(self.internal_data)

    def __call__(self, log_prob: torch.Tensor, encoder_out_lens: torch.Tensor):
        """
        Args:
            log_prob (torch.FloatTensor): GPU tensor of shape `(batch, frame, num_tokens)` storing sequences of
                probability distribution over labels; log_softmax(output of acoustic model).
            lengths (dtype torch.int32): GPU tensor of shape `(batch, )` storing the valid length of
                in time axis of the output Tensor in each batch.

        Returns:
            List[List[CUCTCHypothesis]]:
                List of sorted best hypotheses for each audio sequence in the batch.
        """
        if not encoder_out_lens.dtype == torch.int32:
            raise AssertionError("encoder_out_lens must be torch.int32")
        if not log_prob.dtype == torch.float32:
            raise AssertionError("log_prob must be torch.float32")
        if not (log_prob.is_cuda and encoder_out_lens.is_cuda):
            raise AssertionError("inputs must be cuda tensors")
        if not (log_prob.is_contiguous() and encoder_out_lens.is_contiguous()):
            raise AssertionError("input tensors must be contiguous")
        required_size, score_hyps = cuctc.ctc_beam_search_decoder_batch_gpu_v2(
            self.internal_data,
            self.memory.data_ptr(),
            self.memory.size(0),
            log_prob.data_ptr(),
            encoder_out_lens.data_ptr(),
            log_prob.size(),
            log_prob.stride(),
            self.beam_size,
            self.blank_id,
            self.space_id,
            self.blank_skip_threshold,
        )
        if required_size > 0:
            self.memory = torch.empty(required_size, dtype=torch.int8, device=log_prob.device).contiguous()
            _, score_hyps = cuctc.ctc_beam_search_decoder_batch_gpu_v2(
                self.internal_data,
                self.memory.data_ptr(),
                self.memory.size(0),
                log_prob.data_ptr(),
                encoder_out_lens.data_ptr(),
                log_prob.size(),
                log_prob.stride(),
                self.beam_size,
                self.blank_id,
                self.space_id,
                self.blank_skip_threshold,
            )
        batch_size = len(score_hyps)
        hypos = []
        for i in range(batch_size):
            hypos.append(
                [
                    CUCTCHypothesis(
                        tokens=score_hyps[i][j][1],
                        words=[self.vocab_list[word_id] for word_id in score_hyps[i][j][1]],
                        score=score_hyps[i][j][0],
                    )
                    for j in range(self.nbest)
                ]
            )
        return hypos


def cuda_ctc_decoder(
    tokens: Union[str, List[str]],
    nbest: int = 1,
    beam_size: int = 10,
    blank_skip_threshold: float = _DEFAULT_SKIP_THREASHOLD,
) -> CUCTCDecoder:
    """Builds an instance of :class:`CUCTCDecoder`.

    Args:
        tokens (str or List[str]): File or list containing valid tokens.
            If using a file, the expected format is for tokens mapping to the same index to be on the same line
        beam_size (int, optional): The maximum number of hypos to hold after each decode step (Default: 10)
        nbest (int): The number of best decodings to return
        blank_id (int): The token ID corresopnding to the blank symbol.
        blank_skip_threshold (float): skip frames if log_prob(blank) > blank_skip_threshold, to speed up decoding
            (Default: log(0.95)).

    Returns:
        CUCTCDecoder: decoder

    Example
        >>> decoder = cuda_ctc_decoder(
        >>>     vocab_file="tokens.txt",
        >>>     blank_skip_threshold=math.log(0.95),
        >>> )
        >>> results = decoder(log_probs, encoder_out_lens) # List of shape (B, nbest) of Hypotheses
    """
    if type(tokens) == str:
        tokens = _get_vocab_list(tokens)

    return CUCTCDecoder(vocab_list=tokens, beam_size=beam_size, nbest=nbest, blank_skip_threshold=blank_skip_threshold)
