# read_and_plot.py
import usb.core, usb.util, threading, time, numpy as np
from collections import deque
import argparse

# 可选可用 pyqtgraph 做实时绘图
try:
    from pyqtgraph.Qt import QtGui, QtCore
    import pyqtgraph as pg
    PG_AVAILABLE = True
except Exception:
    PG_AVAILABLE = False

# MQTT 相关（如不需要可不填）
try:
    import paho.mqtt.client as mqtt
    MQTT_AVAILABLE = True
except Exception:
    MQTT_AVAILABLE = False

# -------- 配置区（根据 probe 输出填写） ----------
VID = 0x3063
PID = 0xa132
INTERFACE = 0
BULK_EP_IN_1 = 0x86  # Bulk IN
BULK_EP_IN_2 = 0x88  # 可能第二通道
BULK_EP_OUT_1 = 0x02
BULK_EP_OUT_2 = 0x04

# 解析/换算参数（如实际量程/灵敏度不同请修改）
ADC_DTYPE = np.int16      # 假设 16-bit signed
ADC_MAX = 32767.0         # int16 最大值
V_RANGE = 10.0            # 假设采集卡量程 ±10 V（即满量程对应 ±10V）
SENS_V_PER_G = 0.5        # XY126D500 灵敏度 0.5 V/g (500 mV/g)

# 采样显示设置
BLOCK_SIZE = 4096         # 每次从设备读的字节数（可调）
UI_UPDATE_MS = 50         # GUI 更新频率（ms）
BUFFER_SECONDS = 5        # 绘图缓冲时间（秒）
FS = 50000                # 逻辑采样率（用于元数据，若不确定可按需修改）

# MQTT 配置（可空）
MQTT_BROKER = "localhost"
MQTT_PORT = 1883
MQTT_TOPIC = "daq/art3132/data"

# 如果你从 Windows 抓包得到了“启动采样”的字节包，可以放在这里（示例占位，默认 None）
# 例如 START_COMMAND = b'\x01\x02\x03\x04'  # <- 从抓包得到
START_COMMAND = None
# 发送到 OUT endpoint 的时候，使用 dev.write(BULK_EP_OUT_1, START_COMMAND)

# ---------------------------------------------------

class USBReader(threading.Thread):
    def __init__(self, dev, ep_in, buffer_deque, stop_event):
        super().__init__(daemon=True)
        self.dev = dev
        self.ep_in = ep_in
        self.buffer = buffer_deque
        self.stop_event = stop_event

    def run(self):
        while not self.stop_event.is_set():
            try:
                data = self.dev.read(self.ep_in, BLOCK_SIZE, timeout=1000)
                if len(data) == 0:
                    continue
                ba = bytes(data)
                # 解析为样点数组（根据 ADC_DTYPE）
                arr = np.frombuffer(ba, dtype=ADC_DTYPE)
                # 将样点逐个放到缓冲区（线程安全的 deque）
                for v in arr:
                    self.buffer.append(v)
            except usb.core.USBError as e:
                # 超时或短包可忽略；其它错误打印
                if e.errno is None:
                    # sometimes errno None for timeout
                    pass
                # print("USB read exception:", e)
                time.sleep(0.001)
            except Exception as e:
                print("Reader thread exception:", e)
                time.sleep(0.1)

def adc_to_voltage(adc_arr):
    return (adc_arr.astype(float) / ADC_MAX) * V_RANGE

def voltage_to_g(voltage_arr):
    return voltage_arr / SENS_V_PER_G

def start_device_if_needed(dev):
    # 如果你知道厂商要求 start command，请放到 START_COMMAND 变量中并解注释下面的写入。
    if START_COMMAND:
        try:
            print("发送 START_COMMAND 到 OUT endpoint...")
            # 尝试通过 bulk out 写入
            try:
                dev.write(BULK_EP_OUT_1, START_COMMAND, timeout=500)
                print("写入 BULK_EP_OUT_1 成功")
            except Exception as e:
                print("写入 OUT 失败（尝试 control transfer）:", e)
                try:
                    # 示例 control transfer（bmRequestType, bRequest 需按抓包替换）
                    dev.ctrl_transfer(0x40, 0x01, 0, 0, START_COMMAND)
                    print("control_transfer 成功")
                except Exception as e2:
                    print("control_transfer 也失败：", e2)
    else:
        print("未设置 START_COMMAND（若设备需要启动命令请抓包并填入 START_COMMAND）")

def main():
    dev = usb.core.find(idVendor=VID, idProduct=PID)
    if dev is None:
        print("设备未找到: {:04x}:{:04x}".format(VID, PID))
        return

    print("找到设备，准备初始化...")
    # detach kernel driver if needed
    try:
        if dev.is_kernel_driver_active(INTERFACE):
            dev.detach_kernel_driver(INTERFACE)
            print("已从内核解绑 interface", INTERFACE)
    except Exception as e:
        print("检查内核驱动时发生：", e)

    dev.set_configuration()
    usb.util.claim_interface(dev, INTERFACE)
    print("已 claim interface", INTERFACE)

    # 尝试发送 start command（若有）
    start_device_if_needed(dev)

    # 创建缓冲区并启动读取线程（这里只用一个 IN endpoint；如需并行读两个 endpoint 可再启动一个线程）
    maxlen = int(BUFFER_SECONDS * FS)
    buf = deque(maxlen=maxlen)
    stop_event = threading.Event()
    reader = USBReader(dev, BULK_EP_IN_1, buf, stop_event)
    reader.start()

    # MQTT init（可选）
    if MQTT_AVAILABLE:
        client = mqtt.Client()
        try:
            client.connect(MQTT_BROKER, MQTT_PORT, 60)
            print("已连接 MQTT:", MQTT_BROKER)
        except Exception as e:
            print("连接 MQTT 失败：", e)
            client = None
    else:
        client = None

    # GUI 绘图（若未安装 pyqtgraph，则只做控制台输出）
    if PG_AVAILABLE:
        app = QtGui.QApplication([])
        win = pg.GraphicsLayoutWidget(title="USB3132A 实时加速度")
        p = win.addPlot()
        p.setLabel('left', '加速度', units='g')
        curve = p.plot()
        win.show()

        def update():
            if len(buf) == 0:
                return
            arr = np.array(buf, dtype=ADC_DTYPE)
            volt = adc_to_voltage(arr)
            gvals = voltage_to_g(volt)
            curve.setData(gvals)

        timer = QtCore.QTimer()
        timer.timeout.connect(update)
        timer.start(UI_UPDATE_MS)

        try:
            QtGui.QApplication.instance().exec_()
        finally:
            stop_event.set()
            reader.join(timeout=1)
            if client:
                client.disconnect()
            usb.util.release_interface(dev, INTERFACE)
            usb.util.dispose_resources(dev)
    else:
        print("pyqtgraph 未安装，进入控制台模式。按 Ctrl-C 停止。")
        try:
            while True:
                time.sleep(1)
                if len(buf) >= 512:
                    arr = np.array([buf.popleft() for _ in range(512)], dtype=ADC_DTYPE)
                    volt = adc_to_voltage(arr)
                    gvals = voltage_to_g(volt)
                    print("样点示例 g:", gvals[:8])
                    # 如果需要上传MQTT，可在这里组包并client.publish
                    if client:
                        try:
                            payload = {"ts": int(time.time()*1000), "samples": gvals.tolist()}
                            client.publish(MQTT_TOPIC, str(payload))
                        except Exception as e:
                            print("MQTT publish error:", e)
        except KeyboardInterrupt:
            print("停止中...")
            stop_event.set()
            reader.join(timeout=1)
            if client:
                client.disconnect()
            usb.util.release_interface(dev, INTERFACE)
            usb.util.dispose_resources(dev)

if __name__ == "__main__":
    main()