#!/bin/bash
#
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import triton_python_backend_utils as pb_utils
import numpy as np

import json
import yaml
import time
# 在这里加上utils.py文件包含的各种工具，直接使用
from typing import Dict, List, Union,Iterable,Tuple
from pathlib import Path

from frontend import WavFrontendOnline
from e2evad import E2EVadModel  # WavFrontendOnline


def read_yaml(yaml_path: Union[str, Path]) -> Dict:
        if not Path(yaml_path).exists():
            raise FileExistsError(f"The {yaml_path} does not exist.")

        with open(str(yaml_path), "rb") as f:
            data = yaml.load(f, Loader=yaml.Loader)
        return data

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"]
        vad_config_file = self.model_repository + "/1/config.yaml"
        vad_cmvn_file = self.model_repository + "/1/am.mvn"

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

        vad_config=read_yaml(vad_config_file)
        self.frontend=WavFrontendOnline(cmvn_file=vad_cmvn_file,**vad_config["frontend_conf"])
        self.batch_size = 1
        self.vad_scorer = E2EVadModel(vad_config["model_conf"])
        max_end_sil:int=None
        self.max_end_sil = (
            max_end_sil if max_end_sil is not None else vad_config["model_conf"]["max_end_silence_time"]
        )
        self.encoder_conf = vad_config["encoder_conf"]

        # 这里添加等待时间，如果超过时间就直接结束
        self.last_request_time = {} 

        # 这里加载起来模型，可以不用每次使用都加载一次
        pb_utils.load_model("asr_fsmn-vad_streaming")
        self.in_cache0=[]
        self.in_cache1=[]
        self.in_cache2=[]
        self.in_cache3=[]

        self.out_cache0=[]
        self.out_cache1=[]
        self.out_cache2=[]
        self.out_cache3=[]
        
    # 从原来的list操作变成了单个cache操作
    def prepare_cache(self, in_cache: np.ndarray = None):
        if len(in_cache) > 0:
            return in_cache
        fsmn_layers = self.encoder_conf["fsmn_layers"]
        proj_dim = self.encoder_conf["proj_dim"]
        lorder = self.encoder_conf["lorder"]
        cache = np.zeros((1, proj_dim, lorder - 1, 1), dtype=np.float32)
        return cache
    
    def extract_feat(
        self, waveforms: np.ndarray, is_final: bool = False
    ) -> Tuple[np.ndarray, np.ndarray]:
        waveforms_lens = np.zeros(waveforms.shape[0]).astype(np.int32)
        for idx, waveform in enumerate(waveforms):
            waveforms_lens[idx] = waveform.shape[-1]

        feats, feats_len = self.frontend.extract_fbank(waveforms, waveforms_lens, is_final)

        return feats.astype(np.float32), feats_len.astype(np.int32)

    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`
        """
        # Every Python backend must iterate through list of requests and create
        # an instance of pb_utils.InferenceResponse class for each of them. You
        # should avoid storing any of the input Tensors in the class attributes
        # as they will be overridden in subsequent inference requests. You can
        # make a copy of the underlying NumPy array and store it if it is
        # required.

        total_seq = 0
        text_no_punc_list = []
        responses = []
        response = []

        for request in requests:
            # 获取 sequence_id 和 control 参数
            in_sequence_id = pb_utils.get_input_tensor_by_name(request, "CORRID")
            sequence_id = in_sequence_id.as_numpy()[0]  # shape (1, 1) 取出第一个元素，得到 numpy 数组，dtype=int32，shape=[1]
            in_seq_start = pb_utils.get_input_tensor_by_name(request, "START")
            seq_start = in_seq_start.as_numpy()[0]  # shape (1, 1) 取出第一个元素，得到 numpy 数组，dtype=int32，shape=[1]
            in_seq_end = pb_utils.get_input_tensor_by_name(request, "END")
            seq_end = in_seq_end.as_numpy()[0]  # shape (1, 1) 取出第一个元素，得到 numpy 数组，dtype=int32，shape=[1]
            # Perform inference on the request and append it to responses list...
            in_0 = pb_utils.get_input_tensor_by_name(request, "WAV")

            wav_no_vad = in_0.as_numpy()  # 得到 numpy 数组，dtype=object，shape=[batch_size, 1]

            waveform=np.array(wav_no_vad.tolist(),dtype=object)[0]  # 取出第一个元素，得到 numpy 数组，dtype=float32，shape=[1, 160000]

            total_seq += 1

            current_time = time.time() # 获取当前时间戳
            last_time = self.last_request_time.get(sequence_id.item(), current_time)
            if current_time - last_time > 2.0:
                self.frontend.cache_reset()  # 重置前端处理器
                self.frontend.reset_status()
                self.vad_scorer.AllResetDetection()
                valid_segments = []
                self.in_cache0 = []
                self.in_cache1 = []
                self.in_cache2 = []
                self.in_cache3 = []
                self.out_cache0 = []
                self.out_cache1 = []
                self.out_cache2 = []
                self.out_cache3 = []


            waveforms=np.expand_dims(waveform,axis=0)

            # 初始化/清空缓存
            if seq_start :
                self.frontend.cache_reset()  # 重置前端处理器
                self.frontend.reset_status()
                self.vad_scorer.AllResetDetection()
                valid_segments = []
                self.in_cache0 = []
                self.in_cache1 = []
                self.in_cache2 = []
                self.in_cache3 = []
                self.out_cache0 = []
                self.out_cache1 = []
                self.out_cache2 = []
                self.out_cache3 = []

            # is_final=False
            is_final=bool(seq_end)  # 如果seq_end为True，则表示当前请求是最后一个请求，否则不是
            feats,feats_length=self.extract_feat(waveforms,is_final)     # 这里有问题
        
            if feats.size != 0:
                in_cache0 = self.prepare_cache(self.in_cache0)
                in_cache1 = self.prepare_cache(self.in_cache1)
                in_cache2 = self.prepare_cache(self.in_cache2)
                in_cache3 = self.prepare_cache(self.in_cache3)
                inputs = [
                    pb_utils.Tensor("speech",feats.astype(np.float32)),
                    pb_utils.Tensor("in_cache0",in_cache0.astype(np.float32)),
                    pb_utils.Tensor("in_cache1",in_cache1.astype(np.float32)),
                    pb_utils.Tensor("in_cache2",in_cache2.astype(np.float32)),
                    pb_utils.Tensor("in_cache3",in_cache3.astype(np.float32)),
                ]

                infer_request=pb_utils.InferenceRequest(
                    model_name="asr_fsmn-vad_streaming",
                    requested_output_names=["logits","out_cache0","out_cache1","out_cache2","out_cache3"],
                    inputs=inputs
                )

                response = infer_request.exec()

                segments = []
                segments_cache = []
                if response.has_error():
                    print(f"Inference error: {response.error().message()}")
                else:
                    logits = pb_utils.get_output_tensor_by_name(response, "logits").as_numpy()
                    out_cache0 = pb_utils.get_output_tensor_by_name(response, "out_cache0").as_numpy()
                    out_cache1 = pb_utils.get_output_tensor_by_name(response, "out_cache1").as_numpy()
                    out_cache2 = pb_utils.get_output_tensor_by_name(response, "out_cache2").as_numpy()
                    out_cache3 = pb_utils.get_output_tensor_by_name(response, "out_cache3").as_numpy()
                
                    # 更新缓存
                    self.in_cache0=out_cache0
                    self.in_cache1=out_cache1
                    self.in_cache2=out_cache2
                    self.in_cache3=out_cache3

                
                    # 获取音频数据准备打分
                    waveforms = self.frontend.get_waveforms()
                    segments = self.vad_scorer(
                        logits, waveforms, is_final=is_final, max_end_sil=self.max_end_sil, online=True
                    )

                    valid_segments = []
                    if segments and len(segments[0]) > 0 and segments[0][0]:    # 避免空段
                    # 提取有效语音段
                        for segment in segments[0]:  # batch_size=1
                            if len(segment) == 2:
                                start, end = segment
                                valid_segments.extend([start, end])

                    else:
                        valid_segments.extend([])

                    # 添加进缓存
                    segments_cache.extend(valid_segments)

                    segments_cache_return = segments_cache

            output_data = np.array(valid_segments, dtype=np.int32)  # shape (1, 2)                output_data = np.array([70, -1], dtype=np.int32)  # shape (1, 2)
            out_tensor = pb_utils.Tensor("OUTPUT0", output_data)
            response = pb_utils.InferenceResponse(output_tensors=[out_tensor])
            responses.append(response)
        return responses



        # response=[70,-1]
        # response = []
        # responses = []

        # for request in requests:
        #     # 模拟推理输出
        #     output_data = np.array(valid_segments, dtype=np.int32)  # shape (1, 2)                output_data = np.array([70, -1], dtype=np.int32)  # shape (1, 2)
        #     out_tensor = pb_utils.Tensor("OUTPUT0", output_data)
        #     response = pb_utils.InferenceResponse(output_tensors=[out_tensor])
        #     responses.append(response)

        # return responses

    def finalize(self):
        """`finalize` is called only once when the model is being unloaded.
        Implementing `finalize` function is optional. This function allows
        the model to perform any necessary clean ups before exit.
        """
        # del self.punc_model   # 释放加载好的punc的onnx模型
        print("Cleaning up...")
