path: root/devices/digital_io.py

"""
devices/digital_io.py

Digital I/O device module.  Backends: NI-DAQmx  or  Arduino.

NI backend   – uses nidaqmx digital line tasks (P0.0..P0.7)
Arduino      – uses ArduinoLayer digital pin reads; writes via W:Dxx:val
"""

import random
import time
from typing import Any, Dict

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QGroupBox,
    QCheckBox, QPushButton, QLabel, QFormLayout,
    QComboBox, QLineEdit,
)

from devices.base_device import BaseDevice, ChannelConfig, DeviceInfo, DeviceStatus

_IN_COLORS  = ["#00d4ff", "#4cc9f0", "#90e0ef", "#caf0f8",
               "#0077b6", "#023e8a", "#48cae4", "#ade8f4"]
_OUT_COLORS = ["#ff6b35", "#ffcc00", "#f77f00", "#fcbf49",
               "#d62828", "#e63946", "#fb8500", "#ffd166"]


class DigitalIODevice(BaseDevice):
    DEVICE_TYPE = "digital_io"
    ICON        = "⬛"

    def __init__(
        self,
        device_id:    str  = "dio_0",
        num_inputs:   int  = 8,
        num_outputs:  int  = 8,
        simulate:     bool = True,
        backend:      str  = "nidaqmx",  # "nidaqmx" | "arduino"
        ni_device:    str  = "Dev1",
        ard_port:     str  = "COM3",
        ard_baud:     int  = 115200,
    ):
        self.simulate     = simulate
        self.backend      = backend
        self._ni_device   = ni_device
        self._ard_port    = ard_port
        self._ard_baud    = ard_baud
        self._num_inputs  = num_inputs
        self._num_outputs = num_outputs

        channels = []
        for i in range(num_inputs):
            channels.append(ChannelConfig(
                channel_id=f"di{i}", name=f"DI {i}", unit="",
                min_value=0.0, max_value=1.0,
                color=_IN_COLORS[i % len(_IN_COLORS)],
            ))
        for i in range(num_outputs):
            channels.append(ChannelConfig(
                channel_id=f"do{i}", name=f"DO {i}", unit="",
                min_value=0.0, max_value=1.0,
                color=_OUT_COLORS[i % len(_OUT_COLORS)],
            ))

        info = DeviceInfo(
            device_id=device_id,
            name=f"Digital I/O ({backend.upper()})",
            device_type=self.DEVICE_TYPE,
            description="Digital input/output module",
            icon=self.ICON,
            channels=channels,
        )
        super().__init__(info)

        self._output_state: Dict[str, int] = {
            f"do{i}": 0 for i in range(num_outputs)
        }
        self._sim_toggle: Dict[str, int] = {}
        self._sim_state:  Dict[str, int] = {}

        # Hardware task placeholders
        self._ni_in_task  = None
        self._ni_out_task = None
        self._ard_layer   = None

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

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

        if self.backend == "nidaqmx":
            return self._ni_connect()
        else:
            return self._ard_connect()

    def _ni_connect(self) -> bool:
        try:
            import nidaqmx                                          # type: ignore
            from nidaqmx.constants import LineGrouping              # type: ignore
            n_in  = sum(1 for c in self.info.channels if c.channel_id.startswith("di"))
            n_out = sum(1 for c in self.info.channels if c.channel_id.startswith("do"))

            if n_in:
                self._ni_in_task = nidaqmx.Task()
                for i in range(n_in):
                    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 n_out:
                self._ni_out_task = nidaqmx.Task()
                for i in range(n_out):
                    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()

            self.status = DeviceStatus.CONNECTED
            return True
        except Exception as e:
            print(f"[DigitalIODevice] NI connect failed: {e}")
            self.status = DeviceStatus.ERROR
            return False

    def _ard_connect(self) -> bool:
        from api_layers.port_registry import port_registry
        # Get shared layer — won't open a second connection if AnalogInputDevice
        # is already connected on the same port
        self._ard_layer = port_registry.get_layer(
            port=self._ard_port,
            baud=self._ard_baud,
            simulate=False,
        )
        if self._ard_layer.is_connected:
            # Port already open — just register ourselves and return OK
            self.status = DeviceStatus.CONNECTED
            return True
        ok = self._ard_layer.connect()
        self.status = DeviceStatus.CONNECTED if ok else DeviceStatus.ERROR
        return ok

    def disconnect(self) -> None:
        if self._ni_in_task:
            try: self._ni_in_task.stop(); self._ni_in_task.close()
            except Exception: pass
        if self._ni_out_task:
            try: self._ni_out_task.stop(); self._ni_out_task.close()
            except Exception: pass
        if self._ard_layer and self.backend == "arduino" and not self.simulate:
            from api_layers.port_registry import port_registry
            port_registry.release(self._ard_port, self._ard_baud)
            self._ard_layer = None
        self.status = DeviceStatus.DISCONNECTED

    def read_channels(self) -> Dict[str, float]:
        if self.simulate:
            return self._sim_read()
        if self.backend == "nidaqmx":
            return self._ni_read()
        return self._ard_read()

    def _ni_read(self) -> Dict[str, float]:
        result = {}
        try:
            if self._ni_in_task:
                vals = self._ni_in_task.read()
                for i, ch in enumerate(c for c in self.info.channels if c.channel_id.startswith("di")):
                    result[ch.channel_id] = float(vals[i] if isinstance(vals, list) else vals)
        except Exception as e:
            print(f"[DigitalIODevice] NI read failed: {e}")
        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 _ard_read(self) -> Dict[str, float]:
        if not self._ard_layer:
            return {}
        raw = self._ard_layer.read()
        result = {}
        for ch in self.info.channels:
            if ch.channel_id in raw:
                result[ch.channel_id] = raw[ch.channel_id]
            elif ch.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]:
        result = {}
        for ch in self.info.channels:
            if ch.channel_id.startswith("di"):
                cnt = self._sim_toggle.get(ch.channel_id, 0) + 1
                if cnt >= random.randint(8, 40):
                    self._sim_state[ch.channel_id] = 1 - self._sim_state.get(ch.channel_id, 0)
                    cnt = 0
                self._sim_toggle[ch.channel_id] = cnt
                result[ch.channel_id] = float(self._sim_state.get(ch.channel_id, 0))
            else:
                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:
        self._output_state[channel_id] = int(bool(value))
        if not self.simulate:
            if self.backend == "nidaqmx" and self._ni_out_task:
                try:
                    out_chs = [c for c in self.info.channels if c.channel_id.startswith("do")]
                    idx = next((i for i, c in enumerate(out_chs) if c.channel_id == channel_id), None)
                    if idx is not None:
                        states = [self._output_state.get(c.channel_id, 0) for c in out_chs]
                        self._ni_out_task.write(states)
                except Exception as e:
                    print(f"[DigitalIODevice] NI write failed: {e}")
            elif self.backend == "arduino" and self._ard_layer:
                # Map channel ID to Arduino pin name:
                # do0→D5, do1→D6, do2→D7... (matches DIGITAL_OUT in firmware)
                # Or the channel_id itself if it already looks like D5
                pin = channel_id if channel_id.upper().startswith("D") else f"D{5 + int(channel_id.replace('do',''))}"
                self._ard_layer.digital_write(pin, int(bool(value)))
        return True

    def pwm_channel(self, channel_id: str, duty_pct: float) -> bool:
        """Send a PWM command to an Arduino output pin (0–100%)."""
        if self.backend == "arduino" and self._ard_layer and not self.simulate:
            pin = channel_id if channel_id.upper().startswith("D") else f"D{9 + int(channel_id.replace('do',''))}"
            return self._ard_layer.pwm_write_pct(pin, duty_pct)
        return False

    def set_parameter(self, name: str, value: Any) -> bool:
        """Send a named parameter/setpoint to the Arduino."""
        if self.backend == "arduino" and self._ard_layer and not self.simulate:
            return self._ard_layer.set_parameter(name, value)
        return False

    def switch_backend(self, backend: str, simulate: bool,
                       ni_device: str, ard_port: str, ard_baud: int) -> None:
        was_running = self.status in (DeviceStatus.CONNECTED, DeviceStatus.SIMULATED)
        if was_running:
            self.disconnect()

        self.backend    = backend
        self.simulate   = simulate
        self._ni_device = ni_device
        self._ard_port  = ard_port
        self._ard_baud  = ard_baud
        self._ard_layer = None
        self._ni_in_task  = None
        self._ni_out_task = None

        if was_running:
            self.connect()

    def get_save_config(self) -> dict:
        return {
            "device_type": self.DEVICE_TYPE,
            "device_id":   self.info.device_id,
            "num_inputs":  self._num_inputs,
            "num_outputs": self._num_outputs,
            "simulate":    self.simulate,
            "backend":     self.backend,
            "ni_device":   self._ni_device,
            "ard_port":    self._ard_port,
            "ard_baud":    self._ard_baud,
        }

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


class DigitalIOConfigWidget(QWidget):
    def __init__(self, device: DigitalIODevice):
        super().__init__()
        self.device = device
        self._buttons: Dict[str, QPushButton] = {}
        self._build()

    def _build(self):
        root = QVBoxLayout(self)
        root.setContentsMargins(8, 8, 8, 8)
        self.setMinimumWidth(380)

        be_grp  = QGroupBox("Backend")
        be_form = QFormLayout(be_grp)
        self.be_cb = QComboBox()
        self.be_cb.addItems(["nidaqmx", "arduino"])
        self.be_cb.setCurrentText(self.device.backend)
        be_form.addRow("Backend:", self.be_cb)
        self.sim_chk = QCheckBox("Simulate")
        self.sim_chk.setChecked(self.device.simulate)
        be_form.addRow(self.sim_chk)

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

        self.ard_port_edit = QLineEdit(self.device._ard_port)
        self.ard_port_edit.setPlaceholderText("e.g. /dev/ttyUSB0 or COM3")
        be_form.addRow("Arduino Port:", self.ard_port_edit)

        self.ard_baud_cb = QComboBox()
        self.ard_baud_cb.addItems(["9600", "19200", "57600", "115200", "230400"])
        self.ard_baud_cb.setCurrentText(str(self.device._ard_baud))
        be_form.addRow("Baud Rate:", self.ard_baud_cb)

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

        self.be_cb.currentTextChanged.connect(self._update_visibility)
        root.addWidget(be_grp)
        self._update_visibility(self.device.backend)

        out_grp  = QGroupBox("Digital Outputs")
        out_lay  = QVBoxLayout(out_grp)
        for ch in (c for c in self.device.info.channels if c.channel_id.startswith("do")):
            row = QHBoxLayout()
            lbl = QLabel(ch.name)
            lbl.setMinimumWidth(50)
            btn = QPushButton("OFF")
            btn.setCheckable(True)
            btn.setChecked(bool(self.device._output_state.get(ch.channel_id, 0)))
            btn.setObjectName("digitalOutBtn")
            cid = ch.channel_id

            def _tog(checked, c=cid, b=btn):
                self.device.write_channel(c, checked)
                b.setText("ON" if checked else "OFF")

            btn.toggled.connect(_tog)
            row.addWidget(lbl)
            row.addWidget(btn)
            out_lay.addLayout(row)
            self._buttons[ch.channel_id] = btn

        root.addWidget(out_grp)
        root.addStretch()

    def _update_visibility(self, backend: str = ""):
        backend = backend or self.be_cb.currentText()
        self.ni_dev_edit.setVisible(backend == "nidaqmx")
        self.ard_port_edit.setVisible(backend == "arduino")
        self.ard_baud_cb.setVisible(backend == "arduino")

    def _apply(self):
        self.device.switch_backend(
            backend=self.be_cb.currentText(),
            simulate=self.sim_chk.isChecked(),
            ni_device=self.ni_dev_edit.text().strip() or "Dev1",
            ard_port=self.ard_port_edit.text().strip(),
            ard_baud=int(self.ard_baud_cb.currentText()),
        )