tape_recorder/ui/waveform_widget.py

"""
Waveform Widget — Overview display of an entire audio file.
"""

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

class WaveformWidget(QWidget):
    """
    Displays the static waveform of a loaded audio file
    and an interactive playhead cursor.
    """
    
    # Emitted when user clicks on the waveform to seek
    seek_requested = pyqtSignal(float) # position in seconds
    
    def __init__(self):
        super().__init__()
        self._init_ui()
        self.duration_sec = 0.0

    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.setMouseEnabled(x=False, y=False)
        self.plot_widget.hideAxis('left')
        self.plot_widget.hideAxis('bottom')
        self.plot_widget.setYRange(-1.1, 1.1)
        
        # Waveform curves (L/R)
        self.curve_l = self.plot_widget.plot(pen=pg.mkPen('#3a3a5c', width=1))
        self.curve_r = self.plot_widget.plot(pen=pg.mkPen('#4a4a7a', width=1))
        
        # Playhead line
        self.playhead = pg.InfiniteLine(pos=0, angle=90, pen=pg.mkPen('#ff7b00', width=2))
        self.plot_widget.addItem(self.playhead)
        
        # Mouse click handling for seek
        self.plot_widget.scene().sigMouseClicked.connect(self._on_mouse_clicked)
        
        layout.addWidget(self.plot_widget)

    def load_audio(self, audio_data):
        """Generate waveform overview from AudioData."""
        self.duration_sec = audio_data.duration
        self.plot_widget.setXRange(0, self.duration_sec)
        
        # Decimate data for display (we don't need millions of points)
        # Target ~2000 points for the view
        target_points = 2000
        n_samples = audio_data.num_samples
        step = max(1, n_samples // target_points)
        
        # Instead of just picking every Nth sample, get min/max over blocks 
        # to correctly represent the envelope
        blocks = n_samples // step
        
        # Time axis
        t = np.linspace(0, self.duration_sec, blocks)
        
        samples = audio_data.samples[:blocks * step]
        if audio_data.channels >= 2:
            left = samples[:, 0].reshape(-1, step)
            right = samples[:, 1].reshape(-1, step)
            
            l_max = np.max(left, axis=1)
            l_min = np.min(left, axis=1)
            r_max = np.max(right, axis=1)
            r_min = np.min(right, axis=1)
            
            # Interleave min/max for shaded look, or just use max
            # For simplicity, we just plot the max envelope
            self.curve_l.setData(t, l_max)
            # Offset right channel negative for classic display
            self.curve_r.setData(t, -r_max)
        else:
            mono = samples.reshape(-1, step)
            m_max = np.max(mono, axis=1)
            m_min = np.min(mono, axis=1)
            self.curve_l.setData(t, m_max)
            self.curve_r.setData(t, m_min)
            
        self.set_position(0)

    def set_position(self, pos_sec: float):
        """Update playhead position."""
        if 0 <= pos_sec <= self.duration_sec:
            self.playhead.setValue(pos_sec)

    def _on_mouse_clicked(self, event):
        """Handle click to seek."""
        if event.button() == 1: # Left click
            # Map scene coordinates to plot coordinates
            pos = self.plot_widget.plotItem.vb.mapSceneToView(event.scenePos())
            x_sec = pos.x()
            if 0 <= x_sec <= self.duration_sec:
                self.seek_requested.emit(x_sec)
                self.set_position(x_sec)