
import pycuda.driver as cuda
import pycuda
import tensorrt as trt
import numpy as np
import time
import os.path as path
import ctypes
import os
import time
import copy
    
    
class NEW_TTS_SYNTH:
    def __init__(self):
        self.ctx = cuda.Device(0).make_context()
        self.stream = cuda.Stream()
        self.TRT_LOGGER = trt.Logger(trt.Logger.ERROR) 
        self.unused_conv1d_dropout = np.ascontiguousarray(np.zeros([64, 512, 300 , 1], dtype=np.float32))
        self._prepare_engine_context()

    
    def destroy(self):
        self.engine = None
        self.ctx.pop()


    def load_engine(self, engine_filepath, trt_logger):
        self.runtime = trt.Runtime(trt_logger)
        with open(engine_filepath, "rb") as f:
            engine = self.runtime.deserialize_cuda_engine(f.read())
        return engine


    def _prepare_engine_context(self):
        trt.init_libnvinfer_plugins(self.TRT_LOGGER, "")
        base_dir = "."

        ctypes.cdll.LoadLibrary(os.path.join(base_dir, "plugin_so_tile.so"))
        ctypes.cdll.LoadLibrary(os.path.join(base_dir, "plugin_so_reverse_part.so"))
        ctypes.cdll.LoadLibrary(os.path.join(base_dir, "plugin_so_mask_zero.so"))
            
        encoder = self.load_engine(os.path.join(base_dir, "trt_encoder_engine_1.trt"), self.TRT_LOGGER)      #// 加载序列化的cuda引擎
        encoder_context = encoder.create_execution_context()
        encoder_context.active_optimization_profile = 0 

        self.engine = encoder
        self.context = encoder_context



    def get_binding_idxs(self, engine: trt.ICudaEngine, profile_index: int):
        num_bindings_per_profile = engine.num_bindings // engine.num_optimization_profiles
        start_binding = profile_index * num_bindings_per_profile
        end_binding = start_binding + num_bindings_per_profile
    
        # Separate input and output binding indices for convenience
        input_binding_idxs = []
        output_binding_idxs = []
        for binding_index in range(start_binding, end_binding):
            if engine.binding_is_input(binding_index):
                input_binding_idxs.append(binding_index)
            else:
                output_binding_idxs.append(binding_index)
    
        return input_binding_idxs, output_binding_idxs


    def is_dynamic(self, shape):
        return any(dim is None or dim < 0 for dim in shape)
    

    def bind_input(self, engine, context, input_binding_idxs, inputs_data):
        host_inputs = []
        for binding_index in input_binding_idxs:
            input_shape = context.get_binding_shape(binding_index)
            input_name = engine.get_binding_name(binding_index)
            if self.is_dynamic(input_shape):
                profile_index = context.active_optimization_profile
                profile_shapes = engine.get_profile_shape(profile_index, binding_index)
                input_data=inputs_data[input_name]
            else:
                input_data=inputs_data[input_name]
    
            host_inputs.append(input_data)
    
        return host_inputs


    def setup_binding_shapes(self, context, host_inputs, input_binding_idxs, output_binding_idxs):

        for host_input, binding_index in zip(host_inputs, input_binding_idxs):
            context.set_binding_shape(binding_index, host_input.shape)
    
        assert context.all_binding_shapes_specified
    
        host_outputs = []
        device_outputs = []
        for binding_index in output_binding_idxs:
            output_shape = context.get_binding_shape(binding_index)
            buffer = np.empty(output_shape, dtype=np.float32)
            host_outputs.append(buffer)
            device_outputs.append(cuda.mem_alloc(buffer.nbytes))
    
        return host_outputs, device_outputs
            
        
    def enc_data_prepare(self, engine, context, input_binding_idxs, output_binding_idxs):

        input_seqs = [[2, 13, 4, 5, 34, 2, 1]]
        enc_dropout = self.unused_conv1d_dropout[:1, :, :7 , :]
        enc_dropout = np.ascontiguousarray(enc_dropout)
        
        input_feed_inputs        = np.asarray(input_seqs, dtype=np.int32)
        input_feed_speaker       = np.asarray([3] * 1,dtype=np.int32)
        input_feed_input_lengths = np.asarray([7], dtype=np.float32)
        input_feed_split_infos   = np.asarray([[7,0,0,0]], dtype=np.int32)
    
        inputs_data = {"inputs:0":input_feed_inputs,
                       "speaker:0":input_feed_speaker,
                       "input_lengths:0":input_feed_input_lengths,
                       "split_infos:0":input_feed_split_infos,
                       "enc_dropout": enc_dropout}
        host_inputs = self.bind_input(engine, context, input_binding_idxs, inputs_data)
        return host_inputs
  
  
    def enc_execute(self):
        engine = self.engine
        context = self.context
        input_binding_idxs, output_binding_idxs = self.get_binding_idxs(engine, context.active_optimization_profile)
        host_inputs = self.enc_data_prepare(engine, context, input_binding_idxs, output_binding_idxs)
        device_inputs = [cuda.mem_alloc(h_input.nbytes) for h_input in host_inputs]
        for h_input, d_input in zip(host_inputs, device_inputs):
            cuda.memcpy_htod(d_input, h_input)            
        host_outputs, device_outputs = self.setup_binding_shapes(context, host_inputs, input_binding_idxs, output_binding_idxs) 
        bindings = device_inputs + device_outputs
        context.execute_v2(bindings)
        
        return device_outputs
    
    
    def synth(self, callback):
        self.ctx.push()
        device_outputs = self.enc_execute()
        print(">> encoder output shape : {}".format(np.shape(device_outputs)))
        callback("a")
        self.ctx.pop()

        
        
