path: root/devices/nidaqmx_device.py

"""
devices/nidaqmx_device.py

Combined NI-DAQmx physical device — analog inputs + digital I/O on one NI board.
"""

from typing import Any, Dict

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QFormLayout, QGroupBox,
    QDoubleSpinBox, QCheckBox, QLineEdit, QLabel, QPushButton,
    QListWidget, QListWidgetItem,
)
from PyQt6.QtCore import Qt, QThread, pyqtSignal

from devices.base_device import BaseDevice, ChannelConfig, DeviceInfo, DeviceStatus
from api_layers.nidaqmx_layer import NidaqmxLayer

_ANALOG_COLORS = [
    "#00d4ff", "#ff6b35", "#7fff6e", "#ffcc00",
    "#c77dff", "#ff4d6d", "#4cc9f0", "#f72585",
    "#38b000", "#e9c46a", "#a8dadc", "#e63946",
    "#90e0ef", "#fb8500", "#b5e48c", "#d62828",
]
_DI_COLORS = ["#4cc9f0", "#90e0ef", "#caf0f8", "#0077b6", "#023e8a", "#48cae4", "#ade8f4", "#00b4d8"]
_DO_COLORS = ["#ff6b35", "#ffcc00", "#f77f00", "#fcbf49", "#d62828", "#e63946", "#fb8500", "#ffd166"]


class NidaqmxDevice(BaseDevice):
    DEVICE_TYPE = "nidaqmx"
    ICON        = "🔬"

    def __init__(
        self,
        device_id:   str   = "ni_0",
        num_analog:  int   = 4,
        min_v:       float = -10.0,
        max_v:       float =  10.0,
        num_di:      int   = 2,
        num_do:      int   = 4,
        simulate:    bool  = True,
        ni_device:   str   = "Dev1",
    ):
        self.simulate    = simulate
        self.backend     = "nidaqmx"
        self._ni_device  = ni_device
        self._num_analog = num_analog
        self._min_v      = min_v
        self._max_v      = max_v
        self._num_di     = num_di
        self._num_do     = num_do
        self._last_error = ""

        channels = []
        for i in range(num_analog):
            channels.append(ChannelConfig(
                channel_id=f"ai{i}", name=f"AI{i}", unit="V",
                min_value=min_v, max_value=max_v,
                color=_ANALOG_COLORS[i % len(_ANALOG_COLORS)],
            ))
        for i in range(num_di):
            channels.append(ChannelConfig(
                channel_id=f"di{i}", name=f"DI {i}", unit="",
                min_value=0.0, max_value=1.0,
                color=_DI_COLORS[i % len(_DI_COLORS)],
            ))
        for i in range(num_do):
            channels.append(ChannelConfig(
                channel_id=f"do{i}", name=f"DO {i}", unit="",
                min_value=0.0, max_value=1.0,
                color=_DO_COLORS[i % len(_DO_COLORS)],
            ))

        info = DeviceInfo(
            device_id=device_id,
            name=f"NI-DAQmx ({ni_device})",
            device_type=self.DEVICE_TYPE,
            description=f"NI-DAQmx — {num_analog} AI, {num_di} DI, {num_do} DO",
            manufacturer="National Instruments",
            icon=self.ICON,
            channels=channels,
        )
        super().__init__(info)

        self._output_state: Dict[str, int] = {f"do{i}": 0 for i in range(num_do)}
        self._ai_layer    = None
        self._ni_in_task  = None
        self._ni_out_task = None

    # ── BaseDevice ────────────────────────────────────────────────────────────

    def connect(self) -> bool:
        self._last_error = ""
        if self.simulate:
            self.status = DeviceStatus.SIMULATED
            return True

        ok = True
        try:
            if self._num_analog > 0:
                ai_ids = [ch.channel_id for ch in self.info.channels
                          if ch.channel_id.startswith("ai")]
                self._ai_layer = NidaqmxLayer(
                    device_name=self._ni_device,
                    channels=ai_ids,
                    min_val=self._min_v,
                    max_val=self._max_v,
                    simulate=False,
                )
                ok = ok and self._ai_layer.start()

            if self._num_di > 0 or self._num_do > 0:
                ok = ok and self._ni_digital_connect()
        except Exception as e:
            self._last_error = str(e)
            ok = False

        self.status = DeviceStatus.CONNECTED if ok else DeviceStatus.ERROR
        return ok

    def _ni_digital_connect(self) -> bool:
        try:
            import nidaqmx                                      # type: ignore
            from nidaqmx.constants import LineGrouping          # type: ignore

            if self._num_di > 0:
                self._ni_in_task = nidaqmx.Task()
                for i in range(self._num_di):
                    self._ni_in_task.di_channels.add_di_chan(
                        f"{self._ni_device}/port0/line{i}",
                        line_grouping=LineGrouping.CHAN_PER_LINE,
                    )
                self._ni_in_task.start()

            if self._num_do > 0:
                self._ni_out_task = nidaqmx.Task()
                for i in range(self._num_do):
                    self._ni_out_task.do_channels.add_do_chan(
                        f"{self._ni_device}/port1/line{i}",
                        line_grouping=LineGrouping.CHAN_PER_LINE,
                    )
                self._ni_out_task.start()

            return True
        except Exception as e:
            self._last_error = str(e)
            return False

    def disconnect(self) -> None:
        if self._ai_layer:
            try:
                self._ai_layer.stop()
            except Exception:
                pass
            self._ai_layer = None
        for task in (self._ni_in_task, self._ni_out_task):
            if task:
                try:
                    task.stop()
                    task.close()
                except Exception:
                    pass
        self._ni_in_task  = None
        self._ni_out_task = None
        self.status = DeviceStatus.DISCONNECTED

    def read_channels(self) -> Dict[str, float]:
        if self.simulate:
            return self._sim_read()

        result = {}
        if self._ai_layer:
            try:
                result.update(self._ai_layer.read())
            except Exception:
                pass

        try:
            if self._ni_in_task:
                vals = self._ni_in_task.read()
                di_chs = [c for c in self.info.channels if c.channel_id.startswith("di")]
                for i, ch in enumerate(di_chs):
                    result[ch.channel_id] = float(vals[i] if isinstance(vals, list) else vals)
        except Exception:
            pass

        for ch in (c for c in self.info.channels if c.channel_id.startswith("do")):
            result[ch.channel_id] = float(self._output_state.get(ch.channel_id, 0))

        return result

    def _sim_read(self) -> Dict[str, float]:
        import math, time
        t = time.time()
        result = {}
        for i, ch in enumerate([c for c in self.info.channels if c.channel_id.startswith("ai")]):
            result[ch.channel_id] = math.sin(t + i) * (self._max_v * 0.5)
        for ch in (c for c in self.info.channels if c.channel_id.startswith("di")):
            result[ch.channel_id] = 0.0
        for ch in (c for c in self.info.channels if c.channel_id.startswith("do")):
            result[ch.channel_id] = float(self._output_state.get(ch.channel_id, 0))
        return result

    def write_channel(self, channel_id: str, value: Any) -> bool:
        cid = channel_id.strip()
        if cid.startswith("do"):
            self._output_state[cid] = int(bool(value))
            if not self.simulate and self._ni_out_task:
                try:
                    do_chs = [c for c in self.info.channels if c.channel_id.startswith("do")]
                    states = [self._output_state.get(c.channel_id, 0) for c in do_chs]
                    self._ni_out_task.write(states)
                except Exception as e:
                    print(f"[NidaqmxDevice] write failed: {e}")
            return True
        return False

    def get_save_config(self) -> dict:
        return {
            "device_type": self.DEVICE_TYPE,
            "device_id":   self.info.device_id,
            "name":        self.info.name,
            "num_analog":  self._num_analog,
            "min_v":       self._min_v,
            "max_v":       self._max_v,
            "num_di":      self._num_di,
            "num_do":      self._num_do,
            "simulate":    self.simulate,
            "ni_device":   self._ni_device,
        }

    def get_config_widget(self) -> QWidget:
        return NidaqmxConfigWidget(self)

    def switch_backend(self, simulate: bool, ni_device: str,
                       min_v: float, max_v: float) -> None:
        was_running = self.status in (DeviceStatus.CONNECTED, DeviceStatus.SIMULATED)
        if was_running:
            self.disconnect()
        self.simulate   = simulate
        self._ni_device = ni_device
        self._min_v     = min_v
        self._max_v     = max_v
        if was_running:
            self.connect()


# ── NI scanner thread ─────────────────────────────────────────────────────────

class _NIScanThread(QThread):
    done = pyqtSignal(list)
    def run(self):
        try:
            import nidaqmx                          # type: ignore
            devs = [(d.name, d.product_type) for d in nidaqmx.system.System().devices]
        except Exception:
            devs = []
        self.done.emit(devs)


# ── Config widget ─────────────────────────────────────────────────────────────

class NidaqmxConfigWidget(QWidget):
    def __init__(self, device: NidaqmxDevice):
        super().__init__()
        self.device   = device
        self._scanner = None
        self._build()

    def _build(self):
        root = QVBoxLayout(self)
        root.setContentsMargins(10, 14, 10, 10)
        root.setSpacing(12)
        self.setMinimumWidth(420)

        # ── NI settings ───────────────────────────────────────────────────
        ni_grp  = QGroupBox("NI-DAQmx Settings")
        ni_form = QFormLayout(ni_grp)
        ni_form.setContentsMargins(10, 16, 10, 10)

        self.ni_dev_edit = QLineEdit(self.device._ni_device)
        ni_form.addRow("NI Device:", self.ni_dev_edit)

        self.min_v_spin = QDoubleSpinBox()
        self.min_v_spin.setRange(-100.0, 0.0)
        self.min_v_spin.setValue(self.device._min_v)
        self.min_v_spin.setSuffix(" V")
        ni_form.addRow("Min Voltage:", self.min_v_spin)

        self.max_v_spin = QDoubleSpinBox()
        self.max_v_spin.setRange(0.0, 100.0)
        self.max_v_spin.setValue(self.device._max_v)
        self.max_v_spin.setSuffix(" V")
        ni_form.addRow("Max Voltage:", self.max_v_spin)

        self.sim_chk = QCheckBox("Simulation Mode  (no hardware)")
        self.sim_chk.setChecked(self.device.simulate)
        from core.app_settings import is_developer_mode
        self.sim_chk.setVisible(is_developer_mode())
        ni_form.addRow(self.sim_chk)

        scan_row = QHBoxLayout()
        self._scan_btn = QPushButton("🔍  Scan NI Devices")
        self._scan_btn.setObjectName("addTraceBtn")
        self._scan_btn.clicked.connect(self._scan_ni)
        self._scan_lbl = QLabel("")
        self._scan_lbl.setObjectName("traceSource")
        scan_row.addWidget(self._scan_btn)
        scan_row.addWidget(self._scan_lbl, 1)
        ni_form.addRow(scan_row)

        self._ni_list = QListWidget()
        self._ni_list.setObjectName("portList")
        self._ni_list.setMaximumHeight(80)
        self._ni_list.itemClicked.connect(self._on_ni_selected)
        ni_form.addRow(self._ni_list)

        apply_btn = QPushButton("Apply & Reconnect")
        apply_btn.setObjectName("applyButton")
        apply_btn.clicked.connect(self._apply)
        ni_form.addRow(apply_btn)

        root.addWidget(ni_grp)

        # ── Diagnostics ───────────────────────────────────────────────────
        diag_grp = QGroupBox("Status")
        diag_lay = QVBoxLayout(diag_grp)
        diag_lay.setContentsMargins(10, 16, 10, 10)
        self._diag_lbl = QLabel()
        self._diag_lbl.setObjectName("traceSource")
        self._diag_lbl.setWordWrap(True)
        diag_lay.addWidget(self._diag_lbl)
        root.addWidget(diag_grp)
        root.addStretch()
        self._refresh_diag()

    def _refresh_diag(self):
        lines = [
            f"NI Device: {self.device._ni_device}",
            f"Simulate:  {self.device.simulate}",
            f"Status:    {self.device.status.value}",
        ]
        if self.device._last_error:
            lines.append(f"Error: {self.device._last_error}")
        self._diag_lbl.setText("\n".join(lines))

    def _scan_ni(self):
        self._scan_btn.setEnabled(False)
        self._scan_lbl.setText("Scanning…")
        self._ni_list.clear()
        self._scanner = _NIScanThread()
        self._scanner.done.connect(self._on_ni_found)
        self._scanner.start()

    def _on_ni_found(self, devices):
        self._scan_btn.setEnabled(True)
        self._ni_list.clear()
        if not devices:
            self._scan_lbl.setText("No NI devices found")
            item = QListWidgetItem("  No NI devices detected")
            item.setFlags(item.flags() & ~Qt.ItemFlag.ItemIsSelectable)
            self._ni_list.addItem(item)
        else:
            self._scan_lbl.setText(f"{len(devices)} found — click to select")
            for name, product in devices:
                label = f"  {name}"
                if product:
                    label += f"  —  {product}"
                item = QListWidgetItem(label)
                item.setData(Qt.ItemDataRole.UserRole, name)
                self._ni_list.addItem(item)

    def _on_ni_selected(self, item: QListWidgetItem):
        name = item.data(Qt.ItemDataRole.UserRole)
        if name:
            self.ni_dev_edit.setText(name)

    def _apply(self):
        self.device.switch_backend(
            simulate=self.sim_chk.isChecked(),
            ni_device=self.ni_dev_edit.text().strip() or "Dev1",
            min_v=self.min_v_spin.value(),
            max_v=self.max_v_spin.value(),
        )
        self._refresh_diag()