path: root/picostream/main.py

import os
import sys
from typing import Any, Dict, Optional

from loguru import logger
from PyQt5.QtCore import QObject, QSettings, QThread, pyqtSignal, pyqtSlot, Qt
from PyQt5.QtGui import QCloseEvent
from PyQt5.QtWidgets import (
    QApplication,
    QCheckBox,
    QComboBox,
    QDoubleSpinBox,
    QFileDialog,
    QFormLayout,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QMainWindow,
    QPushButton,
    QSpinBox,
    QVBoxLayout,
    QWidget,
)

from picostream.cli import Streamer, VOLTAGE_RANGE_MAP
from picostream.dfplot import HDF5LivePlotter


class StreamerWorker(QObject):
    """Worker to run the data acquisition in a background thread."""

    finished = pyqtSignal()
    error = pyqtSignal(str)
    stopRequested = pyqtSignal()

    def __init__(self, settings: Dict[str, Any]) -> None:
        """Initialise the worker."""
        super().__init__()
        self.streamer: Optional[Streamer] = None
        self.settings = settings

    def run(self) -> None:
        """Run the data acquisition."""
        try:
            self.streamer = Streamer(**self.settings)
            self.streamer.run()
        except Exception as e:
            self.error.emit(str(e))
            self.finished.emit()

    @pyqtSlot()
    def stop(self) -> None:
        """Signal the acquisition to stop."""
        if self.streamer:
            self.streamer.shutdown()
        self.finished.emit()


class PicoStreamMainWindow(QMainWindow):
    """The main window for the PicoStream GUI application."""

    def __init__(self) -> None:
        """Initialise the main window."""
        super().__init__()
        self.setWindowTitle("PicoStream")
        self.setGeometry(100, 100, 1200, 600)

        self.settings = QSettings("picostream", "PicoStream")
        self.thread: Optional[QThread] = None
        self.worker: Optional[StreamerWorker] = None

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)

        # Left panel for settings
        settings_panel = QWidget()
        settings_panel.setFixedWidth(350)
        settings_layout = QVBoxLayout(settings_panel)
        form_layout = QFormLayout()

        self.sample_rate_input = QDoubleSpinBox()
        self.sample_rate_input.setRange(1, 125)
        self.sample_rate_input.setValue(62.5)
        self.sample_rate_input.setSuffix(" MS/s")
        form_layout.addRow("Sample Rate:", self.sample_rate_input)

        self.resolution_input = QComboBox()
        self.resolution_input.addItems(["12", "14", "15", "16", "8"])
        form_layout.addRow("Resolution (bits):", self.resolution_input)

        self.voltage_range_input = QComboBox()
        self.voltage_range_input.addItems(VOLTAGE_RANGE_MAP.values())
        form_layout.addRow("Voltage Range:", self.voltage_range_input)

        self.output_file_input = QLineEdit()
        self.output_file_input.setText("output.hdf5")
        file_browse_button = QPushButton("Browse...")
        file_browse_button.clicked.connect(self.select_output_file)
        file_layout = QHBoxLayout()
        file_layout.addWidget(self.output_file_input)
        file_layout.addWidget(file_browse_button)
        form_layout.addRow("Output File:", file_layout)

        self.hw_downsample_input = QSpinBox()
        self.hw_downsample_input.setRange(1, 1000)
        form_layout.addRow("HW Downsample:", self.hw_downsample_input)

        self.downsample_mode_input = QComboBox()
        self.downsample_mode_input.addItems(["average", "aggregate"])
        form_layout.addRow("Downsample Mode:", self.downsample_mode_input)

        self.offset_v_input = QDoubleSpinBox()
        self.offset_v_input.setRange(-1.0, 1.0)
        self.offset_v_input.setSingleStep(0.01)
        self.offset_v_input.setDecimals(3)
        self.offset_v_input.setSuffix(" V")
        form_layout.addRow("Offset:", self.offset_v_input)

        self.live_only_checkbox = QCheckBox("Live-only (overwrite buffer)")
        self.max_buffer_input = QDoubleSpinBox()
        self.max_buffer_input.setRange(0.1, 60.0)
        self.max_buffer_input.setValue(1.0)
        self.max_buffer_input.setSuffix(" s")
        self.live_only_checkbox.stateChanged.connect(
            lambda state: self.max_buffer_input.setEnabled(state > 0)
        )
        form_layout.addRow(self.live_only_checkbox, self.max_buffer_input)

        # Plot settings
        self.plot_window_input = QDoubleSpinBox()
        self.plot_window_input.setRange(0.01, 10.0)
        self.plot_window_input.setValue(0.5)
        self.plot_window_input.setSingleStep(0.1)
        self.plot_window_input.setDecimals(2)
        self.plot_window_input.setSuffix(" s")
        form_layout.addRow("Plot Window:", self.plot_window_input)

        self.y_axis_auto_checkbox = QCheckBox("Auto Y-axis")
        self.y_axis_auto_checkbox.setChecked(True)
        self.y_min_input = QDoubleSpinBox()
        self.y_min_input.setRange(-100000, 100000)
        self.y_min_input.setValue(-1000)
        self.y_min_input.setSuffix(" mV")
        self.y_min_input.setEnabled(False)
        self.y_max_input = QDoubleSpinBox()
        self.y_max_input.setRange(-100000, 100000)
        self.y_max_input.setValue(1000)
        self.y_max_input.setSuffix(" mV")
        self.y_max_input.setEnabled(False)
        
        self.y_axis_auto_checkbox.stateChanged.connect(
            lambda state: self.y_min_input.setEnabled(state == 0)
        )
        self.y_axis_auto_checkbox.stateChanged.connect(
            lambda state: self.y_max_input.setEnabled(state == 0)
        )
        
        form_layout.addRow(self.y_axis_auto_checkbox)
        y_limits_layout = QHBoxLayout()
        y_limits_layout.addWidget(QLabel("Min:"))
        y_limits_layout.addWidget(self.y_min_input)
        y_limits_layout.addWidget(QLabel("Max:"))
        y_limits_layout.addWidget(self.y_max_input)
        form_layout.addRow("Y-axis Limits:", y_limits_layout)

        settings_layout.addLayout(form_layout)

        self.start_button = QPushButton("Start Acquisition")
        self.stop_button = QPushButton("Stop Acquisition")
        self.stop_button.setEnabled(False)
        self.apply_plot_settings_button = QPushButton("Apply Plot Settings")

        settings_layout.addWidget(self.start_button)
        settings_layout.addWidget(self.stop_button)
        settings_layout.addWidget(self.apply_plot_settings_button)
        settings_layout.addStretch()

        # Right panel for the plot
        self.plotter = HDF5LivePlotter(hdf5_path=self.output_file_input.text())

        main_layout.addWidget(settings_panel)
        main_layout.addWidget(self.plotter, 1)

        # Connect signals
        self.start_button.clicked.connect(self.start_acquisition)
        self.stop_button.clicked.connect(self.stop_acquisition)
        self.apply_plot_settings_button.clicked.connect(self.apply_plot_settings)

        self.load_settings()

    def select_output_file(self) -> None:
        """Open a dialog to select the output HDF5 file."""
        file_name, _ = QFileDialog.getSaveFileName(
            self, "Select Output File", self.output_file_input.text(), "HDF5 Files (*.hdf5)"
        )
        if file_name:
            self.output_file_input.setText(file_name)

    def apply_plot_settings(self) -> None:
        """Apply plot settings to the live plotter."""
        window_seconds = self.plot_window_input.value()
        self.plotter.set_display_window(window_seconds)
        
        if self.y_axis_auto_checkbox.isChecked():
            self.plotter.set_y_limits(None, None)
        else:
            y_min = self.y_min_input.value()
            y_max = self.y_max_input.value()
            self.plotter.set_y_limits(y_min, y_max)

    def start_acquisition(self) -> None:
        """Start the background data acquisition."""
        self.start_button.setEnabled(False)
        self.stop_button.setEnabled(True)

        output_file = self.output_file_input.text()

        # Stop plotter updates and clear all state BEFORE changing file path
        self.plotter.stop_updates()
        
        # Clear the plotter's display and reset all internal state
        self.plotter.reset_for_new_file()
        
        # Remove old file if it exists
        if os.path.exists(output_file):
            try:
                os.remove(output_file)
                logger.info(f"Removed existing file: {output_file}")
            except OSError as e:
                self.on_acquisition_error(f"Failed to remove old file: {e}")
                self.on_acquisition_finished()
                return

        # Now set the new file path (plotter is stopped and cleared)
        self.plotter.set_hdf5_path(output_file)
        
        # Apply plot settings before starting
        self.apply_plot_settings()

        settings = {
            "sample_rate_msps": self.sample_rate_input.value(),
            "resolution_bits": int(self.resolution_input.currentText()),
            "channel_range_str": self.voltage_range_input.currentText(),
            "output_file": output_file,
            "hardware_downsample": self.hw_downsample_input.value(),
            "downsample_mode": self.downsample_mode_input.currentText(),
            "offset_v": self.offset_v_input.value(),
            "max_buffer_seconds": self.max_buffer_input.value()
            if self.live_only_checkbox.isChecked()
            else None,
            "enable_live_plot": False,
            "is_gui_mode": True,
        }

        self.thread = QThread()
        self.worker = StreamerWorker(settings)
        self.worker.moveToThread(self.thread)
        self.worker.stopRequested.connect(self.worker.stop, Qt.QueuedConnection)

        self.thread.started.connect(self.worker.run)
        self.worker.finished.connect(self.thread.quit)
        self.worker.finished.connect(self.worker.deleteLater)
        self.thread.finished.connect(self.thread.deleteLater)
        self.worker.finished.connect(self.on_acquisition_finished)
        self.worker.error.connect(self.on_acquisition_error)

        self.thread.start()
        
        # Start plotter updates - it will wait for the file to exist with valid data
        self.plotter.start_updates()

    def stop_acquisition(self) -> None:
        """Stop the background data acquisition."""
        self.plotter.stop_updates()
        if self.worker:
            self.worker.stopRequested.emit()
        self.stop_button.setEnabled(False)

    def on_acquisition_finished(self) -> None:
        """Handle acquisition completion (both success and failure)."""
        logger.info("Acquisition finished.")
        self.plotter.stop_updates()
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.thread = None
        self.worker = None

    def on_acquisition_error(self, err_msg: str) -> None:
        """Handle acquisition error."""
        logger.error(f"Acquisition error: {err_msg}")

    def closeEvent(self, event: QCloseEvent) -> None:
        """Handle window close event."""
        self.save_settings()
        if self.thread and self.thread.isRunning():
            self.stop_acquisition()
            self.thread.wait()
        event.accept()

    def save_settings(self) -> None:
        """Save current settings."""
        self.settings.setValue("sample_rate", self.sample_rate_input.value())
        self.settings.setValue("resolution", self.resolution_input.currentText())
        self.settings.setValue("voltage_range", self.voltage_range_input.currentText())
        self.settings.setValue("output_file", self.output_file_input.text())
        self.settings.setValue("hw_downsample", self.hw_downsample_input.value())
        self.settings.setValue("downsample_mode", self.downsample_mode_input.currentText())
        self.settings.setValue("offset_v", self.offset_v_input.value())
        self.settings.setValue("live_only_mode", self.live_only_checkbox.isChecked())
        self.settings.setValue("max_buffer_seconds", self.max_buffer_input.value())
        self.settings.setValue("plot_window", self.plot_window_input.value())
        self.settings.setValue("y_axis_auto", self.y_axis_auto_checkbox.isChecked())
        
        # Only save Y-axis limits if they're actually being used (not in auto mode)
        if not self.y_axis_auto_checkbox.isChecked():
            self.settings.setValue("y_min", self.y_min_input.value())
            self.settings.setValue("y_max", self.y_max_input.value())

    def load_settings(self) -> None:
        """Load settings."""
        self.sample_rate_input.setValue(self.settings.value("sample_rate", 62.5, type=float))
        self.resolution_input.setCurrentText(self.settings.value("resolution", "12"))
        self.voltage_range_input.setCurrentText(
            self.settings.value("voltage_range", "PS5000A_20V")
        )
        self.output_file_input.setText(self.settings.value("output_file", "output.hdf5"))
        self.hw_downsample_input.setValue(self.settings.value("hw_downsample", 1, type=int))
        self.downsample_mode_input.setCurrentText(
            self.settings.value("downsample_mode", "average")
        )
        self.offset_v_input.setValue(self.settings.value("offset_v", 0.0, type=float))
        live_only = self.settings.value("live_only_mode", False, type=bool)
        self.live_only_checkbox.setChecked(live_only)
        self.max_buffer_input.setValue(
            self.settings.value("max_buffer_seconds", 1.0, type=float)
        )
        self.max_buffer_input.setEnabled(live_only)
        
        self.plot_window_input.setValue(self.settings.value("plot_window", 0.5, type=float))
        y_axis_auto = self.settings.value("y_axis_auto", True, type=bool)
        self.y_axis_auto_checkbox.setChecked(y_axis_auto)
        
        # Only load Y-axis limits if they exist in settings (were saved in manual mode)
        # Otherwise use sensible defaults
        if self.settings.contains("y_min") and self.settings.contains("y_max"):
            self.y_min_input.setValue(self.settings.value("y_min", type=float))
            self.y_max_input.setValue(self.settings.value("y_max", type=float))
        else:
            # Use sensible defaults if no saved values
            self.y_min_input.setValue(-1000.0)
            self.y_max_input.setValue(1000.0)
            
        self.y_min_input.setEnabled(not y_axis_auto)
        self.y_max_input.setEnabled(not y_axis_auto)


def main() -> None:
    """GUI entry point."""
    logger.remove()
    logger.add(sys.stderr, level="INFO")
    app = QApplication(sys.argv)
    main_win = PicoStreamMainWindow()
    main_win.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()