from typing import List, Tuple

import torch
from modules.wenet_extractor.utils.common import log_add


class Sequence:
    __slots__ = {"hyp", "score", "cache"}

    def __init__(
        self,
        hyp: List[torch.Tensor],
        score,
        cache: List[torch.Tensor],
    ):
        self.hyp = hyp
        self.score = score
        self.cache = cache


class PrefixBeamSearch:
    def __init__(self, encoder, predictor, joint, ctc, blank):
        self.encoder = encoder
        self.predictor = predictor
        self.joint = joint
        self.ctc = ctc
        self.blank = blank

    def forward_decoder_one_step(
        self, encoder_x: torch.Tensor, pre_t: torch.Tensor, cache: List[torch.Tensor]
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        padding = torch.zeros(pre_t.size(0), 1, device=encoder_x.device)
        pre_t, new_cache = self.predictor.forward_step(
            pre_t.unsqueeze(-1), padding, cache
        )
        x = self.joint(encoder_x, pre_t)  # [beam, 1, 1, vocab]
        x = x.log_softmax(dim=-1)
        return x, new_cache

    def prefix_beam_search(
        self,
        speech: torch.Tensor,
        speech_lengths: torch.Tensor,
        decoding_chunk_size: int = -1,
        beam_size: int = 5,
        num_decoding_left_chunks: int = -1,
        simulate_streaming: bool = False,
        ctc_weight: float = 0.3,
        transducer_weight: float = 0.7,
    ):
        """prefix beam search
        also see wenet.transducer.transducer.beam_search
        """
        assert speech.shape[0] == speech_lengths.shape[0]
        assert decoding_chunk_size != 0
        device = speech.device
        batch_size = speech.shape[0]
        assert batch_size == 1

        # 1. Encoder
        encoder_out, _ = self.encoder(
            speech, speech_lengths, decoding_chunk_size, num_decoding_left_chunks
        )  # (B, maxlen, encoder_dim)
        maxlen = encoder_out.size(1)

        ctc_probs = self.ctc.log_softmax(encoder_out).squeeze(0)
        beam_init: List[Sequence] = []

        # 2. init beam using Sequence to save beam unit
        cache = self.predictor.init_state(1, method="zero", device=device)
        beam_init.append(Sequence(hyp=[self.blank], score=0.0, cache=cache))
        # 3. start decoding (notice: we use breathwise first searching)
        # !!!! In this decoding method: one frame do not output multi units. !!!!
        # !!!!    Experiments show that this strategy has little impact      !!!!
        for i in range(maxlen):
            # 3.1 building input
            # decoder taking the last token to predict the next token
            input_hyp = [s.hyp[-1] for s in beam_init]
            input_hyp_tensor = torch.tensor(input_hyp, dtype=torch.int, device=device)
            # building statement from beam
            cache_batch = self.predictor.cache_to_batch([s.cache for s in beam_init])
            # build score tensor to do torch.add() function
            scores = torch.tensor([s.score for s in beam_init]).to(device)

            # 3.2 forward decoder
            logp, new_cache = self.forward_decoder_one_step(
                encoder_out[:, i, :].unsqueeze(1),
                input_hyp_tensor,
                cache_batch,
            )  # logp: (N, 1, 1, vocab_size)
            logp = logp.squeeze(1).squeeze(1)  # logp: (N, vocab_size)
            new_cache = self.predictor.batch_to_cache(new_cache)

            # 3.3 shallow fusion for transducer score
            #     and ctc score where we can also add the LM score
            logp = torch.log(
                torch.add(
                    transducer_weight * torch.exp(logp),
                    ctc_weight * torch.exp(ctc_probs[i].unsqueeze(0)),
                )
            )

            # 3.4 first beam prune
            top_k_logp, top_k_index = logp.topk(beam_size)  # (N, N)
            scores = torch.add(scores.unsqueeze(1), top_k_logp)

            # 3.5 generate new beam (N*N)
            beam_A = []
            for j in range(len(beam_init)):
                # update seq
                base_seq = beam_init[j]
                for t in range(beam_size):
                    # blank: only update the score
                    if top_k_index[j, t] == self.blank:
                        new_seq = Sequence(
                            hyp=base_seq.hyp.copy(),
                            score=scores[j, t].item(),
                            cache=base_seq.cache,
                        )

                        beam_A.append(new_seq)
                    # other unit: update hyp score statement and last
                    else:
                        hyp_new = base_seq.hyp.copy()
                        hyp_new.append(top_k_index[j, t].item())
                        new_seq = Sequence(
                            hyp=hyp_new, score=scores[j, t].item(), cache=new_cache[j]
                        )
                        beam_A.append(new_seq)

            # 3.6 prefix fusion
            fusion_A = [beam_A[0]]
            for j in range(1, len(beam_A)):
                s1 = beam_A[j]
                if_do_append = True
                for t in range(len(fusion_A)):
                    # notice: A_ can not fusion with A
                    if s1.hyp == fusion_A[t].hyp:
                        fusion_A[t].score = log_add([fusion_A[t].score, s1.score])
                        if_do_append = False
                        break
                if if_do_append:
                    fusion_A.append(s1)

            # 4. second pruned
            fusion_A.sort(key=lambda x: x.score, reverse=True)
            beam_init = fusion_A[:beam_size]

        return beam_init, encoder_out