import re
import logging

import torch
import math
from collections import defaultdict
from typing import List, Optional, Tuple


symbol_str = '[’!"#$%&\'()*+,-./:;<>=?@，。?★、…【】《》？“”‘’！[\\]^_`{|}~\s]+'


def split_mixed_label(input_str):
    tokens = []
    s = input_str.lower()
    while len(s) > 0:
        match = re.match(r'[A-Za-z!?,<>()\']+', s)
        if match is not None:
            word = match.group(0)
        else:
            word = s[0:1]
        tokens.append(word)
        s = s.replace(word, '', 1).strip(' ')
    return tokens


def query_token_set(txt, symbol_table, lexicon_table):
    tokens_str = tuple()
    tokens_idx = tuple()

    if txt in symbol_table:
        tokens_str = tokens_str + (txt, )
        tokens_idx = tokens_idx + (symbol_table[txt], )
        return tokens_str, tokens_idx

    parts = split_mixed_label(txt)
    for part in parts:
        if part == '!sil' or part == '(sil)' or part == '<sil>':
            tokens_str = tokens_str + ('!sil', )
        elif part == '<blank>' or part == '<blank>':
            tokens_str = tokens_str + ('<blank>', )
        elif part == '(noise)' or part == 'noise)' or part == '(noise' or part == '<noise>':
            tokens_str = tokens_str + ('<unk>', )
        elif part in symbol_table:
            tokens_str = tokens_str + (part, )
        elif part in lexicon_table:
            for ch in lexicon_table[part]:
                tokens_str = tokens_str + (ch, )
        else:
            part = re.sub(symbol_str, '', part)
            for ch in part:
                tokens_str = tokens_str + (ch, )

    for ch in tokens_str:
        if ch in symbol_table:
            tokens_idx = tokens_idx + (symbol_table[ch], )
        elif ch == '!sil':
            if 'sil' in symbol_table:
                tokens_idx = tokens_idx + (symbol_table['sil'], )
            else:
                tokens_idx = tokens_idx + (symbol_table['<blank>'], )
        elif ch == '<unk>':
            if '<unk>' in symbol_table:
                tokens_idx = tokens_idx + (symbol_table['<unk>'], )
            else:
                tokens_idx = tokens_idx + (symbol_table['<blank>'], )
        else:
            if '<unk>' in symbol_table:
                tokens_idx = tokens_idx + (symbol_table['<unk>'], )
                logging.info(f'\'{ch}\' is not in token set, replace with <unk>')
            else:
                tokens_idx = tokens_idx + (symbol_table['<blank>'], )
                logging.info(f'\'{ch}\' is not in token set, replace with <blank>')

    return tokens_str, tokens_idx


class KwsCtcPrefixDecoder():
    """Decoder interface wrapper for CTCPrefixDecode."""

    def __init__(
        self,
        ctc: torch.nn.Module,
        keywords: str,
        token_list: list,
        seg_dict: dict,
    ):
        """Initialize class.

        Args:
            ctc (torch.nn.Module): The CTC implementation.
                For example, :class:`espnet.nets.pytorch_backend.ctc.CTC`

        """
        self.ctc = ctc
        self.token_list = token_list

        token_table = {}
        for token in token_list:
            token_table[token] = token_list.index(token)

        self.keywords_idxset = {0}
        self.keywords_token = {}
        self.keywords_str = keywords
        keywords_list = self.keywords_str.strip().replace(' ', '').split(',')
        for keyword in keywords_list:
            strs, indexs = query_token_set(keyword, token_table, seg_dict)
            self.keywords_token[keyword] = {}
            self.keywords_token[keyword]['token_id'] = indexs
            self.keywords_token[keyword]['token_str'] = ''.join('%s ' % str(i) for i in indexs)
            [ self.keywords_idxset.add(i) for i in indexs ]

    def beam_search(
        self,
        logits: torch.Tensor,
        logits_lengths: torch.Tensor,
        keywords_tokenset: set = None,
        score_beam_size: int = 3,
        path_beam_size: int = 20,
    ) -> Tuple[List[List[int]], torch.Tensor]:
        """ CTC prefix beam search inner implementation

        Args:
            logits (torch.Tensor): (1, max_len, vocab_size)
            logits_lengths (torch.Tensor): (1, )
            keywords_tokenset (set): token set for filtering score
            score_beam_size (int): beam size for score
            path_beam_size (int): beam size for path

        Returns:
            List[List[int]]: nbest results
        """

        maxlen = logits.size(0)
        ctc_probs = logits
        cur_hyps = [(tuple(), (1.0, 0.0, []))]

        # CTC beam search step by step
        for t in range(0, maxlen):
            probs = ctc_probs[t]  # (vocab_size,)
            # key: prefix, value (pb, pnb), default value(-inf, -inf)
            next_hyps = defaultdict(lambda: (0.0, 0.0, []))

            # 2.1 First beam prune: select topk best
            top_k_probs, top_k_index = probs.topk(
                score_beam_size)  # (score_beam_size,)

            # filter prob score that is too small
            filter_probs = []
            filter_index = []
            for prob, idx in zip(top_k_probs.tolist(), top_k_index.tolist()):
                if keywords_tokenset is not None:
                    if prob > 0.05 and idx in keywords_tokenset:
                        filter_probs.append(prob)
                        filter_index.append(idx)
                else:
                    if prob > 0.05:
                        filter_probs.append(prob)
                        filter_index.append(idx)

            if len(filter_index) == 0:
                continue

            for s in filter_index:
                ps = probs[s].item()
                # print(f'frame:{t}, token:{s}, score:{ps}')

                for prefix, (pb, pnb, cur_nodes) in cur_hyps:
                    last = prefix[-1] if len(prefix) > 0 else None
                    if s == 0:  # blank
                        n_pb, n_pnb, nodes = next_hyps[prefix]
                        n_pb = n_pb + pb * ps + pnb * ps
                        nodes = cur_nodes.copy()
                        next_hyps[prefix] = (n_pb, n_pnb, nodes)
                    elif s == last:
                        if not math.isclose(pnb, 0.0, abs_tol=0.000001):
                            # Update *ss -> *s;
                            n_pb, n_pnb, nodes = next_hyps[prefix]
                            n_pnb = n_pnb + pnb * ps
                            nodes = cur_nodes.copy()
                            if ps > nodes[-1]['prob']:  # update frame and prob
                                nodes[-1]['prob'] = ps
                                nodes[-1]['frame'] = t
                            next_hyps[prefix] = (n_pb, n_pnb, nodes)

                        if not math.isclose(pb, 0.0, abs_tol=0.000001):
                            # Update *s-s -> *ss, - is for blank
                            n_prefix = prefix + (s, )
                            n_pb, n_pnb, nodes = next_hyps[n_prefix]
                            n_pnb = n_pnb + pb * ps
                            nodes = cur_nodes.copy()
                            nodes.append(dict(token=s, frame=t,
                                              prob=ps))  # to record token prob
                            next_hyps[n_prefix] = (n_pb, n_pnb, nodes)
                    else:
                        n_prefix = prefix + (s, )
                        n_pb, n_pnb, nodes = next_hyps[n_prefix]
                        if nodes:
                            if ps > nodes[-1]['prob']:  # update frame and prob
                                nodes[-1]['prob'] = ps
                                nodes[-1]['frame'] = t
                        else:
                            nodes = cur_nodes.copy()
                            nodes.append(dict(token=s, frame=t,
                                              prob=ps))  # to record token prob
                        n_pnb = n_pnb + pb * ps + pnb * ps
                        next_hyps[n_prefix] = (n_pb, n_pnb, nodes)

            # 2.2 Second beam prune
            next_hyps = sorted(next_hyps.items(),
                               key=lambda x: (x[1][0] + x[1][1]),
                               reverse=True)

            cur_hyps = next_hyps[:path_beam_size]

        hyps = [(y[0], y[1][0] + y[1][1], y[1][2]) for y in cur_hyps]
        return hyps


    def is_sublist(self, main_list, check_list):
        if len(main_list) < len(check_list):
            return -1

        if len(main_list) == len(check_list):
            return 0 if main_list == check_list else -1

        for i in range(len(main_list) - len(check_list)):
            if main_list[i] == check_list[0]:
                for j in range(len(check_list)):
                    if main_list[i + j] != check_list[j]:
                        break
                else:
                    return i
        else:
            return -1


    def _decode_inside(
        self,
        logits: torch.Tensor,
        logits_lengths: torch.Tensor,
    ):
        hyps = self.beam_search(logits, logits_lengths, self.keywords_idxset)

        hit_keyword = None
        hit_score = 1.0
        # start = 0; end = 0
        for one_hyp in hyps:
            prefix_ids = one_hyp[0]
            # path_score = one_hyp[1]
            prefix_nodes = one_hyp[2]
            assert len(prefix_ids) == len(prefix_nodes)
            for word in self.keywords_token.keys():
                lab = self.keywords_token[word]['token_id']
                offset = self.is_sublist(prefix_ids, lab)
                if offset != -1:
                    hit_keyword = word
                    for idx in range(offset, offset + len(lab)):
                        hit_score *= prefix_nodes[idx]['prob']
                    break
            if hit_keyword is not None:
                hit_score = math.sqrt(hit_score)
                break

        if hit_keyword is not None:
            return True, hit_keyword, hit_score
        else:
            return False, None, None


    def decode(self, x: torch.Tensor):
        """Get an initial state for decoding.

        Args:
            x (torch.Tensor): The encoded feature tensor

        Returns: decode result

        """

        raw_logp = self.ctc.softmax(x.unsqueeze(0)).detach().squeeze(0).cpu()
        xlen = torch.tensor([raw_logp.size(1)])

        return self._decode_inside(raw_logp, xlen)