# read_and_plot_fixed.py
"""
USB3132A read -> convert ADC->Voltage->g -> realtime plot / mqtt upload
Auto-fallback: pyqtgraph -> matplotlib -> console
"""

import usb.core
import usb.util
import threading
import time
import numpy as np
from collections import deque

# try pyqtgraph
try:
    from pyqtgraph.Qt import QtGui, QtCore
    import pyqtgraph as pg
    PG_AVAILABLE = True
except Exception:
    PG_AVAILABLE = False

# try matplotlib
try:
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    MPL_AVAILABLE = True
except Exception:
    MPL_AVAILABLE = False

# try mqtt
try:
    import paho.mqtt.client as mqtt
    MQTT_AVAILABLE = True
except Exception:
    MQTT_AVAILABLE = False

# ------------- Configuration -------------
VID = 0x3063
PID = 0xa132
INTERFACE = 0
BULK_EP_IN_1 = 0x86
BULK_EP_IN_2 = 0x88
BULK_EP_OUT_1 = 0x02
BULK_EP_OUT_2 = 0x04

ADC_DTYPE = np.int16
ADC_MAX = 32767.0
V_RANGE = 10.0          # ±10 V assumed
SENS_V_PER_G = 0.5      # 0.5 V/g (500 mV/g) for XY126D500

BLOCK_SIZE = 4096
UI_UPDATE_MS = 50       # ms
BUFFER_SECONDS = 5
FS = 50000              # logical sampling rate used for buffer sizing

# MQTT (optional)
MQTT_BROKER = "localhost"
MQTT_PORT = 1883
MQTT_TOPIC = "daq/art3132/data"

# If you captured a START command from Windows, place it here as bytes, e.g.:
# START_COMMAND = b'\x01\x02\x03\x04'
START_COMMAND = None

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

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 not data:
                    continue
                ba = bytes(data)
                # parse as signed 16-bit little-endian by default
                arr = np.frombuffer(ba, dtype=ADC_DTYPE)
                # append to buffer
                for v in arr:
                    self.buffer.append(v)
            except usb.core.USBError as e:
                # timeouts etc. can be ignored silently
                # print only unexpected errors
                # On some PyUSB versions e.errno can be None for timeout
                # we just sleep a bit and retry
                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):
    if START_COMMAND:
        try:
            # first try bulk out
            try:
                dev.write(BULK_EP_OUT_1, START_COMMAND, timeout=500)
                print("START_COMMAND written to BULK_EP_OUT_1")
                return
            except Exception as e:
                print("Bulk OUT write failed:", e)
            # fallback to control transfer (example bmRequestType/bRequest may need adjustment)
            try:
                dev.ctrl_transfer(0x40, 0x01, 0, 0, START_COMMAND, timeout=500)
                print("START_COMMAND sent via control transfer")
                return
            except Exception as e:
                print("Control transfer for START_COMMAND failed:", e)
        except Exception as e:
            print("Error sending START_COMMAND:", e)
    else:
        print("No START_COMMAND set. If device requires it, capture it from Windows and set START_COMMAND.")

def open_device():
    dev = usb.core.find(idVendor=VID, idProduct=PID)
    if dev is None:
        raise RuntimeError(f"Device not found: {VID:04x}:{PID:04x}")
    # detach kernel driver if available
    try:
        if hasattr(dev, "is_kernel_driver_active") and dev.is_kernel_driver_active(INTERFACE):
            try:
                dev.detach_kernel_driver(INTERFACE)
                print("Detached kernel driver from interface", INTERFACE)
            except Exception as e:
                print("Could not detach kernel driver:", e)
    except Exception:
        # some platforms/devices may not support is_kernel_driver_active
        pass
    dev.set_configuration()
    usb.util.claim_interface(dev, INTERFACE)
    return dev

def run_pyqtgraph(buf, stop_event, client):
    app = QtGui.QApplication([])
    win = pg.GraphicsLayoutWidget(title="USB3132A Real-time Acceleration (g)")
    p = win.addPlot()
    p.setLabel('left', 'Acceleration', 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()
        if client:
            try:
                client.disconnect()
            except Exception:
                pass

def run_matplotlib(buf, stop_event, client):
    fig, ax = plt.subplots()
    line, = ax.plot([], [], lw=1)
    ax.set_ylim(-5, 5)
    ax.set_xlim(0, int(BUFFER_SECONDS * FS))
    ax.set_xlabel("Sample")
    ax.set_ylabel("Acceleration (g)")
    ax.set_title("USB3132A Realtime Acceleration")

    def init():
        line.set_data([], [])
        return (line,)

    def update(frame):
        if len(buf) == 0:
            return (line,)
        arr = np.array(buf, dtype=ADC_DTYPE)
        volt = adc_to_voltage(arr)
        gvals = voltage_to_g(volt)
        x = np.arange(len(gvals))
        line.set_data(x, gvals)
        ax.set_xlim(0, max(1000, len(gvals)))
        return (line,)

    ani = FuncAnimation(fig, update, init_func=init, interval=UI_UPDATE_MS, blit=False)
    try:
        plt.show()
    finally:
        stop_event.set()
        if client:
            try:
                client.disconnect()
            except Exception:
                pass

def run_console(buf, stop_event, client):
    print("Running in console mode. Press Ctrl-C to stop.")
    try:
        while not stop_event.is_set():
            time.sleep(1)
            if len(buf) >= 512:
                chunk = [buf.popleft() for _ in range(512)]
                arr = np.array(chunk, dtype=ADC_DTYPE)
                volt = adc_to_voltage(arr)
                gvals = voltage_to_g(volt)
                print("Sample g (first 8):", np.round(gvals[:8], 4))
                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:
        pass
    finally:
        stop_event.set()
        if client:
            try:
                client.disconnect()
            except Exception:
                pass

def main():
    # prepare buffer and threads
    maxlen = int(BUFFER_SECONDS * FS)
    buf = deque(maxlen=maxlen)
    stop_event = threading.Event()

    # open device
    try:
        dev = open_device()
    except Exception as e:
        print("Failed to open device:", e)
        return

    print("Device opened and interface claimed.")

    # optionally send START command if you have one
    start_device_if_needed(dev)

    # start reader thread (using primary IN endpoint)
    reader = USBReader(dev, BULK_EP_IN_1, buf, stop_event)
    reader.start()

    # setup MQTT client if available
    if MQTT_AVAILABLE:
        client = mqtt.Client()
        try:
            client.connect(MQTT_BROKER, MQTT_PORT, 60)
            print("Connected to MQTT broker:", MQTT_BROKER)
        except Exception as e:
            print("Could not connect to MQTT broker:", e)
            client = None
    else:
        client = None

    # choose UI mode
    try:
        if PG_AVAILABLE:
            run_pyqtgraph(buf, stop_event, client)
        elif MPL_AVAILABLE:
            run_matplotlib(buf, stop_event, client)
        else:
            run_console(buf, stop_event, client)
    finally:
        # cleanup
        stop_event.set()
        reader.join(timeout=1)
        try:
            usb.util.release_interface(dev, INTERFACE)
        except Exception:
            pass
        try:
            usb.util.dispose_resources(dev)
        except Exception:
            pass
        print("Cleaned up and exiting.")

if __name__ == "__main__":
    main()