# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import math
import os
from dataclasses import dataclass, field
from typing import List, Optional, Tuple, Union

import torch

from nemo.collections.asr.parts.context_biasing import BoostingTreeModelConfig, GPUBoostingTreeModel
from nemo.collections.asr.parts.k2.classes import GraphIntersectDenseConfig
from nemo.collections.asr.parts.submodules.ctc_batched_beam_decoding import BatchedBeamCTCComputer
from nemo.collections.asr.parts.submodules.ngram_lm import DEFAULT_TOKEN_OFFSET, NGramGPULanguageModel
from nemo.collections.asr.parts.submodules.wfst_decoder import RivaDecoderConfig, WfstNbestHypothesis
from nemo.collections.asr.parts.utils import rnnt_utils
from nemo.collections.common.parts.optional_cuda_graphs import WithOptionalCudaGraphs
from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
from nemo.core.classes import Typing, typecheck
from nemo.core.neural_types import HypothesisType, LengthsType, LogprobsType, NeuralType
from nemo.utils import logging


def pack_hypotheses(
    hypotheses: List[rnnt_utils.NBestHypotheses],
    logitlen: torch.Tensor,
) -> List[rnnt_utils.NBestHypotheses]:

    if logitlen is not None:
        if hasattr(logitlen, 'cpu'):
            logitlen_cpu = logitlen.to('cpu')
        else:
            logitlen_cpu = logitlen

    for idx, hyp in enumerate(hypotheses):  # type: rnnt_utils.NBestHypotheses
        for candidate_idx, cand in enumerate(hyp.n_best_hypotheses):
            cand.y_sequence = torch.tensor(cand.y_sequence, dtype=torch.long)

            if logitlen is not None:
                cand.length = logitlen_cpu[idx]

            if cand.dec_state is not None:
                cand.dec_state = _states_to_device(cand.dec_state)

    return hypotheses


def pack_wfst_hypotheses(
    hypotheses: List['WfstNbestHypothesis'],
    logits: torch.Tensor,
    logitlen: torch.Tensor,
) -> List[rnnt_utils.NBestHypotheses]:

    logitlen_cpu = logitlen.to('cpu')

    new_hypotheses = []
    for idx, nbest_hyp in enumerate(hypotheses):  # type: WfstNbestHypothesis
        new_hyp = []
        y_sequence = logits[idx, : logitlen[idx]].to('cpu')
        length = logitlen_cpu[idx]
        for candidate_idx, cand in enumerate(nbest_hyp):
            cand_hyp = rnnt_utils.Hypothesis(
                y_sequence=[],
                score=cand.score,
                text=" ".join(cand.words),
                timestamp=list(cand.timesteps),
                alignments=list(cand.alignment),
            )
            cand_hyp.y_sequence = y_sequence

            if logitlen is not None:
                cand_hyp.length = length

            new_hyp.append(cand_hyp)

        new_hypotheses.append(rnnt_utils.NBestHypotheses(new_hyp))

    return new_hypotheses


def _states_to_device(dec_state, device='cpu'):
    if torch.is_tensor(dec_state):
        dec_state = dec_state.to(device)

    elif isinstance(dec_state, (list, tuple)):
        dec_state = tuple(_states_to_device(dec_i, device) for dec_i in dec_state)

    return dec_state


class AbstractBeamCTCInfer(Typing):
    """A beam CTC decoder.

    Provides a common abstraction for sample level beam decoding.

    Args:
        blank_id: int, index of the blank token. Can be 0 or len(vocabulary).
        beam_size: int, size of the beam used in the underlying beam search engine.

    """

    @property
    def input_types(self):
        """Returns definitions of module input ports."""
        return {
            "decoder_output": NeuralType(('B', 'T', 'D'), LogprobsType()),
            "decoder_lengths": NeuralType(tuple('B'), LengthsType()),
        }

    @property
    def output_types(self):
        """Returns definitions of module output ports."""
        return {"predictions": [NeuralType(elements_type=HypothesisType())]}

    def __init__(self, blank_id: int, beam_size: int):
        self.blank_id = blank_id

        if beam_size < 1:
            raise ValueError("Beam search size cannot be less than 1!")

        self.beam_size = beam_size

        # Variables set by corresponding setter methods
        self.vocab = None
        self.decoding_type = None
        self.tokenizer = None

        # Utility maps for vocabulary
        self.vocab_index_map = None
        self.index_vocab_map = None

        # Internal variable, used to prevent double reduction of consecutive tokens (ctc collapse)
        self.override_fold_consecutive_value = None

    def set_vocabulary(self, vocab: List[str]):
        """
        Set the vocabulary of the decoding framework.

        Args:
            vocab: List of str. Each token corresponds to its location in the vocabulary emitted by the model.
                Note that this vocabulary must NOT contain the "BLANK" token.
        """
        self.vocab = vocab
        self.vocab_index_map = {v: i for i, v in enumerate(vocab)}
        self.index_vocab_map = {i: v for i, v in enumerate(vocab)}

    def set_decoding_type(self, decoding_type: str):
        """
        Sets the decoding type of the framework. Can support either char or subword models.

        Args:
            decoding_type: Str corresponding to decoding type. Only supports "char" and "subword".
        """
        decoding_type = decoding_type.lower()
        supported_types = ['char', 'subword']

        if decoding_type not in supported_types:
            raise ValueError(
                f"Unsupported decoding type. Supported types = {supported_types}.\n" f"Given = {decoding_type}"
            )

        self.decoding_type = decoding_type

    def set_tokenizer(self, tokenizer: TokenizerSpec):
        """
        Set the tokenizer of the decoding framework.

        Args:
            tokenizer: NeMo tokenizer object, which inherits from TokenizerSpec.
        """
        self.tokenizer = tokenizer

    @typecheck()
    def forward(
        self,
        decoder_output: torch.Tensor,
        decoder_lengths: torch.Tensor,
    ) -> Tuple[List[Union[rnnt_utils.Hypothesis, rnnt_utils.NBestHypotheses]]]:
        """Returns a list of hypotheses given an input batch of the encoder hidden embedding.
        Output token is generated auto-repressively.

        Args:
            decoder_output: A tensor of size (batch, timesteps, features) or (batch, timesteps) (each timestep is a label).
            decoder_lengths: list of int representing the length of each sequence
                output sequence.

        Returns:
            packed list containing batch number of sentences (Hypotheses).
        """
        raise NotImplementedError()

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)


class BeamCTCInfer(AbstractBeamCTCInfer):
    """A beam CTC decoder.

    Provides a common abstraction for sample level and batch level greedy decoding.

    Args:
        blank_index: int index of the blank token. Can be 0 or len(vocabulary).
        preserve_alignments: Bool flag which preserves the history of logprobs generated during
            decoding (sample / batched). When set to true, the Hypothesis will contain
            the non-null value for `logprobs` in it. Here, `logprobs` is a torch.Tensors.
        compute_timestamps: A bool flag, which determines whether to compute the character/subword, or
                word based timestamp mapping the output log-probabilities to discrite intervals of timestamps.
                The timestamps will be available in the returned Hypothesis.timestep as a dictionary.

    """

    def __init__(
        self,
        blank_id: int,
        beam_size: int,
        search_type: str = "default",
        return_best_hypothesis: bool = True,
        preserve_alignments: bool = False,
        compute_timestamps: bool = False,
        ngram_lm_alpha: float = 0.3,
        beam_beta: float = 0.0,
        ngram_lm_model: str = None,
        flashlight_cfg: Optional['FlashlightConfig'] = None,
        pyctcdecode_cfg: Optional['PyCTCDecodeConfig'] = None,
    ):
        super().__init__(blank_id=blank_id, beam_size=beam_size)

        self.search_type = search_type
        self.return_best_hypothesis = return_best_hypothesis
        self.preserve_alignments = preserve_alignments
        self.compute_timestamps = compute_timestamps

        if self.compute_timestamps:
            raise ValueError("Currently this flag is not supported for beam search algorithms.")

        self.vocab = None  # This must be set by specific method by user before calling forward() !

        if search_type == "default" or search_type == "nemo":
            self.search_algorithm = self.default_beam_search
        elif search_type == "pyctcdecode":
            self.search_algorithm = self._pyctcdecode_beam_search
        elif search_type == "flashlight":
            self.search_algorithm = self.flashlight_beam_search
        else:
            raise NotImplementedError(
                f"The search type ({search_type}) supplied is not supported!\n"
                f"Please use one of : (default, nemo, pyctcdecode)"
            )

        # Log the beam search algorithm
        logging.info(f"Beam search algorithm: {search_type}")

        self.ngram_lm_alpha = ngram_lm_alpha
        self.beam_beta = beam_beta

        # Default beam search args
        self.ngram_lm_model = ngram_lm_model

        # PyCTCDecode params
        if pyctcdecode_cfg is None:
            pyctcdecode_cfg = PyCTCDecodeConfig()
        self.pyctcdecode_cfg = pyctcdecode_cfg  # type: PyCTCDecodeConfig

        if flashlight_cfg is None:
            flashlight_cfg = FlashlightConfig()
        self.flashlight_cfg = flashlight_cfg

        # Default beam search scorer functions
        self.default_beam_scorer = None
        self.pyctcdecode_beam_scorer = None
        self.flashlight_beam_scorer = None
        self.token_offset = 0

    @typecheck()
    def forward(
        self,
        decoder_output: torch.Tensor,
        decoder_lengths: torch.Tensor,
    ) -> Tuple[List[Union[rnnt_utils.Hypothesis, rnnt_utils.NBestHypotheses]]]:
        """Returns a list of hypotheses given an input batch of the encoder hidden embedding.
        Output token is generated auto-repressively.

        Args:
            decoder_output: A tensor of size (batch, timesteps, features).
            decoder_lengths: list of int representing the length of each sequence
                output sequence.

        Returns:
            packed list containing batch number of sentences (Hypotheses).
        """
        if self.vocab is None:
            raise RuntimeError("Please set the vocabulary with `set_vocabulary()` before calling this function.")

        if self.decoding_type is None:
            raise ValueError("Please set the decoding type with `set_decoding_type()` before calling this function.")

        with torch.no_grad(), torch.inference_mode():
            # Process each sequence independently
            prediction_tensor = decoder_output

            if prediction_tensor.ndim != 3:
                raise ValueError(
                    f"`decoder_output` must be a tensor of shape [B, T, V] (log probs, float). "
                    f"Provided shape = {prediction_tensor.shape}"
                )

            # determine type of input - logprobs or labels
            out_len = decoder_lengths if decoder_lengths is not None else None
            hypotheses = self.search_algorithm(prediction_tensor, out_len)

            # Pack results into Hypotheses
            packed_result = pack_hypotheses(hypotheses, decoder_lengths)

            # Pack the result
            if self.return_best_hypothesis and isinstance(packed_result[0], rnnt_utils.NBestHypotheses):
                packed_result = [res.n_best_hypotheses[0] for res in packed_result]  # type: Hypothesis

        return (packed_result,)

    @torch.no_grad()
    def default_beam_search(
        self, x: torch.Tensor, out_len: torch.Tensor
    ) -> List[Union[rnnt_utils.Hypothesis, rnnt_utils.NBestHypotheses]]:
        """
        Open Seq2Seq Beam Search Algorithm (DeepSpeed)

        Args:
            x: Tensor of shape [B, T, V+1], where B is the batch size, T is the maximum sequence length,
                and V is the vocabulary size. The tensor contains log-probabilities.
            out_len: Tensor of shape [B], contains lengths of each sequence in the batch.

        Returns:
            A list of NBestHypotheses objects, one for each sequence in the batch.
        """
        if self.compute_timestamps:
            raise ValueError(
                f"Beam Search with strategy `{self.search_type}` does not support time stamp calculation!"
            )

        if self.default_beam_scorer is None:
            # Check for filepath
            if self.ngram_lm_model is None or not os.path.exists(self.ngram_lm_model):
                raise FileNotFoundError(
                    f"KenLM binary file not found at : {self.ngram_lm_model}. "
                    f"Please set a valid path in the decoding config."
                )

            # perform token offset for subword models
            if self.decoding_type == 'subword':
                vocab = [chr(idx + self.token_offset) for idx in range(len(self.vocab))]
            else:
                # char models
                vocab = self.vocab

            # Must import at runtime to avoid circular dependency due to module level import.
            from nemo.collections.asr.modules.beam_search_decoder import BeamSearchDecoderWithLM

            self.default_beam_scorer = BeamSearchDecoderWithLM(
                vocab=vocab,
                lm_path=self.ngram_lm_model,
                beam_width=self.beam_size,
                alpha=self.ngram_lm_alpha,
                beta=self.beam_beta,
                num_cpus=max(1, os.cpu_count()),
                input_tensor=False,
            )

        x = x.to('cpu')

        with typecheck.disable_checks():
            data = [x[sample_id, : out_len[sample_id], :].softmax(dim=-1) for sample_id in range(len(x))]
            beams_batch = self.default_beam_scorer.forward(log_probs=data, log_probs_length=None)

        # For each sample in the batch
        nbest_hypotheses = []
        for beams_idx, beams in enumerate(beams_batch):
            # For each beam candidate / hypothesis in each sample
            hypotheses = []
            for candidate_idx, candidate in enumerate(beams):
                hypothesis = rnnt_utils.Hypothesis(
                    score=0.0, y_sequence=[], dec_state=None, timestamp=[], last_token=None
                )

                # For subword encoding, NeMo will double encode the subword (multiple tokens) into a
                # singular unicode id. In doing so, we preserve the semantic of the unicode token, and
                # compress the size of the final KenLM ARPA / Binary file.
                # In order to do double encoding, we shift the subword by some token offset.
                # This step is ignored for character based models.
                if self.decoding_type == 'subword':
                    pred_token_ids = [ord(c) - self.token_offset for c in candidate[1]]
                else:
                    # Char models
                    pred_token_ids = [self.vocab_index_map[c] for c in candidate[1]]

                # We preserve the token ids and the score for this hypothesis
                hypothesis.y_sequence = pred_token_ids
                hypothesis.score = candidate[0]

                # If alignment must be preserved, we preserve a view of the output logprobs.
                # Note this view is shared amongst all beams within the sample, be sure to clone it if you
                # require specific processing for each sample in the beam.
                # This is done to preserve memory.
                if self.preserve_alignments:
                    hypothesis.alignments = x[beams_idx][: out_len[beams_idx]]

                hypotheses.append(hypothesis)

            # Wrap the result in NBestHypothesis.
            hypotheses = rnnt_utils.NBestHypotheses(hypotheses)
            nbest_hypotheses.append(hypotheses)

        return nbest_hypotheses

    @torch.no_grad()
    def _pyctcdecode_beam_search(
        self, x: torch.Tensor, out_len: torch.Tensor
    ) -> List[Union[rnnt_utils.Hypothesis, rnnt_utils.NBestHypotheses]]:
        """
        PyCTCDecode Beam Search Algorithm. Should support Char and Subword models.

        Args:
            x: Tensor of shape [B, T, V+1], where B is the batch size, T is the maximum sequence length,
                and V is the vocabulary size. The tensor contains log-probabilities.
            out_len: Tensor of shape [B], contains lengths of each sequence in the batch.

        Returns:
            A list of NBestHypotheses objects, one for each sequence in the batch.
        """
        if self.compute_timestamps:
            raise ValueError(
                f"Beam Search with strategy `{self.search_type}` does not support time stamp calculation!"
            )

        try:
            import pyctcdecode
        except (ImportError, ModuleNotFoundError):
            raise ImportError(
                "Could not load `pyctcdecode` library. Please install it from pip using :\n"
                "pip install --upgrade pyctcdecode"
            )

        if self.pyctcdecode_beam_scorer is None:
            self.pyctcdecode_beam_scorer = pyctcdecode.build_ctcdecoder(
                labels=self.vocab, kenlm_model_path=self.ngram_lm_model, alpha=self.ngram_lm_alpha, beta=self.beam_beta
            )  # type: pyctcdecode.BeamSearchDecoderCTC

        x = x.to('cpu').numpy()

        with typecheck.disable_checks():
            beams_batch = []
            for sample_id in range(len(x)):
                logprobs = x[sample_id, : out_len[sample_id], :]
                result = self.pyctcdecode_beam_scorer.decode_beams(
                    logprobs,
                    beam_width=self.beam_size,
                    beam_prune_logp=self.pyctcdecode_cfg.beam_prune_logp,
                    token_min_logp=self.pyctcdecode_cfg.token_min_logp,
                    prune_history=self.pyctcdecode_cfg.prune_history,
                    hotwords=self.pyctcdecode_cfg.hotwords,
                    hotword_weight=self.pyctcdecode_cfg.hotword_weight,
                    lm_start_state=None,
                )  # Output format: text, last_lm_state, text_frames, logit_score, lm_score
                beams_batch.append(result)

        nbest_hypotheses = []
        for beams_idx, beams in enumerate(beams_batch):
            hypotheses = []
            for candidate_idx, candidate in enumerate(beams):
                # Candidate = (text, last_lm_state, text_frames, logit_score, lm_score)
                hypothesis = rnnt_utils.Hypothesis(
                    score=0.0, y_sequence=[], dec_state=None, timestamp=[], last_token=None
                )

                # TODO: Requires token ids to be returned rather than text.
                if self.decoding_type == 'subword':
                    if self.tokenizer is None:
                        raise ValueError("Tokenizer must be provided for subword decoding. Use set_tokenizer().")

                    pred_token_ids = self.tokenizer.text_to_ids(candidate[0])
                else:
                    if self.vocab is None:
                        raise ValueError("Vocab must be provided for character decoding. Use set_vocab().")

                    chars = list(candidate[0])
                    pred_token_ids = [self.vocab_index_map[c] for c in chars]

                hypothesis.y_sequence = pred_token_ids
                hypothesis.text = candidate[0]  # text
                hypothesis.score = candidate[4]  # score

                # Inject word level timestamps
                hypothesis.timestamp = candidate[2]  # text_frames

                if self.preserve_alignments:
                    hypothesis.alignments = torch.from_numpy(x[beams_idx][: out_len[beams_idx]])

                hypotheses.append(hypothesis)

            hypotheses = rnnt_utils.NBestHypotheses(hypotheses)
            nbest_hypotheses.append(hypotheses)

        return nbest_hypotheses

    @torch.no_grad()
    def flashlight_beam_search(
        self, x: torch.Tensor, out_len: torch.Tensor
    ) -> List[Union[rnnt_utils.Hypothesis, rnnt_utils.NBestHypotheses]]:
        """
        Flashlight Beam Search Algorithm. Should support Char and Subword models.

        Args:
            x: Tensor of shape [B, T, V+1], where B is the batch size, T is the maximum sequence length,
                and V is the vocabulary size. The tensor contains log-probabilities.
            out_len: Tensor of shape [B], contains lengths of each sequence in the batch.

        Returns:
            A list of NBestHypotheses objects, one for each sequence in the batch.
        """
        if self.compute_timestamps:
            raise ValueError(
                f"Beam Search with strategy `{self.search_type}` does not support time stamp calculation!"
            )

        if self.flashlight_beam_scorer is None:
            # Check for filepath
            if self.ngram_lm_model is None or not os.path.exists(self.ngram_lm_model):
                raise FileNotFoundError(
                    f"KenLM binary file not found at : {self.ngram_lm_model}. "
                    "Please set a valid path in the decoding config."
                )

            # perform token offset for subword models
            # if self.decoding_type == 'subword':
            #    vocab = [chr(idx + self.token_offset) for idx in range(len(self.vocab))]
            # else:
            #    # char models
            #    vocab = self.vocab

            # Must import at runtime to avoid circular dependency due to module level import.
            from nemo.collections.asr.modules.flashlight_decoder import FlashLightKenLMBeamSearchDecoder

            self.flashlight_beam_scorer = FlashLightKenLMBeamSearchDecoder(
                lm_path=self.ngram_lm_model,
                vocabulary=self.vocab,
                tokenizer=self.tokenizer,
                lexicon_path=self.flashlight_cfg.lexicon_path,
                boost_path=self.flashlight_cfg.boost_path,
                beam_size=self.beam_size,
                beam_size_token=self.flashlight_cfg.beam_size_token,
                beam_threshold=self.flashlight_cfg.beam_threshold,
                lm_weight=self.ngram_lm_alpha,
                word_score=self.beam_beta,
                unk_weight=self.flashlight_cfg.unk_weight,
                sil_weight=self.flashlight_cfg.sil_weight,
            )

        x = x.to('cpu')

        with typecheck.disable_checks():
            beams_batch = self.flashlight_beam_scorer.forward(log_probs=x)

        # For each sample in the batch
        nbest_hypotheses = []
        for beams_idx, beams in enumerate(beams_batch):
            # For each beam candidate / hypothesis in each sample
            hypotheses = []
            for candidate_idx, candidate in enumerate(beams):
                hypothesis = rnnt_utils.Hypothesis(
                    score=0.0, y_sequence=[], dec_state=None, timestamp=[], last_token=None
                )

                # We preserve the token ids and the score for this hypothesis
                hypothesis.y_sequence = candidate['tokens'].tolist()
                hypothesis.score = candidate['score']

                # If alignment must be preserved, we preserve a view of the output logprobs.
                # Note this view is shared amongst all beams within the sample, be sure to clone it if you
                # require specific processing for each sample in the beam.
                # This is done to preserve memory.
                if self.preserve_alignments:
                    hypothesis.alignments = x[beams_idx][: out_len[beams_idx]]

                hypotheses.append(hypothesis)

            # Wrap the result in NBestHypothesis.
            hypotheses = rnnt_utils.NBestHypotheses(hypotheses)
            nbest_hypotheses.append(hypotheses)

        return nbest_hypotheses

    def set_decoding_type(self, decoding_type: str):
        super().set_decoding_type(decoding_type)

        # Please check train_kenlm.py in scripts/asr_language_modeling/ to find out why we need
        # TOKEN_OFFSET for BPE-based models
        if self.decoding_type == 'subword':
            self.token_offset = DEFAULT_TOKEN_OFFSET


class WfstCTCInfer(AbstractBeamCTCInfer):
    """A WFST-based beam CTC decoder.

    Provides a common abstraction for sample level and batch level beam decoding.

    Args:
        TBD

    """

    def __init__(
        self,
        blank_id: int,
        beam_size: int,
        search_type: str = "riva",  # 'riva', 'k2'
        return_best_hypothesis: bool = True,
        preserve_alignments: bool = False,
        compute_timestamps: bool = False,
        decoding_mode: str = 'nbest',  # 'nbest', 'mbr' ('mbr' works only for search_type == 'riva' and beam_size == 1)
        open_vocabulary_decoding: bool = False,
        beam_width: float = 10.0,
        lm_weight: float = 1.0,
        device: str = "cuda",
        arpa_lm_path: str = None,
        wfst_lm_path: str = None,
        riva_decoding_cfg: Optional['RivaDecoderConfig'] = None,
        k2_decoding_cfg: Optional['GraphIntersectDenseConfig'] = None,
    ):
        super().__init__(blank_id=blank_id, beam_size=beam_size)

        self.search_type = search_type
        self.return_best_hypothesis = return_best_hypothesis
        self.preserve_alignments = preserve_alignments
        self.compute_timestamps = compute_timestamps

        self.decoding_algorithm = None
        if search_type in ("default", "riva"):
            self.decoding_algorithm = self._riva_decoding
        elif search_type == "k2":
            self.decoding_algorithm = self._k2_decoding

        # Log the WFST search_type
        logging.info(f"WFST beam search search_type: {search_type}")
        self.search_type = search_type

        if beam_size > 1 and decoding_mode != 'nbest':
            logging.warning(
                f"`beam_size` > 1 is supported only for `decoding_mode` == `nbest`\n"
                f"(provided: `{decoding_mode}`).\n"
                f"`beam_size` rewritten as 1"
            )
            self.beam_size = 1
        self.decoding_mode = decoding_mode

        self.open_vocabulary_decoding = open_vocabulary_decoding
        self._tokenword_disambig_id = -1
        self.beam_width = beam_width
        self.lm_weight = lm_weight
        self.device = device

        # Default beam search args
        self.arpa_lm_path = arpa_lm_path
        self.wfst_lm_path = wfst_lm_path

        self.riva_decoding_cfg = riva_decoding_cfg
        self.k2_decoding_cfg = k2_decoding_cfg

        # Default beam search scorer functions
        self.riva_decoder = None
        self.k2_decoder = None

    @typecheck()
    def forward(
        self,
        decoder_output: torch.Tensor,
        decoder_lengths: torch.Tensor,
    ) -> Tuple[List[Union[rnnt_utils.Hypothesis, rnnt_utils.NBestHypotheses]]]:
        """Returns a list of hypotheses given an input batch of the encoder hidden embedding.
        Output token is generated auto-repressively.

        Args:
            decoder_output: A tensor of size (batch, timesteps, features).
            decoder_lengths: list of int representing the length of each sequence
                output sequence.

        Returns:
            packed list containing batch number of sentences (Hypotheses).
        """
        if self.vocab is None:
            raise RuntimeError("Please set the vocabulary with `set_vocabulary()` before calling this function.")

        if self.decoding_type != 'subword':
            raise ValueError(
                f"`decoding_type` other than `subword` is not supported. Provided: `{self.decoding_type}`"
            )
        elif self.tokenizer is None:
            raise ValueError("Tokenizer must be provided for subword decoding. Use set_tokenizer().")
        if self.decoding_algorithm is None:
            raise NotImplementedError(
                f"The decoding search_type ({self.search_type}) supplied is not supported!\n"
                f"Please use one of : (default, riva, k2)"
            )

        with torch.no_grad(), torch.inference_mode():
            # Process each sequence independently
            prediction_tensor = decoder_output

            if prediction_tensor.ndim != 3:
                raise ValueError(
                    f"`decoder_output` must be a tensor of shape [B, T, V] (log probs, float). "
                    f"Provided shape = {prediction_tensor.shape}"
                )

            hypotheses = self.decoding_algorithm(prediction_tensor, decoder_lengths)

            # Pack results into Hypotheses
            packed_result = pack_wfst_hypotheses(hypotheses, prediction_tensor, decoder_lengths)

            # Pack the result
            if self.return_best_hypothesis and isinstance(packed_result[0], rnnt_utils.NBestHypotheses):
                packed_result = [res.n_best_hypotheses[0] for res in packed_result]  # type: Hypothesis

        return (packed_result,)

    def _prepare_decoding_lm_wfst(self) -> Union[str, 'kaldifst.StdFst', 'k2.Fsa']:  # noqa: F821
        """TBD"""
        arpa_lm_path_exists = self.arpa_lm_path is not None and os.path.exists(self.arpa_lm_path)
        wfst_lm_path_exists = self.wfst_lm_path is not None and os.path.exists(self.wfst_lm_path)
        lm_fst = None
        if wfst_lm_path_exists:
            if self.search_type == "riva" and not self.wfst_lm_path.endswith(".fst"):
                raise ValueError(
                    f"Search type `riva` expects WFSTs in the `.fst` format. Provided: `{self.wfst_lm_path}`"
                )
            if self.search_type == "k2" and not self.wfst_lm_path.endswith(".pt"):
                raise ValueError(
                    f"Search type `k2` expects WFSTs in the `.pt` format. Provided: `{self.wfst_lm_path}`"
                )
            if arpa_lm_path_exists:
                logging.warning(
                    "Both `arpa_lm_path` and `wfst_lm_path` are provided and not empty. The latter will be used."
                )
            lm_fst = self.wfst_lm_path
        elif not arpa_lm_path_exists:
            raise FileNotFoundError(
                f"Arpa LM file not found at `{self.arpa_lm_path}` and WFST LM is not found at `{self.wfst_lm_path}`.\n"
                f"Please set a valid path in the decoding config for at least one of those."
            )
        else:
            logging.warning(
                f"Since WFST LM is not found at `{self.wfst_lm_path}`, "
                f"it will be made from the Arpa LM at `{self.arpa_lm_path}`.\n"
                f"This procedure will take some time."
            )
            if self.wfst_lm_path is not None:
                logging.info(f"WFST LM will be buffered at `{self.wfst_lm_path}`.")
                write_tlg_path = self.wfst_lm_path
            else:
                logging.warning("Consider providing a write-permitted `wfst_lm_path` for WFST LM buffering.")
                write_tlg_path = None
            ctc_topology = "default"  # there is no way to indicate the need of other topologies
            target = "kaldi" if self.search_type == "riva" else "k2"

            from nemo.collections.asr.parts.utils.wfst_utils import mkgraph_ctc_ov

            lm_fst, tokenword_disambig_id = mkgraph_ctc_ov(
                tokenizer=self.tokenizer,
                lm_path=self.arpa_lm_path,
                topology_name=ctc_topology,
                write_tlg_path=write_tlg_path,
                open_vocabulary=self.open_vocabulary_decoding,
                target=target,
            )
            self._tokenword_disambig_id = tokenword_disambig_id

        return lm_fst

    @torch.no_grad()
    def _riva_decoding(self, x: torch.Tensor, out_len: torch.Tensor) -> List['WfstNbestHypothesis']:
        """
        Riva Asrlib WFST decoder Algorithm.

        Args:
            x: Tensor of shape [B, T, V+1], where B is the batch size, T is the maximum sequence length,
                and V is the vocabulary size. The tensor contains log-probabilities.
            out_len: Tensor of shape [B], contains lengths of each sequence in the batch.

        Returns:
            A list of WfstNbestHypothesis objects, one for each sequence in the batch.
        """
        if self.riva_decoder is None:
            lm_fst = self._prepare_decoding_lm_wfst()
            if self.open_vocabulary_decoding and self._tokenword_disambig_id == -1:
                # trying to extract tokenword_disambig_id from the lm_fst
                if isinstance(lm_fst, str):
                    # use importer instead of direct import to possibly get an installation message
                    from nemo.collections.asr.parts.utils.wfst_utils import kaldifst_importer

                    kaldifst = kaldifst_importer()
                    lm_fst = kaldifst.StdVectorFst.read(self.wfst_lm_path)
                tokenword_disambig_id = lm_fst.output_symbols.find("#1")
                if tokenword_disambig_id == -1:
                    raise ValueError(
                        "Cannot determine `tokenword_disambig_id` "
                        "which is required if `open_vocabulary_decoding` == True"
                    )
                self._tokenword_disambig_id = tokenword_disambig_id
            if not self.device.startswith("cuda"):
                raise ValueError(f"Riva decoder does not support non-cuda device. Provided: `{self.device}`")

            from nemo.collections.asr.parts.submodules.wfst_decoder import RivaGpuWfstDecoder

            self.riva_decoder = RivaGpuWfstDecoder(
                lm_fst=lm_fst,
                decoding_mode=self.decoding_mode,
                beam_size=self.beam_width,
                config=self.riva_decoding_cfg,
                tokenword_disambig_id=self._tokenword_disambig_id,
                lm_weight=self.lm_weight,
                nbest_size=self.beam_size,
            )

        return self.riva_decoder.decode(x.to(device=self.device), out_len.to(device=self.device))

    @torch.no_grad()
    def _k2_decoding(self, x: torch.Tensor, out_len: torch.Tensor) -> List['WfstNbestHypothesis']:
        """
        K2 WFST decoder Algorithm.

        Args:
            x: Tensor of shape [B, T, V+1], where B is the batch size, T is the maximum sequence length,
                and V is the vocabulary size. The tensor contains log-probabilities.
            out_len: Tensor of shape [B], contains lengths of each sequence in the batch.

        Returns:
            A list of WfstNbestHypothesis objects, one for each sequence in the batch.
        """
        if self.k2_decoder is None:
            lm_fst = self._prepare_decoding_lm_wfst()
            if self.open_vocabulary_decoding and self._tokenword_disambig_id == -1:
                if isinstance(lm_fst, str):
                    from nemo.collections.asr.parts.k2.utils import load_graph

                    with torch.inference_mode(False):
                        lm_fst = load_graph(lm_fst)
                try:
                    tokenword_disambig_id = lm_fst.aux_labels_sym.get("#1")
                    self._tokenword_disambig_id = tokenword_disambig_id
                except KeyError:
                    raise ValueError(
                        "Cannot determine `tokenword_disambig_id` "
                        "which is required if `open_vocabulary_decoding` == True"
                    )

            from nemo.collections.asr.parts.k2.graph_decoders import K2WfstDecoder

            self.k2_decoder = K2WfstDecoder(
                lm_fst=lm_fst,
                decoding_mode=self.decoding_mode,
                beam_size=self.beam_width,
                config=self.k2_decoding_cfg,
                tokenword_disambig_id=self._tokenword_disambig_id,
                lm_weight=self.lm_weight,
                nbest_size=self.beam_size,
                device=self.device,
            )

        return self.k2_decoder.decode(x.to(device=self.device), out_len.to(device=self.device))


class BeamBatchedCTCInfer(AbstractBeamCTCInfer, WithOptionalCudaGraphs):
    """
    A batched beam CTC decoder.

    Args:
        blank_index: int index of the blank token. Can be 0 or len(vocabulary).
        beam_size: int size of the beam.
        return_best_hypothesis:  When set to True (default), returns a single Hypothesis.
            When set to False, returns a NBestHypotheses container, which contains a list of Hypothesis.
        preserve_alignments: Bool flag which preserves the history of logprobs generated during
            decoding (sample / batched). When set to true, the Hypothesis will contain
            the non-null value for `logprobs` in it. Here, `logprobs` is a torch.Tensors.
        compute_timestamps: A bool flag, which determines whether to compute the character/subword, or
                word based timestamp mapping the output log-probabilities to discrite intervals of timestamps.
                The timestamps will be available in the returned Hypothesis.timestep as a dictionary.
        ngram_lm_alpha: float, the language model weight.
        beam_beta: float, the word insertion weight.
        beam_threshold: float, the beam pruning threshold.
        ngram_lm_model: str, the path to the ngram model.
        boosting_tree: BoostingTreeModelConfig, the boosting tree model config.
        boosting_tree_alpha: float, the boosting tree alpha.
        allow_cuda_graphs: bool, whether to allow cuda graphs for the beam search algorithm.
    """

    def __init__(
        self,
        blank_index: int,
        beam_size: int,
        return_best_hypothesis: bool = True,
        preserve_alignments: bool = False,
        compute_timestamps: bool = False,
        ngram_lm_alpha: float = 1.0,
        beam_beta: float = 0.0,
        beam_threshold: float = 20.0,
        ngram_lm_model: str = None,
        boosting_tree: BoostingTreeModelConfig = None,
        boosting_tree_alpha: float = 0.0,
        allow_cuda_graphs: bool = True,
        tokenizer: TokenizerSpec = None,
    ):
        super().__init__(blank_id=blank_index, beam_size=beam_size)

        self.return_best_hypothesis = return_best_hypothesis
        self.preserve_alignments = preserve_alignments
        self.compute_timestamps = compute_timestamps
        self.allow_cuda_graphs = allow_cuda_graphs

        if self.compute_timestamps:
            raise ValueError("`Compute timestamps` is not supported for batched beam search.")
        if self.preserve_alignments:
            raise ValueError("`Preserve alignments` is not supported for batched beam search.")

        self.beam_beta = beam_beta
        self.beam_threshold = beam_threshold

        # load fusion models from paths (ngram_lm_model and boosting_tree_model)
        fusion_models, fusion_models_alpha = [], []
        if ngram_lm_model is not None:
            assert blank_index != 0, "Blank should not be the first token in the vocabulary"
            fusion_models.append(NGramGPULanguageModel.from_file(lm_path=ngram_lm_model, vocab_size=blank_index))
            fusion_models_alpha.append(ngram_lm_alpha)
        if boosting_tree and not BoostingTreeModelConfig.is_empty(boosting_tree):
            assert blank_index != 0, "Blank should not be the first token in the vocabulary"
            fusion_models.append(GPUBoostingTreeModel.from_config(boosting_tree, tokenizer=tokenizer))
            fusion_models_alpha.append(boosting_tree_alpha)
        if not fusion_models:
            fusion_models = None
            fusion_models_alpha = None

        # # Default beam search args

        self.search_algorithm = BatchedBeamCTCComputer(
            blank_index=blank_index,
            beam_size=beam_size,
            return_best_hypothesis=return_best_hypothesis,
            preserve_alignments=preserve_alignments,
            compute_timestamps=compute_timestamps,
            fusion_models=fusion_models,
            fusion_models_alpha=fusion_models_alpha,
            beam_beta=beam_beta,
            beam_threshold=beam_threshold,
            allow_cuda_graphs=allow_cuda_graphs,
        )

    def disable_cuda_graphs(self) -> bool:
        """Disable CUDA graphs (e.g., for decoding in training)"""
        if isinstance(self.search_algorithm, WithOptionalCudaGraphs):
            return self.search_algorithm.disable_cuda_graphs()
        return False

    def maybe_enable_cuda_graphs(self) -> bool:
        """Enable CUDA graphs (if allowed)"""
        if isinstance(self.search_algorithm, WithOptionalCudaGraphs):
            return self.search_algorithm.maybe_enable_cuda_graphs()
        return False

    @typecheck()
    def forward(
        self,
        decoder_output: torch.Tensor,
        decoder_lengths: torch.Tensor,
    ) -> Tuple[List[Union[rnnt_utils.Hypothesis, rnnt_utils.NBestHypotheses]]]:
        """Returns a list of hypotheses given an input batch of the encoder hidden embedding.
        Output token is generated auto-repressively.

        Args:
            decoder_output: A tensor of size (batch, timesteps, features).
            decoder_lengths: list of int representing the length of each sequence
                output sequence.

        Returns:
            packed list containing batch number of sentences (Hypotheses).
        """
        with torch.no_grad(), torch.inference_mode():
            if decoder_output.ndim != 3:
                raise ValueError(
                    f"`decoder_output` must be a tensor of shape [B, T, V] (log probs, float). "
                    f"Provided shape = {decoder_output.shape}"
                )

            batched_beam_hyps = self.search_algorithm(decoder_output, decoder_lengths)

            batch_size = decoder_lengths.shape[0]
            if self.return_best_hypothesis:
                hyps = batched_beam_hyps.to_hyps_list(score_norm=False)[:batch_size]
            else:
                hyps = batched_beam_hyps.to_nbest_hyps_list(score_norm=False)[:batch_size]

        return (hyps,)


@dataclass
class PyCTCDecodeConfig:
    # These arguments cannot be imported from pyctcdecode (optional dependency)
    # Therefore we copy the values explicitly
    # Taken from pyctcdecode.constant
    beam_prune_logp: float = -10.0
    token_min_logp: float = -5.0
    prune_history: bool = False
    hotwords: Optional[List[str]] = None
    hotword_weight: float = 10.0


@dataclass
class FlashlightConfig:
    lexicon_path: Optional[str] = None
    boost_path: Optional[str] = None
    beam_size_token: int = 16
    beam_threshold: float = 20.0
    unk_weight: float = -math.inf
    sil_weight: float = 0.0


@dataclass
class BeamCTCInferConfig:
    beam_size: int
    search_type: str = 'default'
    preserve_alignments: bool = False
    compute_timestamps: bool = False
    return_best_hypothesis: bool = True
    allow_cuda_graphs: bool = True

    beam_alpha: Optional[float] = None  # Deprecated
    beam_beta: float = 1.0
    beam_threshold: float = 20.0
    kenlm_path: Optional[str] = None  # Deprecated, default should be None
    ngram_lm_alpha: Optional[float] = 1.0
    ngram_lm_model: Optional[str] = None
    boosting_tree: BoostingTreeModelConfig = field(default_factory=BoostingTreeModelConfig)
    boosting_tree_alpha: Optional[float] = 0.0

    flashlight_cfg: Optional[FlashlightConfig] = field(default_factory=lambda: FlashlightConfig())
    pyctcdecode_cfg: Optional[PyCTCDecodeConfig] = field(default_factory=lambda: PyCTCDecodeConfig())


@dataclass
class WfstCTCInferConfig:
    beam_size: int
    search_type: str = "riva"  # 'riva', 'k2'
    return_best_hypothesis: bool = True
    preserve_alignments: bool = False
    compute_timestamps: bool = False
    decoding_mode: str = 'nbest'  # 'nbest', 'mbr' ('mbr' works only for search_type == 'riva' and beam_size == 1)
    open_vocabulary_decoding: bool = False
    beam_width: float = 10.0
    lm_weight: float = 1.0
    device: str = "cuda"
    arpa_lm_path: Optional[str] = None
    wfst_lm_path: Optional[str] = None
    riva_decoding_cfg: Optional['RivaDecoderConfig'] = field(default_factory=lambda: RivaDecoderConfig())
    k2_decoding_cfg: Optional['GraphIntersectDenseConfig'] = field(default_factory=lambda: GraphIntersectDenseConfig())
