import torch
from transformers import is_torch_npu_available

import time
import numpy as np
from app.log import logger
from funasr import AutoModel
from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
from app import config
from tqdm import tqdm
from app.util import utils

if utils.is_not_blank(config.device) and "npu" in config.device.lower() and is_torch_npu_available():
    import torch_npu
    from ais_bench.infer.interface import InferSession

class SenseVoiceOnnxModel():
    def __init__(self):
        super().__init__()
        self.blank_id = 0
        self.lid_dict = {"auto": 0, "zh": 3, "en": 4, "yue": 7, "ja": 11, "ko": 12, "nospeech": 13}
        self.textnorm_dict = {'withitn': 14, "woitn": 15}
        self.current_device_index = config.current_device_index
        torch_npu.npu.set_device(self.current_device_index)
        _, kwargs = AutoModel.build_model(model=config.model.auto_speech_recognition.sense_voice.path,
                                          trust_remote_code=False)
        self.om_path = config.model.auto_speech_recognition.sense_voice.om_path
        self.om_sess = InferSession(self.current_device_index, self.om_path)
        self.kwargs = kwargs
        self.transform = True

    def __call__(
            self,
            data_in,
            language=None,
            use_itn=None,
            tokenizer=None,
            frontend=None,
            show_progress=True,
    ):
        results = []
        try:
            torch_npu.npu.set_device(self.current_device_index)
            key = ["wav_file_tmp_name"]
            if use_itn == None:
                use_itn = self.kwargs.get('use_itn', False)
            kwargs = self.kwargs
            # 直接传音频波形的时候transform=True
            if self.transform:
                audio_sample_list = data_in
            else:
                audio_sample_list = load_audio_text_image_video(
                    data_in,
                    fs=frontend.fs,
                    audio_fs=kwargs.get("fs", 16000),
                    data_type=kwargs.get("data_type", "sound"),
                    tokenizer=tokenizer,
                )

            data_type = kwargs.get("data_type", "sound")

            if frontend == None:
                frontend = kwargs.get('frontend')
            frontend_conf = kwargs.get('frontend_conf')
            if tokenizer == None:
                tokenizer = kwargs.get('tokenizer')
            tokenizer_conf = kwargs.get('tokenizer_conf')
            if show_progress:
                progress_bar = tqdm(
                    total=1,
                    desc="segment for asr processing",
                    unit="frames",
                    bar_format="{l_bar}{bar}{r_bar} | {percentage:.2f}%",
                )
            speech, speech_lengths = extract_fbank(
                audio_sample_list, data_type=kwargs.get("data_type", "sound"), frontend=frontend
            )
            speech = speech.to(device=kwargs["device"])
            speech_lengths = speech_lengths.to(device=kwargs["device"])

            if language == None:
                language = kwargs.get("language", "auto")
            language = torch.LongTensor([self.lid_dict[language] if language in self.lid_dict else 0]).to(speech.device)

            textnorm = kwargs.get("text_norm", None)
            if textnorm is None:
                textnorm = "withitn" if use_itn else "woitn"
            textnorm = torch.LongTensor([self.textnorm_dict[textnorm]]).to(speech.device)

            s = time.time()
            feed = [speech.cpu().detach().numpy().astype(np.float32),
                    speech_lengths.cpu().detach().numpy().astype(np.int32),
                    language.cpu().detach().numpy().astype(np.int32),
                    textnorm.cpu().detach().numpy().astype(np.int32)]

            ctc_logits, encoder_out_lens = self.om_sess.infer(feed, mode='dymshape', custom_sizes=300000000)
            ctc_logits = torch.from_numpy(ctc_logits).npu()
            encoder_out_lens = torch.from_numpy(encoder_out_lens).npu()
            e = time.time()
            cost_time = e - s



            x = ctc_logits[0, : encoder_out_lens[0].item(), :]
            yseq = x.argmax(dim=-1)
            yseq = torch.unique_consecutive(yseq, dim=-1)
            mask = yseq != self.blank_id
            token_int = yseq[mask].tolist()

            # Change integer-ids to tokens
            text = tokenizer.decode(token_int)
            result_i = {"key": key[0], "text": text}
            results.append(result_i)
            if show_progress:
                progress_bar.update(1)
        except Exception as e:
            logger.error(f"Sensevoice Inference failed: {e}")
        return results