import itertools
from typing import List, NamedTuple, Optional

import torch
from flashlight.lib.text.decoder import (
    create_emitting_model_state,
    CriterionType,
    DecodeResult,
    EmittingModelState,
    get_obj_from_emitting_model_state,
    LexiconDecoder,
    LexiconDecoderOptions,
    LexiconFreeDecoder,
    LexiconFreeDecoderOptions,
    LexiconFreeSeq2SeqDecoder,
    LexiconFreeSeq2SeqDecoderOptions,
    LexiconSeq2SeqDecoder,
    LexiconSeq2SeqDecoderOptions,
    LM,
    LMState,
    SmearingMode,
    Trie,
    TrieNode,
    ZeroLM,
    # KenLM
)

class CTCHypothesis(NamedTuple):
    r"""Represents hypothesis generated by CTC beam search decoder :py:func:`CTCDecoder`.

    :ivar torch.LongTensor tokens: Predicted sequence of token IDs. Shape `(L, )`, where
        `L` is the length of the output sequence
    :ivar List[str] words: List of predicted words
    :ivar float score: Score corresponding to hypothesis
    :ivar torch.IntTensor timesteps: Timesteps corresponding to the tokens. Shape `(L, )`,
        where `L` is the length of the output sequence
    """
    tokens: torch.LongTensor
    words: List[str]
    score: float
    timesteps: torch.IntTensor


class SeqHypothesis(NamedTuple):
    """
    A simulation of model state. These are synthetically created for the test
    but store information about model scores for the next timestep (i.e.
    "hidden states" in an autoregressive sense)
    """

    timestep: int
    token_idx: int
    score: float

    # def __init__(self, timestep, token_idx, score):
    #     self.timestep = timestep
    #     self.token_idx = token_idx
    #     self.score = score

class BeamSearchInterface():
    """
    .. devices:: CPU

    beam search decoder interface.

    Note:
        To build the decoder, please use the factory function :py:class:`BeamSearchInterface`.

    Args:
        nbest (int): number of best decodings to return
        beam_size (int, optional): max number of hypos to hold after each decode step (Default: 50)
        beam_size_token (int, optional): max number of tokens to consider at each decode step.
            If `None`, it is set to the total number of tokens (Default: None)
        beam_threshold (float, optional): threshold for pruning hypothesis (Default: 50)
        lm_weight (float, optional): weight of language model (Default: 2)
        sil_score (float, optional): silence insertion score (Default: 0)
        log_add (bool, optional): whether or not to use logadd when merging hypotheses (Default: False)
        blank (int) blank idx in vocabulary (Default: 0)
        silence (int) silence idx in vocabulary (Default: 0)
    """

    def __init__(
        self,
        nbest: int,
        beam_size: int = 50,
        beam_size_token: Optional[int] = 20,
        beam_threshold: float = 50,
        lm_weight: float = 2,
        sil_score: float = 0,
        log_add: bool = False,
        blank: int = 0,
        silence: int = 0,
    ) -> None:
        
        self.nbest = nbest
        self.blank = blank
        self.silence = silence
        self.decoder = None

    def decode(
        self, emissions: torch.FloatTensor, lengths: Optional[torch.Tensor] = None
    ):
        # Overriding the signature so that the return type is correct on Sphinx
        """decode(self, emissions: torch.FloatTensor, lengths: Optional[torch.Tensor] = None) -> \
            List[List[torchaudio.models.decoder.CTCHypothesis]]

        Args:
            emissions (torch.FloatTensor): CPU tensor of shape `(batch, frame, num_tokens)` storing sequences of
                probability distribution over labels; output of acoustic model.
            lengths (Tensor or None, optional): CPU tensor of shape `(batch, )` storing the valid length of
                in time axis of the output Tensor in each batch.
        """
        if emissions.dtype != torch.float32:
            raise ValueError("emissions must be float32.")
        if emissions.is_cuda:
            raise RuntimeError("emissions must be a CPU tensor.")
        if lengths is not None and lengths.is_cuda:
            raise RuntimeError("lengths must be a CPU tensor.")


class BeamSearchCtc(BeamSearchInterface):
    """
    .. devices:: CPU

    CTC beam search decoder from *Flashlight* [:footcite:`kahn2022flashlight`].

    Note:
        To build the decoder, please use the factory function :py:func:`ctc_decoder`.

    Args:
        nbest (int): number of best decodings to return
        beam_size (int, optional): max number of hypos to hold after each decode step (Default: 50)
        beam_size_token (int, optional): max number of tokens to consider at each decode step.
            If `None`, it is set to the total number of tokens (Default: None)
        beam_threshold (float, optional): threshold for pruning hypothesis (Default: 50)
        lm_weight (float, optional): weight of language model (Default: 2)
        sil_score (float, optional): silence insertion score (Default: 0)
        log_add (bool, optional): whether or not to use logadd when merging hypotheses (Default: False)
        blank (int) blank idx in vocabulary (Default: 0)
        silence (int) silence idx in vocabulary (Default: 0)
    """

    def __init__(
        self,
        nbest: int,
        beam_size: int = 50,
        beam_size_token: Optional[int] = 20,
        beam_threshold: float = 50,
        lm_weight: float = 2,
        sil_score: float = 0,
        log_add: bool = False,
        blank: int = 0,
        silence: int = 0,
    ) -> None:
        super().__init__(
            nbest = nbest,
            beam_size = beam_size,
            beam_size_token = beam_size_token,
            beam_threshold = beam_threshold,
            lm_weight = lm_weight,
            sil_score = sil_score,
            log_add = log_add,
            blank = blank,
            silence = silence,
        )
        decoder_options = LexiconFreeDecoderOptions(
            beam_size=beam_size,
            beam_size_token=beam_size_token,
            beam_threshold=beam_threshold,
            lm_weight=lm_weight,
            sil_score=sil_score,
            log_add=log_add,
            criterion_type=CriterionType.CTC,
        )
        self.decoder = LexiconFreeDecoder(decoder_options, ZeroLM(), self.silence, self.blank, [])

    def _get_tokens(self, idxs: torch.IntTensor) -> torch.LongTensor:
        idxs = (g[0] for g in itertools.groupby(idxs))
        idxs = filter(lambda x: x != self.blank, idxs)
        return torch.LongTensor(list(idxs))

    def _get_timesteps(self, idxs: torch.IntTensor) -> torch.IntTensor:
        """Returns frame numbers corresponding to non-blank tokens."""

        timesteps = []
        for i, idx in enumerate(idxs):
            if idx == self.blank:
                continue
            if i == 0 or idx != idxs[i - 1]:
                timesteps.append(i)
        return torch.IntTensor(timesteps)

    def decode(
        self, emissions: torch.FloatTensor, lengths: Optional[torch.Tensor] = None
    ) -> List[List[CTCHypothesis]]:
        # Overriding the signature so that the return type is correct on Sphinx
        """decode(self, emissions: torch.FloatTensor, lengths: Optional[torch.Tensor] = None) -> \
            List[List[torchaudio.models.decoder.CTCHypothesis]]

        Args:
            emissions (torch.FloatTensor): CPU tensor of shape `(batch, frame, num_tokens)` storing sequences of
                probability distribution over labels; output of acoustic model.
            lengths (Tensor or None, optional): CPU tensor of shape `(batch, )` storing the valid length of
                in time axis of the output Tensor in each batch.

        Returns:
            List[List[CTCHypothesis]]:
                List of sorted best hypotheses for each audio sequence in the batch.
        """
        super().decode(emissions=emissions, lengths=lengths)

        B, T, N = emissions.size()
        if lengths is None:
            lengths = torch.full((B,), T)

        float_bytes = 4
        hypos = []

        for b in range(B):
            emissions_ptr = emissions.data_ptr() + float_bytes * b * emissions.stride(0)

            results = self.decoder.decode(emissions_ptr, lengths[b], N)

            nbest_results = results[: self.nbest]
            hypos.append(
                [
                    CTCHypothesis(
                        tokens=self._get_tokens(result.tokens),
                        # words=[self.word_dict.get_entry(x) for x in result.words if x >= 0],
                        words=[],
                        score=result.score,
                        timesteps=self._get_timesteps(result.tokens),
                    )
                    for result in nbest_results
                ]
            )

        return hypos





class BeamSearchSequence(BeamSearchInterface):
    """
    .. devices:: CPU

    sequence beam search decoder from *Flashlight* [:footcite:`kahn2022flashlight`].

    Note:
        To build the decoder, please use the factory function :py:func:`ctc_decoder`.

    Args:
        nbest (int): number of best decodings to return
        beam_size (int, optional): max number of hypos to hold after each decode step (Default: 50)
        beam_size_token (int, optional): max number of tokens to consider at each decode step.
            If `None`, it is set to the total number of tokens (Default: None)
        beam_threshold (float, optional): threshold for pruning hypothesis (Default: 50)
        lm_weight (float, optional): weight of language model (Default: 2)
        sil_score (float, optional): silence insertion score (Default: 0)
        log_add (bool, optional): whether or not to use logadd when merging hypotheses (Default: False)
        blank (int) blank idx in vocabulary (Default: 0)
        silence (int) silence idx in vocabulary (Default: 0)
    """

    def __init__(
        self,
        nbest: int,
        vocab_size: int,
        beam_size: int = 50,
        beam_size_token: Optional[int] = 20,
        beam_threshold: float = 50,
        lm_weight: float = 0.0,
        sil_score: float = -1.0,
        log_add: bool = False,
        blank: int = 0,
        silence: int = 0,
    ) -> None:
        super().__init__(
            nbest = nbest,
            beam_size = beam_size,
            beam_size_token = beam_size_token,
            beam_threshold = beam_threshold,
            lm_weight = lm_weight,
            sil_score = sil_score,
            log_add = log_add,
            blank = blank,
            silence = silence,
        )
        decoder_options = LexiconFreeDecoderOptions(
            beam_size=beam_size,
            beam_size_token=beam_size_token,
            beam_threshold=beam_threshold,
            lm_weight=lm_weight,
            # word_score = 2.0,
            # unk_score = -float("inf"),
            sil_score=sil_score,
            log_add=log_add,
            criterion_type=CriterionType.ASG,
        )
        # transitions = [1.0/vocab_size] * (vocab_size**2)
        class Transition():
            def __init__(self) -> None:
                pass
            
            def __getitem__(self, index) -> float:
                return 0.0

            def __len__(self) -> int:
                return vocab_size**2
        
        # transitions = [0.0] * (vocab_size**2)
        transitions = Transition()
        self.decoder = LexiconFreeDecoder(decoder_options, ZeroLM(), self.silence, self.blank, transitions)

    def _get_tokens(self, idxs: torch.IntTensor) -> torch.LongTensor:
        idxs = (g[0] for g in itertools.groupby(idxs))
        idxs = filter(lambda x: x != self.blank, idxs)
        return torch.LongTensor(list(idxs))

    def _get_timesteps(self, idxs: torch.IntTensor) -> torch.IntTensor:
        """Returns frame numbers corresponding to non-blank tokens."""

        timesteps = []
        for i, idx in enumerate(idxs):
            if idx == self.blank:
                continue
            if i == 0 or idx != idxs[i - 1]:
                timesteps.append(i)
        return torch.IntTensor(timesteps)

    def decode(
        self, emissions: torch.FloatTensor, lengths: Optional[torch.Tensor] = None
    ) -> List[List[CTCHypothesis]]:
        # Overriding the signature so that the return type is correct on Sphinx
        """decode(self, emissions: torch.FloatTensor, lengths: Optional[torch.Tensor] = None) -> \
            List[List[torchaudio.models.decoder.CTCHypothesis]]

        Args:
            emissions (torch.FloatTensor): CPU tensor of shape `(batch, frame, num_tokens)` storing sequences of
                probability distribution over labels; output of acoustic model.
            lengths (Tensor or None, optional): CPU tensor of shape `(batch, )` storing the valid length of
                in time axis of the output Tensor in each batch.

        Returns:
            List[List[CTCHypothesis]]:
                List of sorted best hypotheses for each audio sequence in the batch.
        """
        super().decode(emissions=emissions, lengths=lengths)

        B, T, N = emissions.size()
        if lengths is None:
            lengths = torch.full((B,), T)

        float_bytes = 4
        hypos = []

        for b in range(B):
            emissions_ptr = emissions.data_ptr() + float_bytes * b * emissions.stride(0)
            results = self.decoder.decode(emissions_ptr, lengths[b], N)
            nbest_results = results[: self.nbest]
            hypos.append(
                [
                    CTCHypothesis(
                        tokens=self._get_tokens(result.tokens),
                        # words=[self.word_dict.get_entry(x) for x in result.words if x >= 0],
                        words=[],
                        score=result.score,
                        # timesteps=self._get_timesteps(result.tokens),
                        timesteps=[],
                    )
                    for result in nbest_results
                ]
            )

        return hypos

