# Copyright (c) 2025, 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.

# Copyright 2017 Johns Hopkins University (Shinji Watanabe)
#
# 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 typing import List, Optional, Tuple

import numpy as np
import torch
from tqdm import tqdm

from nemo.collections.asr.modules import rnnt_abstract
from nemo.collections.asr.parts.submodules.ngram_lm import NGramGPULanguageModel
from nemo.collections.asr.parts.submodules.rnnt_beam_decoding import pack_hypotheses
from nemo.collections.asr.parts.submodules.tdt_malsd_batched_computer import ModifiedALSDBatchedTDTComputer
from nemo.collections.asr.parts.utils.asr_confidence_utils import ConfidenceMethodMixin
from nemo.collections.asr.parts.utils.batched_beam_decoding_utils import BlankLMScoreMode, PruningMode
from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis, NBestHypotheses, is_prefix
from nemo.collections.common.parts.optional_cuda_graphs import WithOptionalCudaGraphs
from nemo.core.classes import Typing, typecheck
from nemo.core.neural_types import AcousticEncodedRepresentation, HypothesisType, LengthsType, NeuralType
from nemo.utils import logging

try:
    import kenlm

    KENLM_AVAILABLE = True
except (ImportError, ModuleNotFoundError):
    KENLM_AVAILABLE = False


class BeamTDTInfer(Typing):
    """
    Beam search implementation for Token-andDuration Transducer (TDT) models.

    Sequence level beam decoding or batched-beam decoding, performed auto-repressively
    depending on the search type chosen.

    Args:
        decoder_model: rnnt_utils.AbstractRNNTDecoder implementation.
        joint_model: rnnt_utils.AbstractRNNTJoint implementation.
        durations: list of duration values from TDT model.

        beam_size: number of beams for beam search. Must be a positive integer >= 1.
            If beam size is 1, defaults to stateful greedy search.
            For accurate greedy results, please use GreedyRNNTInfer or GreedyBatchedRNNTInfer.

        search_type: str representing the type of beam search to perform.
            Must be one of ['beam', 'maes'].

            Algorithm used:

                `default` - basic beam search strategy. Larger beams generally result in better decoding,
                    however the time required for the search also grows steadily.

                `maes` = modified adaptive expansion search. Please refer to the paper:
                    [Accelerating RNN Transducer Inference via Adaptive Expansion Search]
                    (https://ieeexplore.ieee.org/document/9250505)

                    Modified Adaptive Synchronous Decoding (mAES) execution time is adaptive w.r.t the
                    number of expansions (for tokens) required per timestep. The number of expansions can usually
                    be constrained to 1 or 2, and in most cases 2 is sufficient.

                    This beam search technique can possibly obtain superior WER while sacrificing some evaluation time.

        score_norm: bool, whether to normalize the scores of the log probabilities.

        return_best_hypothesis: bool, decides whether to return a single hypothesis (the best out of N),
            or return all N hypothesis (sorted with best score first). The container class changes based
            this flag -
            When set to True (default), returns a single Hypothesis.
            When set to False, returns a NBestHypotheses container, which contains a list of Hypothesis.

        # The following arguments are specific to the chosen `search_type`

        # mAES flags
        maes_num_steps: Number of adaptive steps to take. From the paper, 2 steps is generally sufficient. int > 1.

        maes_prefix_alpha: Maximum prefix length in prefix search. Must be an integer, and is advised to keep this as 1
            in order to reduce expensive beam search cost later. int >= 0.

        maes_expansion_beta: Maximum number of prefix expansions allowed, in addition to the beam size.
            Effectively, the number of hypothesis = beam_size + maes_expansion_beta. Must be an int >= 0,
            and affects the speed of inference since large values will perform large beam search in the next step.

        maes_expansion_gamma: Float pruning threshold used in the prune-by-value step when computing the expansions.
            The default (2.3) is selected from the paper. It performs a comparison
            (max_log_prob - gamma <= log_prob[v]) where v is all vocabulary indices in the Vocab set and max_log_prob
            is the "most" likely token to be predicted. Gamma therefore provides a margin of additional tokens which
            can be potential candidates for expansion apart from the "most likely" candidate.
            Lower values will reduce the number of expansions (by increasing pruning-by-value, thereby improving speed
            but hurting accuracy). Higher values will increase the number of expansions (by reducing pruning-by-value,
            thereby reducing speed but potentially improving accuracy). This is a hyper parameter to be experimentally
            tuned on a validation set.

        softmax_temperature: Scales the logits of the joint prior to computing log_softmax.

        preserve_alignments: Bool flag which preserves the history of alignments generated during
            beam decoding (sample). When set to true, the Hypothesis will contain
            the non-null value for `alignments` in it. Here, `alignments` is a List of List of Tensor (of length V + 1)

            The length of the list corresponds to the Acoustic Length (T).
            Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary.
            U is the number of target tokens for the current timestep Ti.

            NOTE: `preserve_alignments` is an invalid argument for any `search_type`
            other than basic beam search.

        ngram_lm_model: str
            The path to the N-gram LM.
        ngram_lm_alpha: float
            Alpha weight of N-gram LM.
    """

    @property
    def input_types(self):
        """Returns definitions of module input ports."""
        return {
            "encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
            "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
            "partial_hypotheses": [NeuralType(elements_type=HypothesisType(), optional=True)],  # must always be last
        }

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

    def __init__(
        self,
        decoder_model: rnnt_abstract.AbstractRNNTDecoder,
        joint_model: rnnt_abstract.AbstractRNNTJoint,
        durations: list,
        beam_size: int,
        search_type: str = 'default',
        score_norm: bool = True,
        return_best_hypothesis: bool = True,
        maes_num_steps: int = 2,
        maes_prefix_alpha: int = 1,
        maes_expansion_gamma: float = 2.3,
        maes_expansion_beta: int = 2,
        softmax_temperature: float = 1.0,
        preserve_alignments: bool = False,
        ngram_lm_model: Optional[str] = None,
        ngram_lm_alpha: float = 0.3,
        max_symbols_per_step: Optional[int] = None,
        blank_lm_score_mode: Optional[str] = "no_score",
        pruning_mode: Optional[str] = "early",
        allow_cuda_graphs: bool = False,
    ):
        self.joint = joint_model
        self.decoder = decoder_model
        self.durations = durations

        self.token_offset = 0
        self.search_type = search_type
        self.blank = decoder_model.blank_idx
        self.vocab_size = decoder_model.vocab_size
        self.return_best_hypothesis = return_best_hypothesis

        self.beam_size = beam_size
        self.score_norm = score_norm
        self.max_candidates = beam_size
        self.softmax_temperature = softmax_temperature
        self.preserve_alignments = preserve_alignments

        if preserve_alignments:
            raise ValueError("Alignment preservation has not been implemented.")
        if beam_size < 1:
            raise ValueError("Beam search size cannot be less than 1!")

        if self.preserve_alignments:
            raise NotImplementedError("Preserving alignments is not implemented.")

        if search_type == "default":
            if self.beam_size == 1:
                logging.info(
                    """If beam size is 1, defaults to stateful greedy search.
                     For accurate greedy results, please use GreedyTDTInfer or GreedyBatchedTDTInfer."""
                )
            self.search_algorithm = self.default_beam_search
        elif search_type == "tsd":
            raise NotImplementedError("`tsd` (Time Synchronous Decoding) has not been implemented.")
        elif search_type == "alsd":
            raise NotImplementedError("`alsd` (Alignment Length Synchronous Decoding) has not been implemented.")
        elif search_type == "nsc":
            raise NotImplementedError("`nsc` (Constrained Beam Search) has not been implemented.")
        elif search_type == "maes":
            self.search_algorithm = self.modified_adaptive_expansion_search
        else:
            raise NotImplementedError(
                f"The search type ({search_type}) supplied is not supported!\n" f"Please use one of : (default, maes)"
            )

        if max_symbols_per_step is not None:
            logging.warning(
                f"Not supported parameter `max_symbols_per_step` for decoding strategy {self.search_algorithm }"
            )

        if allow_cuda_graphs:
            logging.warning(
                f"""Cuda Graphs are not supported for the decoding strategy {self.search_algorithm}.
                                Decoding will proceed without Cuda Graphs."""
            )

        strategies = ["default", "maes"]
        strategies_batch = ["malsd_batch"]
        if (pruning_mode, blank_lm_score_mode) != ("early", "no_score"):
            logging.warning(
                f"""Decoding strategies {strategies} support early pruning and the 'no_score' blank scoring mode.
                Please choose a strategy from {strategies_batch} for {pruning_mode} pruning 
                and {blank_lm_score_mode} blank scoring mode." 
                """
            )

        if self.search_type == 'maes':
            self.maes_num_steps = int(maes_num_steps)
            self.maes_prefix_alpha = int(maes_prefix_alpha)
            self.maes_expansion_beta = int(maes_expansion_beta)
            self.maes_expansion_gamma = float(maes_expansion_gamma)

            self.max_candidates += maes_expansion_beta

            if self.maes_prefix_alpha < 0:
                raise ValueError("`maes_prefix_alpha` must be a positive integer.")

            if self.vocab_size < beam_size + maes_expansion_beta:
                raise ValueError(
                    f"beam_size ({beam_size}) + expansion_beta ({maes_expansion_beta}) "
                    f"should be smaller or equal to vocabulary size ({self.vocab_size})."
                )

            if self.maes_num_steps < 1:
                raise ValueError("`maes_num_steps` must be greater than 0.")

        try:
            self.zero_duration_idx = self.durations.index(0)
        except ValueError:
            self.zero_duration_idx = None
        self.min_non_zero_duration_idx = int(
            np.argmin(np.ma.masked_where(np.array(self.durations) == 0, self.durations))
        )

        if ngram_lm_model:
            if search_type != "maes":
                raise ValueError("For decoding with language model `maes` decoding strategy must be chosen.")
            self.ngram_lm = ngram_lm_model
            self.ngram_lm_alpha = ngram_lm_alpha
        else:
            self.ngram_lm = None

    @typecheck()
    def __call__(
        self,
        encoder_output: torch.Tensor,
        encoded_lengths: torch.Tensor,
        partial_hypotheses: tuple[list[Hypothesis | NBestHypotheses],] = None,
    ) -> tuple[list[Hypothesis | NBestHypotheses],]:
        """Perform general beam search.

        Args:
            encoder_output: encoder outputs (batch, features, timesteps).
            encoded_lengths: lengths of the encoder outputs.

        Returns:
            Either a list containing a single Hypothesis (when `return_best_hypothesis=True`,
            otherwise a list containing a single NBestHypotheses, which itself contains a list of
            Hypothesis. This list is sorted such that the best hypothesis is the first element.
        """
        # Preserve decoder and joint training state
        decoder_training_state = self.decoder.training
        joint_training_state = self.joint.training

        with torch.inference_mode():
            # Apply optional preprocessing
            encoder_output = encoder_output.transpose(1, 2)  # (B, T, D)

            self.decoder.eval()
            self.joint.eval()

            hypotheses = []
            with tqdm(
                range(encoder_output.size(0)),
                desc='Beam search progress:',
                total=encoder_output.size(0),
                unit='sample',
            ) as idx_gen:

                _p = next(self.joint.parameters())
                dtype = _p.dtype

                # Decode every sample in the batch independently.
                for batch_idx in idx_gen:
                    inseq = encoder_output[batch_idx : batch_idx + 1, : encoded_lengths[batch_idx], :]  # [1, T, D]
                    logitlen = encoded_lengths[batch_idx]

                    if inseq.dtype != dtype:
                        inseq = inseq.to(dtype=dtype)

                    # Extract partial hypothesis if exists
                    partial_hypothesis = partial_hypotheses[batch_idx] if partial_hypotheses is not None else None

                    # Execute the specific search strategy
                    nbest_hyps = self.search_algorithm(
                        inseq, logitlen, partial_hypotheses=partial_hypothesis
                    )  # sorted list of hypothesis

                    # Prepare the list of hypotheses
                    nbest_hyps = pack_hypotheses(nbest_hyps)

                    # Pack the result
                    if self.return_best_hypothesis:
                        best_hypothesis: Hypothesis = nbest_hyps[0]
                    else:
                        best_hypothesis: NBestHypotheses = NBestHypotheses(nbest_hyps)
                    hypotheses.append(best_hypothesis)

        self.decoder.train(decoder_training_state)
        self.joint.train(joint_training_state)

        return (hypotheses,)

    def default_beam_search(
        self,
        encoder_outputs: torch.Tensor,
        encoded_lengths: torch.Tensor,
        partial_hypotheses: Optional[Hypothesis] = None,
    ) -> List[Hypothesis]:
        """Default Beam search implementation for TDT models.

        Args:
            encoder_outputs: encoder outputs (batch, features, timesteps).
            encoded_lengths: lengths of the encoder outputs.
            partial_hypotheses: partial hypoteses.

        Returns:
            nbest_hyps: N-best decoding results
        """
        if partial_hypotheses is not None:
            raise NotImplementedError("Support for `partial_hypotheses` is not implemented.")

        beam = min(self.beam_size, self.vocab_size)
        beam_k = min(beam, (self.vocab_size - 1))
        durations_beam_k = min(beam, len(self.durations))

        # Initialize zero vector states.
        decoder_state = self.decoder.initialize_state(encoder_outputs)
        # Cache decoder results to avoid duplicate computations.
        cache = {}

        # Initialize hypothesis array with blank hypothesis.
        start_hyp = Hypothesis(
            score=0.0, y_sequence=[self.blank], dec_state=decoder_state, timestamp=[-1], length=0, last_frame=0
        )
        kept_hyps = [start_hyp]

        for time_idx in range(int(encoded_lengths)):
            # Retrieve hypotheses for current and future frames
            hyps = [hyp for hyp in kept_hyps if hyp.last_frame == time_idx]  # hypotheses for current frame
            kept_hyps = [hyp for hyp in kept_hyps if hyp.last_frame > time_idx]  # hypothesis for future frames

            # Loop over hypotheses of current frame
            while len(hyps) > 0:
                max_hyp = max(hyps, key=lambda x: x.score)
                hyps.remove(max_hyp)

                # Update decoder state and get probability distribution over vocabulary and durations.
                encoder_output = encoder_outputs[:, time_idx : time_idx + 1, :]  # [1, 1, D]
                decoder_output, decoder_state, _ = self.decoder.score_hypothesis(max_hyp, cache)  # [1, 1, D]
                logits = (
                    self.joint.joint(encoder_output, decoder_output) / self.softmax_temperature
                )  # [1, 1, 1, V + NUM_DURATIONS + 1]
                logp = torch.log_softmax(logits[0, 0, 0, : -len(self.durations)], dim=-1)  # [V + 1]
                durations_logp = torch.log_softmax(logits[0, 0, 0, -len(self.durations) :], dim=-1)  # [NUM_DURATIONS]

                # Proccess non-blank tokens
                # Retrieve the top `beam_k` most probable tokens and the top `duration_beam_k` most probable durations.
                # Then, select the top `beam_k` pairs of (token, duration) based on the highest combined probabilities.
                # Note that indices are obtained in the flattened array.
                logp_topks, logp_topk_idxs = logp[:-1].topk(beam_k, dim=-1)  # topk of tokens without blank token
                durations_logp_topks, durations_logp_topk_idxs = durations_logp.topk(durations_beam_k, dim=-1)
                total_logp_topks, total_logp_topk_idxs = (
                    torch.cartesian_prod(durations_logp_topks, logp_topks).sum(dim=-1).topk(beam_k, dim=-1)
                )

                # Loop over pairs of (token, duration) with highest combined log prob
                for total_logp_topk, total_logp_topk_idx in zip(total_logp_topks, total_logp_topk_idxs):
                    # Restore indices from flattened array indices
                    token_idx = int(logp_topk_idxs[total_logp_topk_idx % beam_k])
                    duration_idx = int(durations_logp_topk_idxs[total_logp_topk_idx // beam_k])

                    duration = self.durations[duration_idx]
                    # Construct hypothesis for non-blank token
                    new_hyp = Hypothesis(
                        score=float(max_hyp.score + total_logp_topk),  # update score
                        y_sequence=max_hyp.y_sequence + [token_idx],  # update hypothesis sequence
                        dec_state=decoder_state,  # update decoder state
                        timestamp=max_hyp.timestamp + [time_idx + duration],  # update timesteps
                        length=encoded_lengths,
                        last_frame=max_hyp.last_frame + duration,
                    )  # update frame idx where last token appeared

                    # Update current frame hypotheses if duration is zero and future frame hypotheses otherwise
                    if duration == 0:
                        hyps.append(new_hyp)
                    else:
                        kept_hyps.append(new_hyp)

                # Update future frames with blank tokens
                # Note: blank token can have only non-zero duration
                for duration_idx in durations_logp_topk_idxs:
                    duration_idx = int(duration_idx)
                    # If zero is the only duration in topk, switch to closest non-zero duration to continue
                    if duration_idx == self.zero_duration_idx:
                        if durations_logp_topk_idxs.shape[0] == 1:
                            duration_idx = self.min_non_zero_duration_idx
                        else:
                            continue

                    duration = self.durations[duration_idx]
                    new_hyp = Hypothesis(
                        score=float(max_hyp.score + logp[self.blank] + durations_logp[duration_idx]),  # update score
                        y_sequence=max_hyp.y_sequence[:],  # no need to update sequence
                        dec_state=max_hyp.dec_state,  # no need to update decoder state
                        timestamp=max_hyp.timestamp[:],  # no need to update timesteps
                        length=encoded_lengths,
                        last_frame=max_hyp.last_frame + duration,
                    )  # update frame idx where last token appeared
                    kept_hyps.append(new_hyp)

                # Merge duplicate hypotheses.
                # If two consecutive blank tokens are predicted and their duration values sum up to the same number,
                # it will produce two hypotheses with the same token sequence but different scores.
                kept_hyps = self.merge_duplicate_hypotheses(kept_hyps)

                if len(hyps) > 0:
                    # Keep those hypothesis that have scores greater than next search generation
                    hyps_max = float(max(hyps, key=lambda x: x.score).score)
                    kept_most_prob = sorted(
                        [hyp for hyp in kept_hyps if hyp.score > hyps_max],
                        key=lambda x: x.score,
                    )
                    # If enough hypotheses have scores greater than next search generation,
                    # stop beam search.
                    if len(kept_most_prob) >= beam:
                        kept_hyps = kept_most_prob
                        break
                else:
                    # If there are no hypotheses in a current frame,
                    # keep only `beam` best hypotheses for the next search generation.
                    kept_hyps = sorted(kept_hyps, key=lambda x: x.score, reverse=True)[:beam]
        return self.sort_nbest(kept_hyps)

    def modified_adaptive_expansion_search(
        self,
        encoder_outputs: torch.Tensor,
        encoded_lengths: torch.Tensor,
        partial_hypotheses: Optional[Hypothesis] = None,
    ) -> List[Hypothesis]:
        """
        Modified Adaptive Exoansion Search algorithm for TDT models.
        Based on/modified from https://ieeexplore.ieee.org/document/9250505.
        Supports N-gram language model shallow fusion.

        Args:
            encoder_outputs: encoder outputs (batch, features, timesteps).
            encoded_lengths: lengths of the encoder outputs.
            partial_hypotheses: partial hypotheses.

        Returns:
            nbest_hyps: N-best decoding results
        """
        if partial_hypotheses is not None:
            raise NotImplementedError("Support for `partial_hypotheses` is not implemented.")

        beam = min(self.beam_size, self.vocab_size)
        beam_state = self.decoder.initialize_state(
            torch.zeros(1, device=encoder_outputs.device, dtype=encoder_outputs.dtype)
        )  # [L, B, H], [L, B, H] for LSTMS

        # Initialize first hypothesis for the beam (blank).
        start_hyp = Hypothesis(
            y_sequence=[self.blank],
            score=0.0,
            dec_state=self.decoder.batch_select_state(beam_state, 0),
            timestamp=[-1],
            length=0,
            last_frame=0,
        )
        init_tokens = [start_hyp]

        # Cache decoder results to avoid duplicate computations.
        cache = {}

        # Decode a batch of beam states and scores
        beam_decoder_output, beam_state = self.decoder.batch_score_hypothesis(init_tokens, cache)
        state = beam_state[0]

        # Initialize first hypothesis for the beam (blank) for kept hypotheses
        start_hyp_kept = Hypothesis(
            y_sequence=[self.blank],
            score=0.0,
            dec_state=state,
            dec_out=[beam_decoder_output[0]],
            timestamp=[-1],
            length=0,
            last_frame=0,
        )

        kept_hyps = [start_hyp_kept]

        # Setup ngram LM:
        if self.ngram_lm:
            init_lm_state = kenlm.State()
            self.ngram_lm.BeginSentenceWrite(init_lm_state)
            start_hyp_kept.ngram_lm_state = init_lm_state

        for time_idx in range(encoded_lengths):
            # Select current iteration hypotheses
            hyps = [x for x in kept_hyps if x.last_frame == time_idx]
            kept_hyps = [x for x in kept_hyps if x.last_frame > time_idx]

            if len(hyps) == 0:
                continue

            beam_encoder_output = encoder_outputs[:, time_idx : time_idx + 1]  # [1, 1, D]
            # Perform prefix search to update hypothesis scores.
            if self.zero_duration_idx is not None:
                hyps = self.prefix_search(
                    sorted(hyps, key=lambda x: len(x.y_sequence), reverse=True),
                    beam_encoder_output,
                    prefix_alpha=self.maes_prefix_alpha,
                )

            list_b = []  # List that contains the blank token emissions
            list_nb = []  # List that contains the non-zero duration non-blank token emissions
            # Repeat for number of mAES steps
            for n in range(self.maes_num_steps):
                # Pack the decoder logits for all current hypotheses
                beam_decoder_output = torch.stack([h.dec_out[-1] for h in hyps])  # [H, 1, D]

                # Extract the log probabilities
                beam_logits = self.joint.joint(beam_encoder_output, beam_decoder_output) / self.softmax_temperature
                beam_logp = torch.log_softmax(beam_logits[:, 0, 0, : -len(self.durations)], dim=-1)
                beam_duration_logp = torch.log_softmax(beam_logits[:, 0, 0, -len(self.durations) :], dim=-1)

                # Retrieve the top `max_candidades` most probable tokens.
                # Then, select the top `max_candidates` pairs of (token, duration)
                # based on the highest combined probabilities.
                # Note that indices are obtained in flattened array.
                beam_logp_topks, beam_idx_topks = beam_logp.topk(self.max_candidates, dim=-1)
                beam_total_logp = (beam_duration_logp[:, :, None] + beam_logp_topks[:, None, :]).view(
                    len(hyps), -1
                )  # [B, MAX_CANDIDATES*DURATION_BEAM]
                beam_total_logp_topks, beam_total_logp_topk_idxs = beam_total_logp.topk(
                    self.max_candidates, dim=-1
                )  # [B, MAX_CANDIDATES]

                # Prune hypothesis to obtain k expansions
                beam_best_expansion_scores = beam_total_logp_topks.max(dim=-1, keepdim=True).values
                beam_masks = beam_total_logp_topks >= beam_best_expansion_scores - self.maes_expansion_gamma
                beam_kexpansions_idxs = [
                    sum_logp_topk_idxs[mask] for sum_logp_topk_idxs, mask in zip(beam_total_logp_topk_idxs, beam_masks)
                ]

                list_exp = []  # List that contains the hypothesis expansion
                list_nb_exp = []  # List that contains the hypothesis expansion
                for hyp_idx, hyp in enumerate(hyps):  # For all hypothesis
                    for idx in beam_kexpansions_idxs[hyp_idx]:  # For all expansions within this hypothesis
                        # Restore indices in logp and durations_logp arrays from flattened indices.
                        k = int(beam_idx_topks[hyp_idx][idx % self.max_candidates])
                        duration = self.durations[int(idx // self.max_candidates)]
                        total_logp = float(beam_total_logp[hyp_idx][idx])

                        # Forcing blank token to have non-zero duration
                        if k == self.blank and duration == 0:
                            duration = self.durations[self.min_non_zero_duration_idx]

                        new_hyp = Hypothesis(
                            score=hyp.score + total_logp,
                            y_sequence=hyp.y_sequence[:],
                            dec_out=hyp.dec_out[:],
                            dec_state=hyp.dec_state,
                            timestamp=hyp.timestamp[:],
                            length=time_idx,
                            last_frame=hyp.last_frame + duration,
                        )

                        if self.ngram_lm:
                            new_hyp.ngram_lm_state = hyp.ngram_lm_state

                        # If the expansion was for blank
                        if k == self.blank:
                            list_b.append(new_hyp)
                        else:
                            new_hyp.y_sequence.append(k)
                            new_hyp.timestamp.append(time_idx + duration)

                            if self.ngram_lm:
                                lm_score, new_hyp.ngram_lm_state = self.compute_ngram_score(hyp.ngram_lm_state, int(k))
                                new_hyp.score += self.ngram_lm_alpha * lm_score

                            # If token duration is 0 adding to expansions list
                            if duration == 0:
                                list_exp.append(new_hyp)
                            else:
                                list_nb_exp.append(new_hyp)

                # Update states for hypothesis that do not end with blank
                hyps_to_update = list_nb_exp + list_exp
                if len(hyps_to_update) > 0:
                    # Decode a batch of beam states and scores
                    beam_decoder_output, beam_state = self.decoder.batch_score_hypothesis(
                        hyps_to_update,
                        cache,
                    )
                    for hyp_idx, hyp in enumerate(hyps_to_update):
                        # Preserve the decoder logits for the current beam
                        hyp.dec_out.append(beam_decoder_output[hyp_idx])
                        hyp.dec_state = beam_state[hyp_idx]

                # If there were no token expansions in any of the hypotheses,
                # Early exit
                list_nb += list_nb_exp
                if not list_exp:
                    kept_hyps = kept_hyps + list_b + list_nb
                    kept_hyps = self.merge_duplicate_hypotheses(kept_hyps)
                    kept_hyps = sorted(kept_hyps, key=lambda x: x.score, reverse=True)[:beam]

                    break
                else:
                    # If this isn't the last mAES step
                    if n < (self.maes_num_steps - 1):
                        # Copy the expanded hypothesis for the next iteration
                        hyps = self.merge_duplicate_hypotheses(list_exp)
                    else:
                        # If this is the last mAES step add probabilities of the blank token to the end.
                        # Extract the log probabilities
                        beam_decoder_output = torch.stack([h.dec_out[-1] for h in list_exp])  # [H, 1, D]
                        beam_logits = (
                            self.joint.joint(beam_encoder_output, beam_decoder_output) / self.softmax_temperature
                        )
                        beam_logp = torch.log_softmax(beam_logits[:, 0, 0, : -len(self.durations)], dim=-1)

                        # Get most probable durations
                        beam_duration_logp = torch.log_softmax(beam_logits[:, 0, 0, -len(self.durations) :], dim=-1)
                        _, beam_max_duration_idx = torch.max(beam_duration_logp, dim=-1)

                        # For all expansions, add the score for the blank label
                        for hyp_idx, hyp in enumerate(list_exp):
                            # If zero duration was obtained, change to the closest non-zero duration
                            duration_idx = int(beam_max_duration_idx[hyp_idx])
                            if duration_idx == self.zero_duration_idx:
                                duration_idx = self.min_non_zero_duration_idx

                            total_logp = float(
                                beam_logp[hyp_idx, self.blank] + beam_duration_logp[hyp_idx, duration_idx]
                            )
                            hyp.score += total_logp
                            hyp.last_frame += self.durations[duration_idx]

                        # Finally, update the kept hypothesis of sorted top Beam candidates
                        kept_hyps = kept_hyps + list_b + list_exp + list_nb
                        kept_hyps = self.merge_duplicate_hypotheses(kept_hyps)
                        kept_hyps = sorted(kept_hyps, key=lambda x: x.score, reverse=True)[:beam]

        # Sort the hypothesis with best scores
        return self.sort_nbest(kept_hyps)

    def merge_duplicate_hypotheses(self, hypotheses):
        """
        Merges hypotheses with identical token sequences and lengths.
        The combined hypothesis's probability is the sum of the probabilities of all duplicates.
        Duplicate hypotheses occur when two consecutive blank tokens are predicted
        and their duration values sum up to the same number.

        Args:
            hypotheses: list of hypotheses.

        Returns:
            hypotheses: list if hypotheses without duplicates.
        """
        sorted_hyps = sorted(hypotheses, key=lambda x: x.score, reverse=True)
        kept_hyps = {}
        for hyp in sorted_hyps:
            hyp_key = (tuple(hyp.y_sequence), int(hyp.last_frame))
            if hyp_key in kept_hyps:
                kept_hyp = kept_hyps[hyp_key]
                kept_hyp.score = float(torch.logaddexp(torch.tensor(kept_hyp.score), torch.tensor(hyp.score)))
            else:
                kept_hyps[hyp_key] = hyp
        return list(kept_hyps.values())

    def set_decoding_type(self, decoding_type: str):
        """
        Sets decoding type. Please check train_kenlm.py in scripts/asr_language_modeling/ to find out why we need
        Args:
            decoding_type: decoding type
        """
        # TOKEN_OFFSET for BPE-based models
        if decoding_type == 'subword':
            from nemo.collections.asr.parts.submodules.ctc_beam_decoding import DEFAULT_TOKEN_OFFSET

            self.token_offset = DEFAULT_TOKEN_OFFSET

    def prefix_search(
        self, hypotheses: List[Hypothesis], encoder_output: torch.Tensor, prefix_alpha: int
    ) -> List[Hypothesis]:
        """
        Performs a prefix search and updates the scores of the hypotheses in place.
        Based on https://arxiv.org/pdf/1211.3711.pdf.

        Args:
            hypotheses: a list of hypotheses sorted by the length from the longest to the shortest.
            encoder_output: encoder output.
            prefix_alpha: maximum allowable length difference between hypothesis and a prefix.

        Returns:
            hypotheses: list of hypotheses with updated scores.
        """
        # Iterate over hypotheses.
        for curr_idx, curr_hyp in enumerate(hypotheses[:-1]):
            # For each hypothesis, iterate over the subsequent hypotheses.
            # If a hypothesis is a prefix of the current one, update current score.
            for pref_hyp in hypotheses[(curr_idx + 1) :]:
                curr_hyp_length = len(curr_hyp.y_sequence)
                pref_hyp_length = len(pref_hyp.y_sequence)

                if (
                    is_prefix(curr_hyp.y_sequence, pref_hyp.y_sequence)
                    and (curr_hyp_length - pref_hyp_length) <= prefix_alpha
                ):
                    # Compute the score of the first token
                    # that follows the prefix hypothesis tokens in current hypothesis.
                    # Use the decoder output, which is stored in the prefix hypothesis.
                    logits = self.joint.joint(encoder_output, pref_hyp.dec_out[-1]) / self.softmax_temperature
                    logp = torch.log_softmax(logits[0, 0, 0, : -len(self.durations)], dim=-1)
                    duration_logp = torch.log_softmax(logits[0, 0, 0, -len(self.durations) :], dim=-1)
                    curr_score = pref_hyp.score + float(
                        logp[curr_hyp.y_sequence[pref_hyp_length]] + duration_logp[self.zero_duration_idx]
                    )

                    if self.ngram_lm:
                        lm_score, next_state = self.compute_ngram_score(
                            pref_hyp.ngram_lm_state, int(curr_hyp.y_sequence[pref_hyp_length])
                        )
                        curr_score += self.ngram_lm_alpha * lm_score

                    for k in range(pref_hyp_length, (curr_hyp_length - 1)):
                        # Compute the score of the next token.
                        # Approximate decoder output with the one that is stored in current hypothesis.
                        logits = self.joint.joint(encoder_output, curr_hyp.dec_out[k]) / self.softmax_temperature
                        logp = torch.log_softmax(logits[0, 0, 0, : -len(self.durations)], dim=-1)
                        duration_logp = torch.log_softmax(logits[0, 0, 0, -len(self.durations) :], dim=-1)
                        curr_score += float(logp[curr_hyp.y_sequence[k + 1]] + duration_logp[self.zero_duration_idx])

                        if self.ngram_lm:
                            lm_score, next_state = self.compute_ngram_score(
                                next_state, int(curr_hyp.y_sequence[k + 1])
                            )
                            curr_score += self.ngram_lm_alpha * lm_score

                    # Update current hypothesis score
                    curr_hyp.score = np.logaddexp(curr_hyp.score, curr_score)
        return hypotheses

    def compute_ngram_score(self, current_lm_state: "kenlm.State", label: int) -> Tuple[float, "kenlm.State"]:
        """
        Computes the score for KenLM Ngram language model.

        Args:
            current_lm_state: current state of the KenLM language model.
            label: next label.

        Returns:
            lm_score: score for `label`.
        """
        if self.token_offset:
            label = chr(label + self.token_offset)
        else:
            label = str(label)

        next_state = kenlm.State()
        lm_score = self.ngram_lm.BaseScore(current_lm_state, label, next_state)
        lm_score *= 1.0 / np.log10(np.e)

        return lm_score, next_state

    def sort_nbest(self, hyps: List[Hypothesis]) -> List[Hypothesis]:
        """Sort hypotheses by score or score given sequence length.

        Args:
            hyps: list of hypotheses

        Return:
            hyps: sorted list of hypotheses
        """
        if self.score_norm:
            return sorted(hyps, key=lambda x: x.score / len(x.y_sequence), reverse=True)
        else:
            return sorted(hyps, key=lambda x: x.score, reverse=True)


class BeamBatchedTDTInfer(Typing, ConfidenceMethodMixin, WithOptionalCudaGraphs):
    @property
    def input_types(self):
        """Returns definitions of module input ports."""
        return {
            "encoder_output": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
            "encoded_lengths": NeuralType(tuple('B'), LengthsType()),
            "partial_hypotheses": [NeuralType(elements_type=HypothesisType(), optional=True)],  # must always be last
        }

    def __init__(
        self,
        decoder_model: rnnt_abstract.AbstractRNNTDecoder,
        joint_model: rnnt_abstract.AbstractRNNTJoint,
        durations: list,
        blank_index: int,
        beam_size: int,
        search_type: str = 'malsd_batch',
        score_norm: bool = True,
        max_symbols_per_step: Optional[int] = None,
        preserve_alignments: bool = False,
        fusion_models: Optional[List[NGramGPULanguageModel]] = None,
        fusion_models_alpha: Optional[List[float]] = None,
        blank_lm_score_mode: Optional[str | BlankLMScoreMode] = BlankLMScoreMode.NO_SCORE,
        pruning_mode: Optional[str | PruningMode] = PruningMode.EARLY,
        allow_cuda_graphs: Optional[bool] = True,
        return_best_hypothesis: Optional[str] = True,
    ):
        """
        Init method.
        Args:
            decoder_model: Prediction network from RNN-T
            joint_model: Joint module from RNN-T
            durations: Token durations tensor for TDT
            blank_index: index of blank symbol
            beam_size: beam size
            max_symbols_per_step: max symbols to emit on each step (to avoid infinite looping)
            preserve_alignments: if alignments are needed
            fusion_models: list of fusion models to use for decoding
            fusion_models_alpha: list of alpha values for fusion models
            blank_lm_score_mode: mode for scoring blank symbol with LM
            pruning_mode: mode for pruning hypotheses with LM
            allow_cuda_graphs: whether to allow CUDA graphs
            score_norm: whether to normalize scores before best hypothesis extraction
        """
        super().__init__()
        self.decoder = decoder_model
        self.joint = joint_model

        self.durations = durations
        self._blank_index = blank_index
        self._SOS = blank_index  # Start of single index
        self.beam_size = beam_size
        self.return_best_hypothesis = return_best_hypothesis
        self.score_norm = score_norm

        if max_symbols_per_step is not None and max_symbols_per_step <= 0:
            raise ValueError(f"Expected max_symbols_per_step > 0 (or None), got {max_symbols_per_step}")
        self.max_symbols = max_symbols_per_step
        self.preserve_alignments = preserve_alignments

        if search_type == "malsd_batch":
            # Depending on availability of `blank_as_pad` support
            # switch between more efficient batch decoding technique
            self._decoding_computer = ModifiedALSDBatchedTDTComputer(
                decoder=self.decoder,
                joint=self.joint,
                durations=durations,
                beam_size=self.beam_size,
                blank_index=self._blank_index,
                max_symbols_per_step=self.max_symbols,
                preserve_alignments=preserve_alignments,
                fusion_models=fusion_models,
                fusion_models_alpha=fusion_models_alpha,
                blank_lm_score_mode=blank_lm_score_mode,
                pruning_mode=pruning_mode,
                allow_cuda_graphs=allow_cuda_graphs,
            )
        else:
            raise Exception(f"Decoding strategy {search_type} nor implemented.")

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

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

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

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

    @typecheck()
    def forward(
        self,
        encoder_output: torch.Tensor,
        encoded_lengths: torch.Tensor,
        partial_hypotheses: Optional[list[Hypothesis]] = None,
    ) -> Tuple[list[Hypothesis] | List[NBestHypotheses]]:
        """Returns a list of hypotheses given an input batch of the encoder hidden embedding.
        Output token is generated auto-regressively.
        Args:
            encoder_output: A tensor of size (batch, features, timesteps).
            encoded_lengths: list of int representing the length of each sequence
                output sequence.
        Returns:
            Tuple[list[Hypothesis] | List[NBestHypotheses]]: Tuple of a list of hypotheses for each batch. Each hypothesis contains
                the decoded sequence, timestamps and associated scores. The format of the returned hypotheses depends
                on the `return_best_hypothesis` attribute:
                    - If `return_best_hypothesis` is True, returns the best hypothesis for each batch.
                    - Otherwise, returns the N-best hypotheses for each batch.
        """
        if partial_hypotheses is not None:
            raise NotImplementedError("Partial hypotheses feature is not yet supported in batched beam search.")
        # Preserve decoder and joint training state
        decoder_training_state = self.decoder.training
        joint_training_state = self.joint.training

        with torch.inference_mode():
            # Apply optional preprocessing
            encoder_output = encoder_output.transpose(1, 2)  # (B, T, D)
            logitlen = encoded_lengths

            self.decoder.eval()
            self.joint.eval()

            inseq = encoder_output  # [B, T, D]
            batched_beam_hyps = self._decoding_computer(x=inseq, out_len=logitlen)

            # Ensures the correct number of hypotheses (batch_size) for CUDA Graphs compatibility
            batch_size = encoder_output.shape[0]
            if self.return_best_hypothesis:
                hyps = batched_beam_hyps.to_hyps_list(score_norm=self.score_norm)[:batch_size]
            else:
                hyps = batched_beam_hyps.to_nbest_hyps_list(score_norm=self.score_norm)[:batch_size]

        self.decoder.train(decoder_training_state)
        self.joint.train(joint_training_state)

        return (hyps,)
