tape_recorder/ui/spectrum_widget.py

"""
Spectrum Widget — Real-time FFT spectrum display using pyqtgraph.
"""

import numpy as np
import pyqtgraph as pg
from PyQt6.QtWidgets import QWidget, QVBoxLayout
from PyQt6.QtCore import QTimer

class SpectrumWidget(QWidget):
    """
    Displays a real-time frequency spectrum.
    """
    def __init__(self, sample_rate: int = 44100, update_interval_ms: int = 50):
        super().__init__()
        self.sample_rate = sample_rate
        self.update_interval_ms = update_interval_ms
        
        # Audio buffers for FFT
        self.fft_size = 4096
        self.buffer_raw = np.zeros(self.fft_size, dtype=np.float32)
        self.buffer_processed = np.zeros(self.fft_size, dtype=np.float32)
        self.buffer_idx_raw = 0
        self.buffer_idx_processed = 0
        
        # Window function for FFT (Hanning)
        self.window = np.hanning(self.fft_size)
        
        self._init_ui()
        
        # Setup timer for UI updates (to avoid overwhelming the GUI thread)
        self.update_timer = QTimer()
        self.update_timer.timeout.connect(self._update_plot)
        self.update_timer.start(self.update_interval_ms)
        
        self._new_data = False

    def _init_ui(self):
        layout = QVBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        
        # Configure pyqtgraph
        pg.setConfigOption('background', '#ffffff')
        pg.setConfigOption('foreground', '#333333')
        
        self.plot_widget = pg.PlotWidget()
        self.plot_widget.setLogMode(x=True, y=False)
        self.plot_widget.setLabel('bottom', 'Frequency', units='Hz')
        self.plot_widget.setLabel('left', 'Magnitude', units='dB')
        self.plot_widget.setYRange(-100, 0)
        self.plot_widget.setXRange(np.log10(20), np.log10(20000))
        self.plot_widget.showGrid(x=True, y=True, alpha=0.3)
        
        from PyQt6.QtCore import Qt
        
        # Plot curves
        self.curve_raw = self.plot_widget.plot(pen=pg.mkPen('#a0a0a0', width=1, style=Qt.PenStyle.DashLine), name="Actual")
        self.curve_processed = self.plot_widget.plot(pen=pg.mkPen('#ff7b00', width=2), name="Corrected")
        self.plot_widget.addLegend()
        
        # Prepare frequency bins
        self.freqs = np.fft.rfftfreq(self.fft_size, d=1.0/self.sample_rate)
        # Avoid log(0) for the very first bin
        self.freqs[0] = 1e-10
        
        layout.addWidget(self.plot_widget)

    def set_sample_rate(self, sr: int):
        if self.sample_rate != sr:
            self.sample_rate = sr
            self.freqs = np.fft.rfftfreq(self.fft_size, d=1.0/self.sample_rate)
            self.freqs[0] = 1e-10

    def add_audio_data(self, raw_data: np.ndarray, processed_data: np.ndarray):
        """
        Add new audio data to the buffers.
        Called from the GUI thread on poll.
        """
        def _add(data, buffer, idx):
            if data.size == 0:
                return idx
            if data.ndim > 1:
                chunk = np.mean(data, axis=1)
            else:
                chunk = data
            n = len(chunk)
            if n >= self.fft_size:
                buffer[:] = chunk[-self.fft_size:]
                return self.fft_size
            else:
                buffer[:-n] = buffer[n:]
                buffer[-n:] = chunk
                return min(idx + n, self.fft_size)
                
        self.buffer_idx_raw = _add(raw_data, self.buffer_raw, self.buffer_idx_raw)
        self.buffer_idx_processed = _add(processed_data, self.buffer_processed, self.buffer_idx_processed)
        self._new_data = True

    def _update_plot(self):
        """Calculate FFT and update the plots on the GUI thread."""
        if not self._new_data or self.buffer_idx_processed < self.fft_size // 4:
            # Decay the plots if no new data
            if not self._new_data:
                for curve in (self.curve_raw, self.curve_processed):
                    current_data = curve.getData()[1]
                    if current_data is not None:
                        decayed = current_data - 2.0 # -2dB per frame decay
                        curve.setData(self.freqs, decayed)
            return
            
        self._new_data = False
        
        def _calc_fft(buffer, curve):
            windowed = buffer * self.window
            mag = np.abs(np.fft.rfft(windowed))
            mag_db = 20 * np.log10(mag / (self.fft_size / 2) + 1e-10)
            
            current_data = curve.getData()[1]
            if current_data is not None:
                alpha = 0.6
                mag_db = alpha * mag_db + (1 - alpha) * current_data
                
            curve.setData(self.freqs, mag_db)

        _calc_fft(self.buffer_raw, self.curve_raw)
        _calc_fft(self.buffer_processed, self.curve_processed)

    def closeEvent(self, event):
        self.update_timer.stop()
        super().closeEvent(event)