import ctypes
import numpy as np
import cv2
import sys
import time

# ------------------- DLL 加载 -------------------
dll = ctypes.CDLL(r"D:\BodyBalanceEvaluation\backend\tests\Wrapper.dll")

# 定义 C 结构体
class FPMS_DEVICE_INFO(ctypes.Structure):
    _fields_ = [
        ("mn", ctypes.c_uint16),
        ("sn", ctypes.c_char * 64),
        ("fwVersion", ctypes.c_uint16),
        ("protoVer", ctypes.c_uint8),
        ("pid", ctypes.c_uint16),
        ("vid", ctypes.c_uint16),
        ("rows", ctypes.c_uint16),
        ("cols", ctypes.c_uint16),
    ]

# 函数声明
dll.fpms_usb_init_wrap.argtypes = [ctypes.c_int]
dll.fpms_usb_init_wrap.restype = ctypes.c_int

dll.fpms_usb_get_device_list_wrap.argtypes = [ctypes.POINTER(FPMS_DEVICE_INFO), ctypes.c_int, ctypes.POINTER(ctypes.c_int)]
dll.fpms_usb_get_device_list_wrap.restype = ctypes.c_int

dll.fpms_usb_open_wrap.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_uint64)]
dll.fpms_usb_open_wrap.restype = ctypes.c_int

dll.fpms_usb_read_frame_wrap.argtypes = [ctypes.c_uint64, ctypes.POINTER(ctypes.c_uint16), ctypes.c_size_t]
dll.fpms_usb_read_frame_wrap.restype = ctypes.c_int

dll.fpms_usb_close_wrap.argtypes = [ctypes.c_uint64]
dll.fpms_usb_close_wrap.restype = ctypes.c_int

# ------------------- 初始化 -------------------
if dll.fpms_usb_init_wrap(0) != 0:
    print("初始化失败")
    sys.exit(1)

# 获取设备列表
count = ctypes.c_int()
devs = (FPMS_DEVICE_INFO * 10)()
r = dll.fpms_usb_get_device_list_wrap(devs, 10, ctypes.byref(count))
if r != 0 or count.value == 0:
    print("未检测到设备")
    sys.exit(1)

print("检测到设备数量:", count.value)
dev = devs[0]
rows, cols = dev.rows, dev.cols
print(f"使用设备 SN={dev.sn.decode(errors='ignore')}  {rows}x{cols}")

# 打开设备
handle = ctypes.c_uint64()
r = dll.fpms_usb_open_wrap(0, ctypes.byref(handle))
if r != 0:
    print("设备打开失败")
    sys.exit(1)
print("设备已打开, 句柄 =", handle.value)

# ------------------- 循环读取帧数据 -------------------
buf_len = rows * cols
FrameArray = (ctypes.c_uint16 * buf_len)()

vmin, vmax = 100, 1000   # 根据实际传感器数据范围调整

try:
    while True:
        r = dll.fpms_usb_read_frame_wrap(handle.value, FrameArray, buf_len)
        if r != 0:
            print("读取帧失败, code=", r)
            time.sleep(0.05)
            continue

        # 转 numpy
        data = np.frombuffer(FrameArray, dtype=np.uint16).reshape((rows, cols))

        # ------------------- 方法1: 固定线性映射 -------------------
        norm1 = np.clip((data - vmin) / (vmax - vmin) * 255, 0, 255).astype(np.uint8)
        heatmap1 = cv2.applyColorMap(norm1, cv2.COLORMAP_JET)
        heatmap1[data <= vmin] = (0,0,0)
        heatmap1 = cv2.resize(heatmap1, (cols*4, rows*4), interpolation=cv2.INTER_NEAREST)
        cv2.imshow("方法1: 固定线性", heatmap1)

        # ------------------- 方法2: 自适应归一化 -------------------
        dmin, dmax = np.min(data), np.max(data)
        norm2 = np.clip((data - dmin) / max(dmax - dmin, 1) * 255, 0, 255).astype(np.uint8)
        heatmap2 = cv2.applyColorMap(norm2, cv2.COLORMAP_JET)
        heatmap2[data <= vmin] = (0,0,0)
        heatmap2 = cv2.resize(heatmap2, (cols*4, rows*4), interpolation=cv2.INTER_NEAREST)
        cv2.imshow("方法2: 自适应归一化", heatmap2)

        # ------------------- 方法3: 对数映射 -------------------
        log_data = np.log1p(data - vmin)
        log_max = np.max(log_data)
        norm3 = np.clip(log_data / max(log_max, 1) * 255, 0, 255).astype(np.uint8)
        heatmap3 = cv2.applyColorMap(norm3, cv2.COLORMAP_JET)
        heatmap3[data <= vmin] = (0,0,0)
        heatmap3 = cv2.resize(heatmap3, (cols*4, rows*4), interpolation=cv2.INTER_NEAREST)
        cv2.imshow("方法3: 对数映射", heatmap3)

        # ------------------- 方法4: 平方根映射 -------------------
        sqrt_data = np.sqrt(np.clip(data - vmin, 0, None))
        sqrt_max = np.max(sqrt_data)
        norm4 = np.clip(sqrt_data / max(sqrt_max, 1) * 255, 0, 255).astype(np.uint8)
        heatmap4 = cv2.applyColorMap(norm4, cv2.COLORMAP_JET)
        heatmap4[data <= vmin] = (0,0,0)
        heatmap4 = cv2.resize(heatmap4, (cols*4, rows*4), interpolation=cv2.INTER_NEAREST)
        cv2.imshow("方法4: 平方根映射", heatmap4)

        # ------------------- 方法5: 分段颜色映射 -------------------
        heatmap5 = np.zeros((rows, cols, 3), dtype=np.uint8)
        # 定义压力区间和对应颜色 (B,G,R)
        bins = [vmin, 250, 400, 600, 800, vmax]
        colors = [
            (0,0,0),
            (0,0,255),     # 蓝
            (0,255,0),     # 绿
            (0,255,255),   # 黄
            (255,0,0),     # 红
        ]
        for i in range(len(bins)-1):
            mask = (data > bins[i]) & (data <= bins[i+1])
            heatmap5[mask] = colors[i]
        heatmap5 = cv2.resize(heatmap5, (cols*4, rows*4), interpolation=cv2.INTER_NEAREST)
        cv2.imshow("方法5: 分段映射", heatmap5)

        # ------------------- 退出 -------------------
        if cv2.waitKey(1) & 0xFF == 27:  # ESC 退出
            break

except KeyboardInterrupt:
    pass

# ------------------- 清理 -------------------
dll.fpms_usb_close_wrap(handle.value)
cv2.destroyAllWindows()
print("已退出")