import ctypes
import numpy as np
import matplotlib.pyplot as plt

class FPMS_DEVICE_C(ctypes.Structure):
    _fields_ = [
        ("mn", ctypes.c_uint8),
        ("sn", ctypes.c_char * 32),
        ("swVersion", ctypes.c_char * 32),
        ("rows", ctypes.c_uint16),
        ("cols", ctypes.c_uint16)
    ]

class SMiTSenseSensor:
    def __init__(self, dll_path = r"D:\Trae_space\BodyBalanceEvaluation\backend\tests\SMiTSenseUsbWrapper.dll"):
        self.dll = ctypes.WinDLL(dll_path)

        # 定义函数接口
        self.dll.fpms_usb_init_wrapper.argtypes = [ctypes.c_int]
        self.dll.fpms_usb_init_wrapper.restype = ctypes.c_int

        self.dll.fpms_usb_get_device_list_wrapper.argtypes = [ctypes.POINTER(FPMS_DEVICE_C), ctypes.POINTER(ctypes.c_int)]
        self.dll.fpms_usb_get_device_list_wrapper.restype = ctypes.c_int

        self.dll.fpms_usb_open_wrapper.argtypes = [FPMS_DEVICE_C, ctypes.POINTER(ctypes.c_void_p)]
        self.dll.fpms_usb_open_wrapper.restype = ctypes.c_int

        self.dll.fpms_usb_read_frame_wrapper.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint16)]
        self.dll.fpms_usb_read_frame_wrapper.restype = ctypes.c_int

        self.dll.fpms_usb_config_light_wrapper.argtypes = [ctypes.c_void_p, ctypes.c_uint8, ctypes.c_uint8, ctypes.c_uint8, ctypes.c_uint8]
        self.dll.fpms_usb_config_light_wrapper.restype = ctypes.c_int

        self.dll.fpms_usb_close_wrapper.argtypes = [ctypes.c_void_p]
        self.dll.fpms_usb_close_wrapper.restype = ctypes.c_int

        self.device_handle = ctypes.c_void_p()
        self.device_info = None

    def initialize(self):
        # 参数0是debug标志，参考示例
        return self.dll.fpms_usb_init_wrapper(0) == 0

    def get_device_list(self, max_devices=10):
        devices = (FPMS_DEVICE_C * max_devices)()
        device_count = ctypes.c_int()
        ret = self.dll.fpms_usb_get_device_list_wrapper(devices, ctypes.byref(device_count))
        if ret != 0 or device_count.value == 0:
            return []
        return devices[:device_count.value]

    def open_device(self, device: FPMS_DEVICE_C):
        handle = ctypes.c_void_p()
        ret = self.dll.fpms_usb_open_wrapper(device, ctypes.byref(handle))
        if ret == 0:
            self.device_handle = handle
            self.device_info = device
            return True
        return False

    def config_light(self, r=1, g=56, b=183, a=214):
        if not self.device_handle:
            return False
        ret = self.dll.fpms_usb_config_light_wrapper(self.device_handle, r, g, b, a)
        return ret == 0

    def read_frame(self):
        if not self.device_handle or not self.device_info:
            return None
        rows = self.device_info.rows
        cols = self.device_info.cols
        frame_buffer = (ctypes.c_uint16 * (rows * cols))()
        ret = self.dll.fpms_usb_read_frame_wrapper(self.device_handle, frame_buffer)
        if ret != 0:
            return None
        # 转成 numpy 二维数组方便处理
        data = np.ctypeslib.as_array(frame_buffer)
        return data.reshape((rows, cols))

    def close_device(self):
        if self.device_handle:
            ret = self.dll.fpms_usb_close_wrapper(self.device_handle)
            self.device_handle = None
            self.device_info = None
            return ret == 0
        return True

def main():
    sensor = SMiTSenseSensor()

    if not sensor.initialize():
        print("初始化失败")
        return

    devices = sensor.get_device_list(20)
    if not devices:
        print("未发现设备")
        return
    print(f"发现设备数量: {len(devices)}")
    for idx, dev in enumerate(devices):
        print(f"设备{idx}: mn={dev.mn}, sn={dev.sn.decode(errors='ignore').strip()}, swVersion={dev.swVersion.decode(errors='ignore').strip()}, rows={dev.rows}, cols={dev.cols}")

    # 依次尝试打开设备
    opened = False
    for device in devices:
        print(f"尝试打开设备 SN: {device.sn.decode(errors='ignore').strip()}")
        if sensor.open_device(device):
            print("打开成功")
            opened = True
            break
        else:
            print("打开失败")

    if not opened:
        print("所有设备打开失败")
        return

    # 配置灯光
    if sensor.config_light():
        print("配置灯光成功")
    else:
        print("配置灯光失败")

    # 准备显示图像
    plt.ion()
    fig, ax = plt.subplots()
    img = None

    try:
        while True:
            frame = sensor.read_frame()
            if frame is None:
                print("读取数据失败")
                break

            if img is None:
                img = ax.imshow(frame, cmap='coolwarm', vmin=0, vmax=1023)
                plt.title("SMiTSense 足底压力分布")
                plt.colorbar(img, ax=ax)
            else:
                img.set_data(frame)
            plt.pause(0.05)
    except KeyboardInterrupt:
        print("用户中断，退出")
    finally:
        sensor.close_device()
        print("设备已关闭")

if __name__ == "__main__":
    main()