# transforms.py
from typing import List, Optional
import random

import torch
from torch import Tensor
import torchvision
import torchvision.transforms as T


__torchvision_need_compat_flag = float(torchvision.__version__.split(".")[1]) < 7
if __torchvision_need_compat_flag:
    from torchvision.ops import _new_empty_tensor
    from torchvision.ops.misc import _output_size

import torchvision.transforms.functional as F


def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
    """
    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
    This will eventually be supported natively by PyTorch, and this
    class can go away.
    """
    if __torchvision_need_compat_flag < 0.7:
        if input.numel() > 0:
            return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners)

        output_shape = _output_size(2, input, size, scale_factor)
        output_shape = list(input.shape[:-2]) + list(output_shape)
        return _new_empty_tensor(input, output_shape)
    else:
        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)

def resize(image, target, size, max_size=None):
    # size can be min_size (scalar) or (w, h) tuple

    def get_size_with_aspect_ratio(image_size, size, max_size=None):
        w, h = image_size
        if max_size is not None:
            min_original_size = float(min((w, h)))
            max_original_size = float(max((w, h)))
            if max_original_size / min_original_size * size > max_size:
                size = int(round(max_size * min_original_size / max_original_size))

        if (w <= h and w == size) or (h <= w and h == size):
            return (h, w)

        if w < h:
            ow = size
            oh = int(size * h / w)
        else:
            oh = size
            ow = int(size * w / h)

        return (oh, ow)

    def get_size(image_size, size, max_size=None):
        if isinstance(size, (list, tuple)):
            return size[::-1]
        else:
            return get_size_with_aspect_ratio(image_size, size, max_size)

    size = get_size(image.size, size, max_size)
    rescaled_image = F.resize(image, size)

    if target is None:
        return rescaled_image, None

    ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(rescaled_image.size, image.size))
    ratio_width, ratio_height = ratios

    target = target.copy()
    if "boxes" in target:
        boxes = target["boxes"]
        scaled_boxes = boxes * torch.as_tensor(
            [ratio_width, ratio_height, ratio_width, ratio_height]
        )
        target["boxes"] = scaled_boxes

    if "area" in target:
        area = target["area"]
        scaled_area = area * (ratio_width * ratio_height)
        target["area"] = scaled_area

    h, w = size
    target["size"] = torch.tensor([h, w])

    if "masks" in target:
        target["masks"] = (
            interpolate(target["masks"][:, None].float(), size, mode="nearest")[:, 0] > 0.5
        )

    return rescaled_image, target

class RandomResize(object):
    def __init__(self, sizes, max_size=None):
        assert isinstance(sizes, (list, tuple))
        self.sizes = sizes
        self.max_size = max_size

    def __call__(self, img, target=None):
        size = random.choice(self.sizes)
        return resize(img, target, size, self.max_size)

class ToTensor(object):
    def __call__(self, img, target):
        return F.to_tensor(img), target


def box_xyxy_to_cxcywh(x):
    x0, y0, x1, y1 = x.unbind(-1)
    b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)]
    return torch.stack(b, dim=-1)


class Normalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, target=None):
        image = F.normalize(image, mean=self.mean, std=self.std)
        if target is None:
            return image, None
        target = target.copy()
        h, w = image.shape[-2:]
        if "boxes" in target:
            boxes = target["boxes"]
            boxes = box_xyxy_to_cxcywh(boxes)
            boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
            target["boxes"] = boxes
        return image, target

# acl_infer.py
from typing import Dict, List
import acl
import numpy as np
from threading import Lock

from PIL import Image
import torch
import torch.nn.functional as F
import transforms as T

ACL_MEM_MALLOC_HUGE_FIRST = 0
ACL_MEMCPY_HOST_TO_DEVICE = 1
ACL_MEMCPY_DEVICE_TO_HOST = 2

def check_ret(message, ret_int):
    """Check int value is 0 or not
    Args:
        message: output log str
        ret_int: check value that type is int
    """
    if ret_int != 0:
        raise Exception("{} failed ret_int={}".format(message, ret_int))

class ACLEngine(object):
    def __init__(self, model_path, device_id=0):
        self._lock = Lock()

        self._device_id = device_id
        # self._context = None
        # self._stream = None
        self._model_id = None
        self._model_desc = None

        self._input_num = 0
        self._output_num = 0
        self._input_dataset = None
        self._output_dataset = None
        self._input_data:List[Dict[str,]] = []
        self._output_data:List[Dict[str,]] =  []
        
        print('Start init resource')
        self.initResource()
        self.loadModel(model_path)
        self.allocateMem()
        print('Init resource success')

    def inference(self, inputs:List[np.ndarray]) -> List[np.ndarray]:
        if len(inputs) != self._input_num:
            print("Current input data num %d unequal to model "
                      "input num %d" % (len(inputs), self._input_num))
            return None
        # self._lock.acquire()
        # 遍历所有输入，拷贝到对应的buffer内存中
        for i, input_data in enumerate(inputs):
            bytes_data = input_data.tobytes()
            bytes_ptr = acl.util.bytes_to_ptr(bytes_data)
            ret = acl.rt.memcpy(self._input_data[i]["buffer"],      # 目标地址 device 
                                self._input_data[i]["size"],        # 目标地址大小
                                bytes_ptr,                          # 源地址 host 
                                self._input_data[i]["size"],        # 源地址大小
                                ACL_MEMCPY_HOST_TO_DEVICE)          # 模式:从host到device
            check_ret("acl.rt.memcpy", ret)
        print('Start Inference')
        # 执行模型推理。
        ret = acl.mdl.execute(self._model_id, self._input_dataset, self._output_dataset)
        check_ret("acl.rt.execute", ret)
        print(f'End Inference ret')
        
        inference_result = []
        for i, item in enumerate(self._output_data):
            # 将推理输出数据从Device传输到Host
            ret = acl.rt.memcpy(item["buffer_host"],                # 目标地址 host
                                item["size"],                       # 目标地址大小
                                item["buffer"],                     # 源地址 device 
                                item["size"],                       # 源地址大小
                                ACL_MEMCPY_DEVICE_TO_HOST)          # 模式：从device到host 
            check_ret("output acl.rt.memcpy ACL_MEMCPY_DEVICE_TO_HOST ret =", ret)
            # 从内存地址获取bytes对象
            bytes_out = acl.util.ptr_to_bytes(item["buffer_host"], item["size"])
            # 按照float32格式将数据转为numpy数组 
            data = np.frombuffer(bytes_out, dtype=np.float32)
            inference_result.append(data)
        # self._lock.release()
        return inference_result

    def initResource(self):
        # step1: 初始化
        ret = acl.init()
        check_ret("acl.init", ret)
        # 指定运算的Device 
        ret = acl.rt.set_device(self._device_id)
        check_ret("acl.rt.set_device", ret)
        # self._context, ret = acl.rt.create_context(self._device_id)
        # check_ret("acl.rt.create_context", ret)
        # self._stream, ret = acl.rt.create_stream()
        # check_ret("acl.rt.create_stream", ret)

    def loadModel(self, model_path):
        # step2: 加载模型
        # 加载离线模型文件，返回标识模型的ID 
        self._model_id, ret = acl.mdl.load_from_file(model_path)
        check_ret("acl.mdl.load_from_file", ret)
        # 创建空白模型描述信息，获取模型描述信息的指针地址
        self._model_desc = acl.mdl.create_desc()
        # 通过模型的ID，将模型的描述信息填充到model_desc
        ret = acl.mdl.get_desc(self._model_desc, self._model_id)
        check_ret("acl.rt.get_desc", ret)
        # 获得模型输入的个数 
        self._input_num = acl.mdl.get_num_inputs(self._model_desc)
        # 获得模型输出的个数 
        self._output_num = acl.mdl.get_num_outputs(self._model_desc)

    def allocateMem(self):
        # step3：创建输入输出数据集
        # 创建aclmdlDataset类型的数据，描述模型推理的输入
        self._input_dataset = acl.mdl.create_dataset()
        self._input_data = []
        for i in range(self._input_num):
            # 获取所需的buffer内存大小 
            buffer_size = acl.mdl.get_input_size_by_index(self._model_desc, i)
            # 申请buffer内存 
            buffer, ret = acl.rt.malloc(buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)
            check_ret("acl.rt.malloc", ret)
            # 从内存创建buffer数据 
            data_buffer = acl.create_data_buffer(buffer, buffer_size)
            # 将buffer数据添加到数据集 
            _, ret = acl.mdl.add_dataset_buffer(self._input_dataset, data_buffer)
            check_ret("acl.rt.add_dataset_buffer", ret)
            self._input_data.append({"buffer": buffer, "size": buffer_size})

        # 创建aclmdlDataset类型的数据，描述模型推理的输出
        self._output_dataset = acl.mdl.create_dataset()
        self._output_data = []
        for i in range(self._output_num):
            # 获取所需的buffer内存大小 
            buffer_size = acl.mdl.get_output_size_by_index(self._model_desc, i)
            # 申请buffer内存 
            buffer, ret = acl.rt.malloc(buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)
            check_ret("acl.rt.malloc", ret)
            # 从内存创建buffer数据
            data_buffer = acl.create_data_buffer(buffer, buffer_size)
            # 将buffer数据添加到数据集 
            _, ret = acl.mdl.add_dataset_buffer(self._output_dataset, data_buffer)
            check_ret("acl.rt.add_dataset_buffer", ret)
            buffer_host, ret = acl.rt.malloc_host(buffer_size)
            check_ret("acl.rt.malloc_host", ret)
            self._output_data.append({"buffer": buffer, "size": buffer_size, "buffer_host": buffer_host})

    def releaseResource(self):
        # 析构函数 按照初始化资源的相反顺序释放资源。
        print('Start Resource destroyed')
        self.freeMem()
        self.unloadModel()
        self.destroyResource()
        print('Resource destroyed successfully')

    def freeMem(self):
        # 释放输入资源，包括数据结构和内存。
        while self._input_data:
            item = self._input_data.pop()
            ret = acl.rt.free(item["buffer"])                   # 释放buffer内存 
            check_ret("acl.rt.free", ret)
        for i in range(self._input_num):
            data_buf = acl.mdl.get_dataset_buffer(self._input_dataset, i)
            if data_buf:
                ret = acl.destroy_data_buffer(data_buf)
                check_ret("acl.destroy_data_buffer", ret)
        ret = acl.mdl.destroy_dataset(self._input_dataset)      # 销毁输入数据集 
        check_ret("acl.mdl.destroy_dataset", ret)
        # 释放输出资源，包括数据结构和内存。
        while self._output_data:
            item = self._output_data.pop()
            ret = acl.rt.free_host(item["buffer_host"])
            check_ret("acl.rt.free_host", ret)
            ret = acl.rt.free(item["buffer"])                   # 释放buffer内存 
            check_ret("acl.rt.free", ret)
        for i in range(self._output_num):
            data_buf = acl.mdl.get_dataset_buffer(self._output_dataset, i)
            if data_buf:
                ret = acl.destroy_data_buffer(data_buf)
                check_ret("acl.destroy_data_buffer", ret)
        ret = acl.mdl.destroy_dataset(self._output_dataset)     # 销毁输出数据集 
        check_ret("acl.mdl.destroy_dataset", ret)

    def unloadModel(self):
        self._input_num = self._output_num = 0
        # 销毁模型描述 
        if self._model_desc:
            ret  = acl.mdl.destroy_desc(self._model_desc)
            check_ret("acl.mdl.destroy_desc", ret)
            self._model_desc = None
        # 卸载模型 
        ret = acl.mdl.unload(self._model_id)
        check_ret("acl.mdl.unload", ret)

        # if self._stream:
        #     ret = acl.rt.destroy_stream(self._stream)
        #     check_ret("acl.rt.destroy_stream", ret)
        # if self._context:
        #     ret = acl.rt.destroy_context(self._context)
        #     check_ret("acl.rt.destroy_context", ret)
        #     self.cont_contextext = None

    def destroyResource(self):
        # 释放device 
        ret = acl.rt.reset_device(self._device_id)
        check_ret("acl.rt.reset_device", ret)
        # acl去初始化
        ret = acl.finalize()
        check_ret("acl.finalize", ret)


def load_image(image_path, size=1280):
    # load image
    image_pil = Image.open(image_path).convert("RGB")  # load image

    transform_resize = T.RandomResize([size], max_size=size)
    transform_tensor = T.Compose(
        [
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    image_pil, _ = transform_resize(image_pil, None)  # PIL
    image, _ = transform_tensor(image_pil, None) # 3, H, W

    h, w = image.size(1), image.size(2)
    assert h <= size and w <= size

    # pad image
    padded_image = F.pad(image, (0, size-w, 0, size-h))

    # mask
    mask = torch.zeros((1, h, w))
    padded_mask = F.pad(mask, (0, size-w, 0, size-h), value=1)

    return image_pil, padded_image, padded_mask.bool()


if __name__ == "__main__":
    # phrase
    phrases = ['cat', 'dog']

    # text embedding
    text_embedding_dict_file = '/root/gdino/models/embeddings.bin'
    print('Read text embedding from ' + text_embedding_dict_file)
    text_embedding = torch.load(text_embedding_dict_file, map_location='cpu')
    caption = [text_embedding[x] for x in phrases]
    caption = torch.stack(caption, dim=0).to(torch.float32)

    image_pil, image, mask = load_image("./.asset/cats.png")
    
    model_path = '/root/gdino/models/7.0.0/engine.om'
    model = ACLEngine(model_path, device_id=1)

    output = model.inference([image[None].numpy(), mask.numpy(), caption.numpy()])
    if output != None:
        for i, out in enumerate(output):
            print("{}\t{}".format(i, out))
    model.releaseResource()

    del model