# Created on 2024-01-01
# Author: GuAn Zhu

# Modified from NVIDIA(https://github.com/wenet-e2e/wenet/blob/main/runtime/gpu/
# model_repo_stateful/feature_extractor/1/model.py)

import triton_python_backend_utils as pb_utils
import json
import numpy as np
import yaml
from collections import OrderedDict

import torch
from torch.utils.dlpack import from_dlpack


class LimitedDict(OrderedDict):
    def __init__(self, max_length):
        super().__init__()
        self.max_length = max_length

    def __setitem__(self, key, value):
        if len(self) >= self.max_length:
            self.popitem(last=False)
        super().__setitem__(key, value)


class Fbank:
    def __init__(self, opts):
        self.opts = opts
        self._initialized = False

    def _ensure_initialized(self):
        if not self._initialized:
            import kaldifeat
            self.kaldifeat = kaldifeat
            self.fbank = kaldifeat.Fbank(self.opts)
            self._initialized = True

    def __call__(self, waves):
        self._ensure_initialized()
        return self.fbank(waves)

    def forward(self, waves):
        self._ensure_initialized()
        return self.fbank(waves)


class Feat(object):
    def __init__(self, seqid, offset_ms, sample_rate, frame_stride, device="cpu"):
        self.seqid = seqid
        self.sample_rate = sample_rate
        self._device = device
        self._get_torch()
        self.wav = self.torch.tensor([], device=device)
        self.offset = int(offset_ms / 1000 * sample_rate)
        self.frames = None
        self.frame_stride = int(frame_stride)
        self.device = device
        self.lfr_m = 7

    def _get_torch(self):
        if not hasattr(self, 'torch'):
            import torch
            self.torch = torch

    def add_wavs(self, wav):
        self._get_torch()
        wav = wav.to(self.device)
        self.wav = self.torch.cat((self.wav, wav), axis=0)

    def get_seg_wav(self):
        seg = self.wav[:]
        self.wav = self.wav[-self.offset :]
        return seg

    def add_frames(self, frames):
        """
        frames: seq_len x feat_sz
        """
        self._get_torch()
        if self.frames is None:
            self.frames = self.torch.cat((frames[0, :].repeat(((self.lfr_m - 1) // 2, 1)), frames), axis=0)
        else:
            self.frames = self.torch.cat([self.frames, frames], axis=0)

    def get_frames(self, num_frames: int):
        seg = self.frames[0:num_frames]
        self.frames = self.frames[self.frame_stride :]
        return seg


class TritonPythonModel:
    """Your Python model must use the same class name. Every Python model
    that is created must have "TritonPythonModel" as the class name.
    """

    def initialize(self, args):
        """`initialize` is called only once when the model is being loaded.
        Implementing `initialize` function is optional. This function allows
        the model to initialize any state associated with this model.
        Parameters
        ----------
        args : dict
          Both keys and values are strings. The dictionary keys and values are:
          * model_config: A JSON string containing the model configuration
          * model_instance_kind: A string containing model instance kind
          * model_instance_device_id: A string containing model instance device ID
          * model_repository: Model repository path
          * model_version: Model version
          * model_name: Model name
        """
        self.model_config = model_config = json.loads(args["model_config"])
        self.max_batch_size = max(model_config["max_batch_size"], 1)
        self.model_repository = args["model_repository"]

        if "GPU" in model_config["instance_group"][0]["kind"]:
            self.device = "cuda"
        else:
            self.device = "cpu"

        # Get OUTPUT0 configuration
        output0_config = pb_utils.get_output_config_by_name(model_config, "speech")
        # Convert Triton types to numpy types
        self.output0_dtype = pb_utils.triton_string_to_numpy(output0_config["data_type"])

        # 暂时保存numpy类型，后续会转为torch类型
        if self.output0_dtype == np.float32:
            self.dtype = np.float32
        else:
            self.dtype = np.float16

        self.feature_size = output0_config["dims"][-1]
        self.decoding_window = output0_config["dims"][-2]

        params = self.model_config["parameters"]
        for li in params.items():
            key, value = li
            value = value["string_value"]
            if key == "config_path":
                value = self.model_repository + "/" + value
                with open(str(value), "rb") as f:
                    config = yaml.load(f, Loader=yaml.Loader)

        # 保存配置信息，但不立即初始化torch和kaldifeat组件
        self.config = config
        self.feature_extractor = None
        self.opts = None
        self.seq_feat = None
        self.chunk_size = None
        self.frame_stride = None
        self.offset_ms = None
        self.sample_rate = None

        # 标记尚未初始化特征提取器
        self.initialized = False

        # 确保已初始化
        self._lazy_initialize()

    def _lazy_initialize(self):
        """延迟初始化特征提取器"""
        if self.initialized:
            return

        import torch
        import kaldifeat

        # 创建opts对象
        opts = kaldifeat.FbankOptions()
        opts.frame_opts.dither = 0.0
        opts.frame_opts.window_type = self.config["frontend_conf"]["window"]
        opts.mel_opts.num_bins = int(self.config["frontend_conf"]["n_mels"])
        opts.frame_opts.frame_shift_ms = float(self.config["frontend_conf"]["frame_shift"])
        opts.frame_opts.frame_length_ms = float(self.config["frontend_conf"]["frame_length"])
        opts.frame_opts.samp_freq = int(self.config["frontend_conf"]["fs"])
        opts.device = torch.device(self.device)
        self.opts = opts
        self.feature_extractor = Fbank(self.opts)

        self.seq_feat = LimitedDict(1024)
        chunk_size_s = float(self.model_config["parameters"]["chunk_size_s"]["string_value"]) # 0.6s

        sample_rate = opts.frame_opts.samp_freq # 16000Hz
        frame_shift_ms = opts.frame_opts.frame_shift_ms # 10ms
        frame_length_ms = opts.frame_opts.frame_length_ms # 25ms

        self.chunk_size = int(chunk_size_s * sample_rate) # 9600
        self.frame_stride = (chunk_size_s * 1000) // frame_shift_ms # 60
        self.offset_ms = self.get_offset(frame_length_ms, frame_shift_ms) # 5ms
        self.sample_rate = sample_rate # 16000Hz

        # 更新输出数据类型为torch类型
        if self.dtype == np.float32:
            self.dtype = torch.float32
        else:
            self.dtype = torch.float16

        self.initialized = True

    def get_offset(self, frame_length_ms, frame_shift_ms):
        offset_ms = 0
        while offset_ms + frame_shift_ms < frame_length_ms:
            offset_ms += frame_shift_ms
        return offset_ms

    def execute(self, requests):
        """`execute` must be implemented in every Python model. `execute`
        function receives a list of pb_utils.InferenceRequest as the only
        argument. This function is called when an inference is requested
        for this model.
        Parameters
        ----------
        requests : list
          A list of pb_utils.InferenceRequest
        Returns
        -------
        list
          A list of pb_utils.InferenceResponse. The length of this list must
          be the same as `requests`
        """

        total_waves = []
        responses = []
        batch_seqid = []
        end_seqid = {}
        for request in requests:
            input0 = pb_utils.get_input_tensor_by_name(request, "wav")
            wav = from_dlpack(input0.to_dlpack())[0]
            # input1 = pb_utils.get_input_tensor_by_name(request, "wav_lens")
            # wav_len = from_dlpack(input1.to_dlpack())[0]
            wav_len = len(wav)
            if wav_len < self.chunk_size:
                temp = torch.zeros(self.chunk_size, dtype=torch.float32, device=self.device)
                temp[0:wav_len] = wav[:]
                wav = temp

            in_start = pb_utils.get_input_tensor_by_name(request, "START")
            start = in_start.as_numpy()[0][0]
            in_ready = pb_utils.get_input_tensor_by_name(request, "READY")
            ready = in_ready.as_numpy()[0][0]
            in_corrid = pb_utils.get_input_tensor_by_name(request, "CORRID")
            corrid = in_corrid.as_numpy()[0][0]
            in_end = pb_utils.get_input_tensor_by_name(request, "END")
            end = in_end.as_numpy()[0][0]

            if start:
                self.seq_feat[corrid] = Feat(
                    corrid, self.offset_ms, self.sample_rate, self.frame_stride, self.device
                )
            if ready:
                self.seq_feat[corrid].add_wavs(wav)

            batch_seqid.append(corrid)
            if end:
                end_seqid[corrid] = 1

            wav = self.seq_feat[corrid].get_seg_wav() * 32768
            total_waves.append(wav)
        features = self.feature_extractor(total_waves)
        for corrid, frames in zip(batch_seqid, features):
            self.seq_feat[corrid].add_frames(frames)
            print(f"self.decoding_window: {self.decoding_window}, {len(frames)},  {len(self.seq_feat[corrid].frames)}")
            speech = self.seq_feat[corrid].get_frames(self.decoding_window)
            out_tensor0 = pb_utils.Tensor("speech", torch.unsqueeze(speech, 0).to("cpu").numpy())
            output_tensors = [out_tensor0]
            response = pb_utils.InferenceResponse(output_tensors=output_tensors)
            responses.append(response)
            if corrid in end_seqid:
                del self.seq_feat[corrid]
        return responses

    def finalize(self):
        print("Remove feature extractor!")
