import cv2
import numpy as np
import platform
from rknnlite.api import RKNNLite

# ------------------- 全局变量 -------------------
_global_rknn_instance = None

labels = {0: '夹具未夹紧', 1: '夹具夹紧'}

DEVICE_COMPATIBLE_NODE = '/proc/device-tree/compatible'

def get_host():
    system = platform.system()
    machine = platform.machine()
    os_machine = system + '-' + machine
    if os_machine == 'Linux-aarch64':
        try:
            with open(DEVICE_COMPATIBLE_NODE) as f:
                device_compatible_str = f.read()
                if 'rk3562' in device_compatible_str:
                    host = 'RK3562'
                elif 'rk3576' in device_compatible_str:
                    host = 'RK3576'
                elif 'rk3588' in device_compatible_str:
                    host = 'RK3588'
                else:
                    host = 'RK3566_RK3568'
        except IOError:
            print('Read device node {} failed.'.format(DEVICE_COMPATIBLE_NODE))
            exit(-1)
    else:
        host = os_machine
    return host

# ------------------- 图像预处理 -------------------
def preprocess(raw_image, target_size=(640, 640)):
    img = cv2.resize(raw_image, target_size)
    img = np.expand_dims(img, 0)  # 添加 batch 维度
    return img

# ------------------- RKNN 模型初始化 -------------------
def init_rknn_model(model_path):
    global _global_rknn_instance
    if _global_rknn_instance is None:
        rknn_lite = RKNNLite(verbose=False)
        ret = rknn_lite.load_rknn(model_path)
        if ret != 0:
            print(f'[ERROR] Load model failed (code: {ret})')
            exit(ret)
        ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
        if ret != 0:
            print(f'[ERROR] Init runtime failed (code: {ret})')
            exit(ret)
        _global_rknn_instance = rknn_lite
        print(f'[INFO] Model loaded successfully: {model_path}')
    return _global_rknn_instance

# ------------------- 推理 -------------------
def yolov11_cls_inference(model_path, raw_image, target_size=(640, 640)):
    """
    返回：(class_id, boolean)
    类别 0 -> False
    类别 1 -> True
    """
    img = preprocess(raw_image, target_size)
    rknn = init_rknn_model(model_path)
    outputs = rknn.inference([img])

    # 获取类别 ID
    output = outputs[0].reshape(-1)
    class_id = int(np.argmax(output))
    bool_value = True if class_id == 1 else False

    return class_id, bool_value

# ------------------- 测试 -------------------
if __name__ == '__main__':
    image_path = "12.png"
    bgr_image = cv2.imread(image_path)
    if bgr_image is None:
        print(f"Failed to read image from {image_path}")
        exit(-1)

    rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
    class_id, bool_value = yolov11_cls_inference(
        model_path="yolov11_cls.rknn",
        raw_image=rgb_image,
        target_size=(640, 640)
    )
    print(f"类别ID: {class_id}, 布尔值: {bool_value}")