tape_recorder/ui/eq_panel.py

"""
EQ Panel — Interactive UI for the Parametric EQ Engine.
"""

from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, 
                             QPushButton, QLabel, QSlider, QGroupBox,
                             QCheckBox, QScrollArea, QComboBox, QFileDialog)
from PyQt6.QtCore import Qt, QSettings
import os
import pyqtgraph as pg
import numpy as np
import scipy.signal as signal
from core.eq_engine import EQEngine, EQBand
from core.tape_processor import TapeProcessor

class EQBandWidget(QGroupBox):
    """Widget for controlling a single EQ band."""
    def __init__(self, band: EQBand, index: int, parent_panel):
        super().__init__(f"Band {index + 1}")
        self.band = band
        self.parent_panel = parent_panel
        
        self.setFixedWidth(80)
        self._init_ui()
        self._update_ui_from_band()
        
    def _init_ui(self):
        layout = QVBoxLayout(self)
        layout.setContentsMargins(5, 10, 5, 5)
        layout.setSpacing(10)
        
        # Enabled toggle
        self.enabled_cb = QCheckBox("On")
        self.enabled_cb.clicked.connect(self._on_changed)
        layout.addWidget(self.enabled_cb, alignment=Qt.AlignmentFlag.AlignHCenter)
        
        # Gain slider (-15 to +15 dB)
        self.gain_slider = QSlider(Qt.Orientation.Vertical)
        self.gain_slider.setRange(-150, 150) # -15.0 to 15.0 dB
        self.gain_slider.setTickPosition(QSlider.TickPosition.TicksBothSides)
        self.gain_slider.setTickInterval(50)
        self.gain_slider.valueChanged.connect(self._on_gain_changed)
        layout.addWidget(self.gain_slider, alignment=Qt.AlignmentFlag.AlignHCenter)
        
        # Gain Label
        self.gain_lbl = QLabel("0.0 dB")
        self.gain_lbl.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.gain_lbl)
        
        # Freq
        freq_lbl = QLabel("Freq (Hz):")
        freq_lbl.setAlignment(Qt.AlignmentFlag.AlignCenter)
        freq_lbl.setStyleSheet("font-size: 8pt; color: #808090;")
        layout.addWidget(freq_lbl)
        
        self.freq_val = QLabel(f"{self.band.freq:.0f}")
        self.freq_val.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.freq_val)
        
        # Q
        q_lbl = QLabel("Q:")
        q_lbl.setAlignment(Qt.AlignmentFlag.AlignCenter)
        q_lbl.setStyleSheet("font-size: 8pt; color: #808090;")
        layout.addWidget(q_lbl)
        
        self.q_val = QLabel(f"{self.band.q:.2f}")
        self.q_val.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(self.q_val)
        
    def _update_ui_from_band(self):
        """Update widgets when the underlying band object changes (e.g. after calibration)."""
        self.enabled_cb.setChecked(self.band.enabled)
        # Block signals to avoid infinite loop
        self.gain_slider.blockSignals(True)
        self.gain_slider.setValue(int(self.band.gain * 10))
        self.gain_slider.blockSignals(False)
        self.gain_lbl.setText(f"{self.band.gain:.1f} dB")
        self.freq_val.setText(f"{self.band.freq:.0f}")
        self.q_val.setText(f"{self.band.q:.2f}")
        
    def _on_gain_changed(self, value):
        self.band.gain = value / 10.0
        self.gain_lbl.setText(f"{self.band.gain:.1f} dB")
        self.parent_panel.on_band_changed()
        
    def _on_changed(self):
        self.band.enabled = self.enabled_cb.isChecked()
        self.parent_panel.on_band_changed()

class EQPanel(QWidget):
    def __init__(self, eq_engine: EQEngine, tape_processor: TapeProcessor):
        super().__init__()
        self.eq_engine = eq_engine
        self.tape_processor = tape_processor
        self.band_widgets = []
        self.settings = QSettings("TapeRecorder", "EQ")
        self._init_ui()
        
        last_profile = self.settings.value("last_profile", "")
        if last_profile and os.path.exists(last_profile):
            try:
                self.eq_engine.load_bands(last_profile)
            except Exception as e:
                print(f"Failed to load last profile: {e}")
                
        self.refresh_ui()

    def _init_ui(self):
        layout = QVBoxLayout(self)
        
        # --- Top: Response Plot ---
        pg.setConfigOption('background', '#ffffff')
        pg.setConfigOption('foreground', '#333333')
        self.plot_widget = pg.PlotWidget(title="Combined EQ Response")
        self.plot_widget.setLogMode(x=True, y=False)
        self.plot_widget.setXRange(np.log10(20), np.log10(20000))
        self.plot_widget.setYRange(-20, 20)
        self.plot_widget.setLabel('bottom', 'Frequency', units='Hz')
        self.plot_widget.setLabel('left', 'Gain', units='dB')
        self.plot_widget.showGrid(x=True, y=True, alpha=0.3)
        
        self.curve = self.plot_widget.plot(pen=pg.mkPen('#ff7b00', width=2))
        self.plot_widget.setFixedHeight(250)
        
        layout.addWidget(self.plot_widget)
        
        # --- Middle: Band Controls ---
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        scroll.setStyleSheet("border: none;")
        
        bands_container = QWidget()
        self.bands_layout = QHBoxLayout(bands_container)
        self.bands_layout.setAlignment(Qt.AlignmentFlag.AlignLeft)
        
        scroll.setWidget(bands_container)
        layout.addWidget(scroll)
        
        # --- Bottom: Presets & Tools & Tape Processing ---
        bottom_layout = QHBoxLayout()
        
        # Tools Group
        tools_grp = QGroupBox("EQ Tools")
        tools_layout = QVBoxLayout(tools_grp)
        row1 = QHBoxLayout()
        
        self.bypass_cb = QCheckBox("Bypass EQ")
        self.bypass_cb.setStyleSheet("font-weight: bold; color: #F44336;")
        self.bypass_cb.clicked.connect(self._on_bypass)
        row1.addWidget(self.bypass_cb)
        row1.addStretch()
        
        row1.addWidget(QLabel("Presets:"))
        self.preset_combo = QComboBox()
        self.preset_combo.addItems(["Custom", "Flat", "Type I (Ferric) Boost", "Type II (Chrome) Cut"])
        self.preset_combo.currentIndexChanged.connect(self._on_preset)
        row1.addWidget(self.preset_combo)
        
        reset_btn = QPushButton("Reset to Flat")
        reset_btn.clicked.connect(self._reset_flat)
        row1.addWidget(reset_btn)
        tools_layout.addLayout(row1)
        
        row2 = QHBoxLayout()
        save_btn = QPushButton("Save Profile...")
        save_btn.clicked.connect(self._save_profile)
        load_btn = QPushButton("Load Profile...")
        load_btn.clicked.connect(self._load_profile)
        row2.addWidget(save_btn)
        row2.addWidget(load_btn)
        row2.addStretch()
        tools_layout.addLayout(row2)
        bottom_layout.addWidget(tools_grp)
        
        # Tape Processing Group
        tape_grp = QGroupBox("Tape Processing (Record Mode)")
        tape_layout = QVBoxLayout(tape_grp)
        
        t_row1 = QHBoxLayout()
        self.tape_hp_cb = QCheckBox("High-Pass (Sub-Bass Cut)")
        self.tape_hp_cb.setChecked(self.tape_processor.enable_highpass)
        self.tape_hp_cb.clicked.connect(self._on_tape_proc_changed)
        t_row1.addWidget(self.tape_hp_cb)
        
        self.tape_lp_cb = QCheckBox("Low-Pass (Ultrasonic Cut)")
        self.tape_lp_cb.setChecked(self.tape_processor.enable_lowpass)
        self.tape_lp_cb.clicked.connect(self._on_tape_proc_changed)
        t_row1.addWidget(self.tape_lp_cb)
        tape_layout.addLayout(t_row1)
        
        t_row2 = QHBoxLayout()
        self.tape_sat_cb = QCheckBox("Soft Saturation (Limiter)")
        self.tape_sat_cb.setChecked(self.tape_processor.enable_saturation)
        self.tape_sat_cb.clicked.connect(self._on_tape_proc_changed)
        t_row2.addWidget(self.tape_sat_cb)
        
        self.sat_thresh_slider = QSlider(Qt.Orientation.Horizontal)
        self.sat_thresh_slider.setRange(10, 200) # 0.1 to 2.0
        self.sat_thresh_slider.setValue(int(self.tape_processor.saturation_threshold * 100))
        self.sat_thresh_slider.setToolTip("Saturation Threshold")
        self.sat_thresh_slider.valueChanged.connect(self._on_sat_slider)
        t_row2.addWidget(self.sat_thresh_slider)
        tape_layout.addLayout(t_row2)
        
        bottom_layout.addWidget(tape_grp)
        
        layout.addLayout(bottom_layout)

    def refresh_ui(self):
        """Rebuild the UI to match the EQEngine's current bands."""
        # Clear existing widgets
        for i in reversed(range(self.bands_layout.count())): 
            self.bands_layout.itemAt(i).widget().setParent(None)
            
        self.band_widgets = []
        
        # Create new widgets
        for i, band in enumerate(self.eq_engine.bands):
            bw = EQBandWidget(band, i, self)
            self.bands_layout.addWidget(bw)
            self.band_widgets.append(bw)
            
        self.update_plot()
            
    def on_band_changed(self):
        """Called by a child BandWidget when a parameter changes."""
        self.preset_combo.blockSignals(True)
        self.preset_combo.setCurrentIndex(0) # Custom
        self.preset_combo.blockSignals(False)
        
        if not self.bypass_cb.isChecked():
            self.eq_engine.update_filters()
            
        self.update_plot()
        
    def _on_bypass(self):
        if self.bypass_cb.isChecked():
            # Clear filters in engine to bypass
            self.eq_engine._sos = None
            self.eq_engine._zi = None
        else:
            self.eq_engine.update_filters()
        self.update_plot()

    def _on_tape_proc_changed(self):
        self.tape_processor.enable_highpass = self.tape_hp_cb.isChecked()
        self.tape_processor.enable_lowpass = self.tape_lp_cb.isChecked()
        self.tape_processor.enable_saturation = self.tape_sat_cb.isChecked()
        self.tape_processor._update_filters()
        self.update_plot()
        
    def _on_sat_slider(self, val):
        self.tape_processor.saturation_threshold = val / 100.0

    def update_plot(self):
        """Recalculate and redraw the frequency response curve."""
        freqs = np.logspace(np.log10(20), np.log10(20000), 500)
        
        if self.bypass_cb.isChecked():
            response = np.zeros_like(freqs)
        else:
            # Need to temporarily generate SOS for the plot if bypassed, but here we just
            # Use the active SOS or calculate it from the bands.
            engine = EQEngine(self.eq_engine.sample_rate)
            
            all_bands = list(self.eq_engine.bands)
            if self.tape_processor.enable_highpass or self.tape_processor.enable_lowpass:
                all_bands.extend(self.tape_processor._filter_engine.bands)
                
            engine.bands = all_bands
            engine.update_filters()
            
            if engine._sos is None:
                response = np.zeros_like(freqs)
            else:
                w = 2 * np.pi * freqs / engine.sample_rate
                _, h = signal.sosfreqz(engine._sos, worN=w)
                response = 20 * np.log10(np.abs(h) + 1e-10)
                
        self.curve.setData(freqs, response)

    def _reset_flat(self):
        for band in self.eq_engine.bands:
            band.gain = 0.0
            
        for bw in self.band_widgets:
            bw._update_ui_from_band()
            
        self.on_band_changed()
        
    def _on_preset(self, index):
        if index == 1: # Flat
            self._reset_flat()
        elif index == 2: # Type I Boost (example)
            self._reset_flat()
            # Boost highs to compensate for tape roll-off
            if len(self.eq_engine.bands) > 0:
                self.eq_engine.bands[-1].gain = 6.0 # High shelf boost
        elif index == 3: # Type II Cut (example)
            self._reset_flat()
            # Cut highs if playing 70us tape on 120us setting
            if len(self.eq_engine.bands) > 0:
                self.eq_engine.bands[-1].gain = -4.5 # High shelf cut
                
        for bw in self.band_widgets:
            bw._update_ui_from_band()
            
        self.eq_engine.update_filters()
        self.update_plot()
        
    def _save_profile(self):
        path, _ = QFileDialog.getSaveFileName(self, "Save EQ Profile", "", "JSON Files (*.json)")
        if path:
            self.eq_engine.save_bands(path)
            self.settings.setValue("last_profile", path)
            
    def _load_profile(self):
        path, _ = QFileDialog.getOpenFileName(self, "Load EQ Profile", "", "JSON Files (*.json)")
        if path:
            try:
                self.eq_engine.load_bands(path)
                self.settings.setValue("last_profile", path)
                self.refresh_ui()
            except Exception as e:
                from PyQt6.QtWidgets import QMessageBox
                QMessageBox.warning(self, "Error", f"Failed to load profile:\n{e}")