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 buffer for FFT
        self.fft_size = 4096
        self.buffer = np.zeros(self.fft_size, dtype=np.float32)
        self.buffer_idx = 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', '#161625')
        pg.setConfigOption('foreground', '#a0a0b0')
        
        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)
        
        # Plot curve (Amber color for analog feel)
        self.curve = self.plot_widget.plot(pen=pg.mkPen('#ff7b00', width=2))
        
        # 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, data: np.ndarray):
        """
        Add new audio data to the buffer.
        Called from the audio callback thread.
        """
        if data.size == 0:
            return
            
        # Convert to mono if stereo
        if data.ndim > 1:
            chunk = np.mean(data, axis=1)
        else:
            chunk = data
            
        n = len(chunk)
        if n >= self.fft_size:
            # Take the latest fft_size samples
            self.buffer[:] = chunk[-self.fft_size:]
            self.buffer_idx = self.fft_size
        else:
            # Shift buffer left and append new data
            self.buffer[:-n] = self.buffer[n:]
            self.buffer[-n:] = chunk
            self.buffer_idx = min(self.buffer_idx + n, self.fft_size)
            
        self._new_data = True

    def _update_plot(self):
        """Calculate FFT and update the plot on the GUI thread."""
        if not self._new_data or self.buffer_idx < self.fft_size // 4:
            # Decay the plot if no new data
            if not self._new_data:
                 current_data = self.curve.getData()[1]
                 if current_data is not None:
                     decayed = current_data - 2.0 # -2dB per frame decay
                     self.curve.setData(self.freqs, decayed)
            return
            
        self._new_data = False
        
        # Apply window and calculate FFT
        windowed = self.buffer * self.window
        mag = np.abs(np.fft.rfft(windowed))
        
        # Normalize and convert to dB
        # Reference is amplitude 1.0 (0 dBFS)
        mag_db = 20 * np.log10(mag / (self.fft_size / 2) + 1e-10)
        
        # Smooth slightly in time domain (optional, but looks better)
        current_data = self.curve.getData()[1]
        if current_data is not None:
             alpha = 0.6
             mag_db = alpha * mag_db + (1 - alpha) * current_data
        
        # Update curve
        self.curve.setData(self.freqs, mag_db)

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