import ctypes
import time
import numpy as np
import cv2

# === DLL 加载 ===
dll = ctypes.WinDLL(r"D:\Trae_space\BodyBalanceEvaluation\backend\tests\SMiTSenseUsbWrapper.dll")

dll.SMiTSenseUsb_Init.argtypes = [ctypes.c_int]
dll.SMiTSenseUsb_Init.restype = ctypes.c_int

dll.SMiTSenseUsb_ScanDevices.argtypes = [ctypes.POINTER(ctypes.c_int)]
dll.SMiTSenseUsb_ScanDevices.restype = ctypes.c_int

dll.SMiTSenseUsb_OpenAndStart.argtypes = [
    ctypes.c_int,
    ctypes.POINTER(ctypes.c_uint16),
    ctypes.POINTER(ctypes.c_uint16)
]
dll.SMiTSenseUsb_OpenAndStart.restype = ctypes.c_int

dll.SMiTSenseUsb_GetLatestFrame.argtypes = [
    ctypes.POINTER(ctypes.c_uint16),
    ctypes.c_int
]
dll.SMiTSenseUsb_GetLatestFrame.restype = ctypes.c_int

dll.SMiTSenseUsb_StopAndClose.argtypes = []
dll.SMiTSenseUsb_StopAndClose.restype = ctypes.c_int

# === 初始化设备 ===
ret = dll.SMiTSenseUsb_Init(0)
if ret != 0:
    raise RuntimeError(f"Init failed: {ret}")

count = ctypes.c_int()
ret = dll.SMiTSenseUsb_ScanDevices(ctypes.byref(count))
if ret != 0 or count.value == 0:
    raise RuntimeError("No devices found")

rows = ctypes.c_uint16()
cols = ctypes.c_uint16()
ret = dll.SMiTSenseUsb_OpenAndStart(0, ctypes.byref(rows), ctypes.byref(cols))
if ret != 0:
    raise RuntimeError("OpenAndStart failed")

rows, cols = rows.value, cols.value
frame_size = rows * cols
buf_type = ctypes.c_uint16 * frame_size
buf = buf_type()

print(f"实时采集: {rows} 行 x {cols} 列")

# 固定映射范围（和原来 matplotlib 一样）
vmin, vmax = 0, 1000

try:
    while True:
        ret = dll.SMiTSenseUsb_GetLatestFrame(buf, frame_size)
        if ret == 0:
            data = np.frombuffer(buf, dtype=np.uint16).reshape((rows, cols))

            # 固定范围映射到 0-255
            norm_data = np.clip((data - vmin) / (vmax - vmin) * 255, 0, 255).astype(np.uint8)

            # 颜色映射（jet）
            heatmap = cv2.applyColorMap(norm_data, cv2.COLORMAP_JET)

            # 放大显示（保持 nearest 效果）
            heatmap = cv2.resize(heatmap, (cols * 4, rows * 4), interpolation=cv2.INTER_NEAREST)

            # 显示
            cv2.imshow("SMiTSense 足底压力分布", heatmap)

            if cv2.waitKey(1) & 0xFF == 27:  # ESC 退出
                break
        else:
            print("读取数据帧失败")
        time.sleep(0.05)  # 20 FPS
except KeyboardInterrupt:
    pass
finally:
    dll.SMiTSenseUsb_StopAndClose()
    cv2.destroyAllWindows()
    print("设备已关闭")