class Model(ABC):
    def __init__(self, model_path):
        print(f"load model {model_path}")
        self.model_path = model_path  # 模型路径
        self.model_id = None  # 模型 id
        self.input_dataset = None  # 输入数据结构
        self.output_dataset = None  # 输出数据结构
        self.model_desc = None  # 模型描述信息
        self._input_num = 0  # 输入数据个数
        self._output_num = 0  # 输出数据个数
        self._output_info = []  # 输出信息列表
        self._is_released = False  # 资源是否被释放
        self._init_resource()

    def _init_resource(self):
        ''' 初始化模型、输出相关资源。相关数据类型: aclmdlDesc aclDataBuffer aclmdlDataset'''
        print("Init model resource")
        # 加载模型文件
        self.model_id, ret = acl.mdl.load_from_file(self.model_path)  # 加载模型
        self.model_desc = acl.mdl.create_desc()  # 初始化模型信息对象
        ret = acl.mdl.get_desc(self.model_desc, self.model_id)  # 根据模型获取描述信息
        print("[Model] Model init resource stage success")

        # 创建模型输出 dataset 结构
        self._gen_output_dataset()  # 创建模型输出dataset结构

    def _gen_output_dataset(self):
        ''' 组织输出数据的dataset结构 '''
        ret = SUCCESS
        self._output_num = acl.mdl.get_num_outputs(self.model_desc)  # 获取模型输出个数
        self.output_dataset = acl.mdl.create_dataset()  # 创建输出dataset结构
        for i in range(self._output_num):
            temp_buffer_size = acl.mdl.get_output_size_by_index(self.model_desc, i)  # 获取模型输出个数
            temp_buffer, ret = acl.rt.malloc(temp_buffer_size, ACL_MEM_MALLOC_NORMAL_ONLY)  # 为每个输出申请device内存
            dataset_buffer = acl.create_data_buffer(temp_buffer, temp_buffer_size)  # 创建输出的data buffer结构,将申请的内存填入data buffer
            _, ret = acl.mdl.add_dataset_buffer(self.output_dataset, dataset_buffer)  # 将 data buffer 加入输出dataset

        if ret == FAILED:
            self._release_dataset(self.output_dataset)   # 失败时释放dataset
        print("[Model] create model output dataset success")

    def _gen_input_dataset(self, input_list):
        ''' 组织输入数据的dataset结构 '''
        ret = SUCCESS
        self._input_num = acl.mdl.get_num_inputs(self.model_desc)  # 获取模型输入个数
        self.input_dataset = acl.mdl.create_dataset()  # 创建输入dataset结构
        for i in range(self._input_num):
            item = input_list[i]  # 获取第 i 个输入数据
            data_ptr = acl.util.bytes_to_ptr(item.tobytes())  # 获取输入数据字节流
            size = item.size * item.itemsize  # 获取输入数据字节数
            dataset_buffer = acl.create_data_buffer(data_ptr, size)  # 创建输入dataset buffer结构, 填入输入数据
            _, ret = acl.mdl.add_dataset_buffer(self.input_dataset, dataset_buffer)  # 将dataset buffer加入dataset

        if ret == FAILED:
            self._release_dataset(self.input_dataset)  # 失败时释放dataset
        print("[Model] create model input dataset success")

    def _unpack_bytes_array(self, byte_array, shape, datatype):
        ''' 将内存不同类型的数据解码为numpy数组 '''
        np_type = None

        # 获取输出数据类型对应的numpy数组类型和解码标记
        if datatype == 0:  # ACL_FLOAT
            np_type = np.float32
        elif datatype == 1:  # ACL_FLOAT16
            np_type = np.float16
        elif datatype == 3:  # ACL_INT32
            np_type = np.int32
        elif datatype == 8:  # ACL_UINT32
            np_type = np.uint32
        else:
            print("unsurpport datatype ", datatype)
            return

        # 将解码后的数据组织为numpy数组,并设置shape和类型
        return np.frombuffer(byte_array, dtype=np_type).reshape(shape)

    def _output_dataset_to_numpy(self):
        ''' 将模型输出解码为numpy数组 '''
        dataset = []
        # 遍历每个输出
        for i in range(self._output_num):
            buffer = acl.mdl.get_dataset_buffer(self.output_dataset, i)  # 从输出dataset中获取buffer
            data_ptr = acl.get_data_buffer_addr(buffer)  # 获取输出数据内存地址
            size = acl.get_data_buffer_size(buffer)  # 获取输出数据字节数
            narray = acl.util.ptr_to_bytes(data_ptr, size)  # 将指针转为字节流数据

            # 根据模型输出的shape和数据类型,将内存数据解码为numpy数组
            dims = acl.mdl.get_output_dims(self.model_desc, i)[0]["dims"]  # 获取每个输出的维度
            datatype = acl.mdl.get_output_data_type(self.model_desc, i)  # 获取每个输出的数据类型
            output_nparray = self._unpack_bytes_array(narray, tuple(dims), datatype)  # 解码为numpy数组
            dataset.append(output_nparray)
        return dataset

    def execute(self, input_list):
        '''创建输入dataset对象, 推理完成后, 将输出数据转换为numpy格式'''
        self._gen_input_dataset(input_list)  # 创建模型输入dataset结构
        ret = acl.mdl.execute(self.model_id, self.input_dataset, self.output_dataset)  # 调用离线模型的execute推理数据
        out_numpy = self._output_dataset_to_numpy()  # 将推理输出的二进制数据流解码为numpy数组, 数组的shape和类型与模型输出规格一致
        return out_numpy

    def release(self):
        ''' 释放模型相关资源 '''
        if self._is_released:
            return

        print("Model start release...")
        self._release_dataset(self.input_dataset)  # 释放输入数据结构
        self.input_dataset = None  # 将输入数据置空
        self._release_dataset(self.output_dataset)  # 释放输出数据结构
        self.output_dataset = None  # 将输出数据置空

        if self.model_id:
            ret = acl.mdl.unload(self.model_id)  # 卸载模型
        if self.model_desc:
            ret = acl.mdl.destroy_desc(self.model_desc)  # 释放模型描述信息
        self._is_released = True
        print("Model release source success")

    def _release_dataset(self, dataset):
        ''' 释放 aclmdlDataset 类型数据 '''
        if not dataset:
            return
        num = acl.mdl.get_dataset_num_buffers(dataset)  # 获取数据集包含的buffer个数
        for i in range(num):
            data_buf = acl.mdl.get_dataset_buffer(dataset, i)  # 获取buffer指针
            if data_buf:
                ret = acl.destroy_data_buffer(data_buf)  # 释放buffer
        ret = acl.mdl.destroy_dataset(dataset)  # 销毁数据集

    @abstractmethod
    def infer(self, inputs): # 保留接口, 子类必须重写
        pass

class YOWO2(Model):
    def __init__(self, model_path):
        super().__init__(model_path)
    
    @torch.no_grad()
    def infer(self, args, transform, class_names, class_colors):
        # path to save 
        save_path = os.path.join(args.save_folder, 'demo', 'videos')
        os.makedirs(save_path, exist_ok=True)

        # path to video
        path_to_video = os.path.join(args.video)

        # video
        video = cv2.VideoCapture(path_to_video)
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        save_size = (960, 720)
        save_name = os.path.join(save_path, 'detection.avi')
        fps = 20.0
        out = cv2.VideoWriter(save_name, fourcc, fps, save_size)

        
        # run
        video_clip = []
        image_list = []
        while(True):
            ret, frame = video.read()
            if ret:
                # to RGB
                frame_rgb = frame[..., (2, 1, 0)]

                # to PIL image
                frame_pil = Image.fromarray(frame_rgb.astype(np.uint8))

                # prepare
                if len(video_clip) <= 0:
                    for _ in range(args.len_clip):
                        video_clip.append(frame_pil)

                video_clip.append(frame_pil)
                del video_clip[0]

                # orig size
                orig_h, orig_w = frame.shape[:2]

                # transform
                x, _ = transform(video_clip)
                # List [T, 3, H, W] -> [3, T, H, W]
                x = torch.stack(x, dim=1)
                x = x.unsqueeze(0) # [B, 3, T, H, W], B=1
                
                x = np.asarray(x)
                # print('x.shape', x.shape, 'x.mean', np.mean(x))
                
                t0 = time.time()
                # inference
                outputs = self.execute([x, ])

                print("inference time ", time.time() - t0, "s")

                # vis detection results
                if args.dataset in ['ava_v2.2']:
                    batch_bboxes = outputs
                    # batch size = 1
                    bboxes = batch_bboxes[0]
                    # multi hot
                    frame = multi_hot_vis(
                        args=args,
                        frame=frame,
                        out_bboxes=bboxes,
                        orig_w=orig_w,
                        orig_h=orig_h,
                        class_names=class_names,
                        act_pose=args.pose
                        )
                elif args.dataset in ['ucf24']:
                    batch_scores, batch_labels, batch_bboxes = outputs
                    # batch size = 1
                    scores = batch_scores[0]
                    labels = batch_labels[0]
                    bboxes = batch_bboxes[0]
                    # rescale
                    bboxes = rescale_bboxes(bboxes, [orig_w, orig_h])
                    # one hot
                    frame = vis_detection(
                        frame=frame,
                        scores=scores,
                        labels=labels,
                        bboxes=bboxes,
                        vis_thresh=args.vis_thresh,
                        class_names=class_names,
                        class_colors=class_colors
                        )
                # save
                frame_resized = cv2.resize(frame, save_size)
                cv2.imwrite('./det_result.png', frame_resized)
                # out.write(frame_resized)

                if args.gif:
                    gif_resized = cv2.resize(frame, (200, 150))
                    gif_resized_rgb = gif_resized[..., (2, 1, 0)]
                    image_list.append(gif_resized_rgb)

                if args.show:
                    # show
                    cv2.imshow('key-frame detection', frame)
                    cv2.waitKey(1)

            else:
                break

        video.release()
        out.release()
        cv2.destroyAllWindows()