import threading
import time
import numpy as np
from enum import Enum
# 导入官方例程的底层依赖
from class_ch341 import *
from class_sensorcmd import *
from class_finger import *

# === 配置区域 ===
DEF_MAX_FINGER_NUM = 3  # 修改为 3 个传感器 (Finger 0, 1, 2)
PCA_ADDR = 0x70  # I2C 多路复用器地址
SAMPLE_RATE_MS = 10  # 采样间隔 (ms)，10ms = 100Hz


class EnumCh341ConnectStatus(Enum):
    CH341_CONNECT_INIT = 0
    CH341_CONNECT_OPEN = 1
    CH341_CONNECT_SET_SPEED = 2
    CH341_CONNECT_SAMPLE_START = 3
    CH341_CONNECT_CHECK = 4
    CH341_CONNECT_SAMPLE_STOP = 5


class TactileSensorDAQ:
    def __init__(self):
        # 1. 硬件初始化
        self.ch341 = ClassCh341()
        self.fingers = list()
        # 初始化3个传感器对象，ID从2开始 (假设硬件拨码是 2,3,4)
        for i in range(DEF_MAX_FINGER_NUM):
            self.fingers.append(ClassFinger(4 + i, self.ch341))

        # 2. 状态机变量
        self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_INIT
        self.ch341CheckTimer = 0
        self.pcaAddr = PCA_ADDR
        self.syncTimer = 0

        # 3. 数据容器 (核心修改)
        # 12维数据: [F0_U1_Fn, F0_U1_Ft, F0_U2_Fn, F0_U2_Ft, F1..., F2...]
        self.raw_data = np.zeros(12, dtype=np.float32)  # 实时读取值
        self.offset = np.zeros(12, dtype=np.float32)  # 去皮偏移量
        self.clean_data = np.zeros(12, dtype=np.float32)  # 输出值 (Raw - Offset)

        # 4. 线程控制
        self.running = False
        self.lock = threading.Lock()  # 线程锁，保证读取安全
        self.thread = None

    def _set_sensor_enable(self, idx):
        """控制 I2C 多路复用器通道"""
        _pack = list()
        _pack.append(idx)
        self.ch341.write(self.pcaAddr, _pack)

    def _update_state_machine(self):
        """维持 CH341 连接状态机 (原 logic 的简化版)"""
        if self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_INIT:
            if self.ch341.init():
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_OPEN

        elif self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_OPEN:
            if self.ch341.open():
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_SET_SPEED
            else:
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_INIT

        elif self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_SET_SPEED:
            if self.ch341.set_speed(self.ch341.IIC_SPEED_400):
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_SAMPLE_START
            else:
                self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_SAMPLE_START  # Retry

        elif self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_SAMPLE_START:
            # 连接建立成功，进入读取循环
            self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_CHECK

        elif self.connectStatus == EnumCh341ConnectStatus.CH341_CONNECT_CHECK:
            # 这里的检查逻辑放入主循环中执行
            pass

    def _read_hardware(self):
        """读取所有传感器数据的核心函数"""
        connectedSensorChan = 0

        # 临时列表存储本轮读取的数据
        temp_data_buffer = []

        for fingerIndex in range(len(self.fingers)):
            # 1. 切通道
            self._set_sensor_enable(1 << (self.fingers[fingerIndex].pcaIdx))
            connectedSensorChan |= 1 << (self.fingers[fingerIndex].pcaIdx)

            current_finger = self.fingers[fingerIndex]

            # 2. 检查连接与读取
            if not current_finger.connect:
                if current_finger.checkSensor():
                    print(f"[System] Finger {fingerIndex} Connected!")
            else:
                current_finger.capRead()

                # 3. 提取数据 (这是修改的关键！)
                # 假设每个传感器有 ydds_num (通常是2) 个单元
                # 这里的 nf 和 tf 应该是数组
                for unit_i in range(current_finger.projectPara.ydds_num):
                    # 提取法向力 Fn
                    fn = current_finger.readData.nf[unit_i]
                    # 提取切向力 Ft
                    ft = current_finger.readData.tf[unit_i]

                    temp_data_buffer.append(fn)
                    temp_data_buffer.append(ft)

        # 4. 更新共享内存
        if len(temp_data_buffer) == 12:  # 确保数据完整
            with self.lock:
                self.raw_data = np.array(temp_data_buffer, dtype=np.float32)
                # 计算去皮后的数据
                self.clean_data = self.raw_data - self.offset
                # 简单滤波：置零负值噪声
                # self.clean_data[self.clean_data < 0] = 0

        # 5. 同步逻辑 (保持原厂逻辑，防止电容漂移)
        if (time.time() - self.syncTimer) > 1.0:  # 1秒同步一次
            self.syncTimer = time.time()
            self._set_sensor_enable(connectedSensorChan)
            for f in self.fingers:
                if f.connect:
                    f.snsCmd.setSensorSync(0)
                    break

    def _thread_worker(self):
        """后台线程主循环"""
        while self.running:
            # 1. 维护连接
            if self.connectStatus != EnumCh341ConnectStatus.CH341_CONNECT_CHECK:
                self._update_state_machine()
                time.sleep(0.1)
                continue

            # 2. 读取数据
            start_time = time.time()
            try:
                self._read_hardware()
            except Exception as e:
                print(f"Read Error: {e}")

            # 3. 检查连接心跳 (保持原厂逻辑)
            self.ch341CheckTimer += (time.time() - start_time) * 1000
            if self.ch341CheckTimer >= 1000:
                self.ch341CheckTimer = 0
                if not self.ch341.connectCheck():
                    print("CH341 Disconnected!")
                    self.connectStatus = EnumCh341ConnectStatus.CH341_CONNECT_INIT

            # 4. 控制采样率
            elapsed = (time.time() - start_time) * 1000
            sleep_time = (SAMPLE_RATE_MS - elapsed) / 1000.0
            if sleep_time > 0:
                time.sleep(sleep_time)

    # === 用户API ===

    def start(self):
        """启动采集"""
        if self.running: return
        self.running = True
        self.thread = threading.Thread(target=self._thread_worker)
        self.thread.daemon = True
        self.thread.start()
        print("Tactile Sensor System Started.")

    def stop(self):
        """停止采集"""
        self.running = False
        if self.thread:
            self.thread.join()
        self.ch341.disconnect()
        print("Tactile Sensor System Stopped.")

    def tare(self):
        """去皮：将当前读数设为零点"""
        print("Taring sensors... please wait.")
        time.sleep(1)  # 等待数据稳定
        with self.lock:
            self.offset = np.copy(self.raw_data)
        print("Tare complete.")

    def get_data(self):
        """获取最新的12维力数据"""
        with self.lock:
            return np.copy(self.clean_data)


# === 调试代码 (直接运行此文件测试) ===
if __name__ == "__main__":
    sensor = TactileSensorDAQ()
    sensor.start()

    try:
        # 等待连接稳定
        print("Waiting for sensors to connect...")
        time.sleep(3)

        sensor.tare()  # 初始去皮

        while True:
            data = sensor.get_data()
            log_str = ""
            for i in range(3):  # 遍历 3 个手指
                base_idx = i * 4
                # 为了显示简洁，我们将 Unit1 和 Unit2 的 Fn 相加，作为一个总压力显示
                # 你也可以根据需要显示全部细节
                f_n_total = data[base_idx] + data[base_idx + 2]
                f_t_total = data[base_idx + 1] + data[base_idx + 3]

                log_str += f"F{i}: N={f_n_total:.1f} T={f_t_total:.1f} | "

            print(log_str)
            # === 修改结束 ===


            time.sleep(0.1)  # 打印频率

    except KeyboardInterrupt:
        sensor.stop()