path: root/picostream/main.py

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

from loguru import logger
from PyQt5.QtCore import QObject, QSettings, QThread, QTimer, pyqtSignal, pyqtSlot, Qt
from PyQt5.QtGui import QCloseEvent, QFont
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.bandwidth_limiter_input = QComboBox()
        self.bandwidth_limiter_input.addItems(["Full", "20 MHz"])
        form_layout.addRow("Bandwidth:", self.bandwidth_limiter_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)

        # Add status display section
        status_group = QWidget()
        status_group_layout = QVBoxLayout(status_group)
        status_title = QLabel("Status")
        status_title.setStyleSheet("font-weight: bold; font-size: 12px;")
        status_group_layout.addWidget(status_title)
        
        # Create status labels with initial values
        self.status_heartbeat = QLabel("UI: -")
        self.status_samples = QLabel("Samples: 0")
        self.status_rate = QLabel("Rate: -")
        self.status_latency = QLabel("Latency: -")
        self.status_errors = QLabel("Errors: 0")
        self.status_saturation = QLabel("Saturation: -")
        self.status_acquisition = QLabel("Waiting for file...")
        
        # Set monospace font for status
        font = QFont()
        font.setFamily("Monospace")
        font.setFixedPitch(True)
        font.setPointSize(9)
        for label in [
            self.status_heartbeat,
            self.status_samples,
            self.status_rate,
            self.status_latency,
            self.status_errors,
            self.status_saturation,
            self.status_acquisition,
        ]:
            label.setFont(font)
            label.setWordWrap(True)
        
        status_group_layout.addWidget(self.status_heartbeat)
        status_group_layout.addWidget(self.status_errors)
        status_group_layout.addWidget(self.status_saturation)
        status_group_layout.addWidget(self.status_samples)
        status_group_layout.addWidget(self.status_latency)
        status_group_layout.addWidget(self.status_rate)
        status_group_layout.addWidget(self.status_acquisition)
        
        settings_layout.addWidget(status_group)
        settings_layout.addStretch()

        # Right panel for the plot
        self.plotter = HDF5LivePlotter(hdf5_path=self.output_file_input.text())
        
        # Hide the plotter's status bar since we're showing it in the sidebar
        self.plotter.hide_status_bar(hide=True)

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

        # Create a timer to sync status from plotter to main window
        self.status_sync_timer = QTimer()
        self.status_sync_timer.timeout.connect(self.sync_status_from_plotter)
        # Timer will be started when acquisition begins

        # 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()
        
        # Ensure status bar remains hidden after reset
        self.plotter.hide_status_bar(hide=True)
        
        # 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,
            "bandwidth_limiter": self.bandwidth_limiter_input.currentText().lower(),
        }

        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()
        
        # Start status sync timer
        self.status_sync_timer.start(100)  # Update every 100ms

    def stop_acquisition(self) -> None:
        """Stop the background data acquisition."""
        self.plotter.stop_updates()
        # Stop the status sync timer to save CPU
        if hasattr(self, 'status_sync_timer'):
            self.status_sync_timer.stop()
        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 sync_status_from_plotter(self) -> None:
        """Sync status from the plotter to the main window status labels."""
        try:
            # Check if plotter exists and has been initialized
            if not hasattr(self, 'plotter') or not self.plotter:
                return
                
            # Check if plotter has the required attributes
            required_labels = [
                'heartbeat_label', 'samples_label', 'rate_label',
                'plotter_latency_label', 'error_label', 
                'saturation_label', 'acq_status_label'
            ]
            
            # Helper function to safely strip HTML tags
            def strip_html(text: str) -> str:
                if not text:
                    return text
                # Remove all HTML tags
                return re.sub(r'<[^>]+>', '', text)
            
            # Safely get text from label or return empty string
            def safe_get_text(label_name: str) -> str:
                if hasattr(self.plotter, label_name):
                    label = getattr(self.plotter, label_name)
                    if hasattr(label, 'text'):
                        return strip_html(label.text())
                return ""
            
            # Update all status labels
            self.status_heartbeat.setText(safe_get_text('heartbeat_label'))
            self.status_samples.setText(safe_get_text('samples_label'))
            self.status_rate.setText(safe_get_text('rate_label'))
            self.status_latency.setText(safe_get_text('plotter_latency_label'))
            self.status_errors.setText(safe_get_text('error_label'))
            self.status_saturation.setText(safe_get_text('saturation_label'))
            self.status_acquisition.setText(safe_get_text('acq_status_label'))
            
        except Exception as e:
            # Don't let errors in status sync crash the app
            logger.debug(f"Error syncing status from plotter: {e}")
            pass

    def closeEvent(self, event: QCloseEvent) -> None:
        """Handle window close event."""
        self.save_settings()
        
        # Stop the status sync timer first
        if hasattr(self, 'status_sync_timer'):
            self.status_sync_timer.stop()
            self.status_sync_timer = None
        
        # Stop plotter updates
        if hasattr(self, 'plotter') and self.plotter:
            self.plotter.stop_updates()
        
        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())
        
        self.settings.setValue("bandwidth_limiter", self.bandwidth_limiter_input.currentText())

    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)
        
        self.bandwidth_limiter_input.setCurrentText(
            self.settings.value("bandwidth_limiter", "Full")
        )


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()