tape_recorder/ui/calibration_panel.py

"""
Calibration Panel — Step-by-step wizard for frequency response measurement and EQ fitting.
"""

from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, 
                             QPushButton, QLabel, QComboBox, 
                             QProgressBar, QStackedWidget,
                             QMessageBox, QGroupBox)
from PyQt6.QtCore import Qt, QTimer, pyqtSignal
import pyqtgraph as pg
import numpy as np
import threading
import time

from core.audio_device import DeviceManager
from core.signal_generator import generate_log_sweep, generate_pink_noise
from core.analyzer import calculate_transfer_function, smooth_curve
from core.calibrator import calculate_correction_curve, fit_eq_bands
from core.eq_engine import EQEngine
from core.tape_processor import TapeProcessor
import sounddevice as sd

class CalibrationPanel(QWidget):
    
    # Signal emitted when calibration finishes and new EQ bands are ready
    calibration_complete = pyqtSignal(list) # list of EQBand
    
    # Internal signals for cross-thread UI updates
    _update_plots_signal = pyqtSignal(np.ndarray, np.ndarray, np.ndarray, np.ndarray)
    _fit_complete_signal = pyqtSignal()
    _set_play_btn_signal = pyqtSignal(bool, str)
    _set_prog_signal = pyqtSignal(bool)
    _set_analyze_btn_signal = pyqtSignal(str)
    _set_apply_btn_signal = pyqtSignal(str)
    
    def __init__(self, device_manager: DeviceManager, eq_engine: EQEngine, tape_processor: TapeProcessor):
        super().__init__()
        self.device_manager = device_manager
        self.eq_engine = eq_engine
        self.tape_processor = tape_processor
        
        # State
        self.reference_signal = None
        self.recorded_signal = None
        self.freqs = None
        self.response_db = None
        self.correction_db = None
        
        self._update_plots_signal.connect(self._update_plots_slot)
        self._fit_complete_signal.connect(self._fit_complete_slot)
        self._set_play_btn_signal.connect(self._set_play_btn_slot)
        self._set_prog_signal.connect(self._set_prog_slot)
        self._set_analyze_btn_signal.connect(self._set_analyze_btn_slot)
        self._set_apply_btn_signal.connect(self._set_apply_btn_slot)
        
        self._init_ui()
        
    def _init_ui(self):
        layout = QVBoxLayout(self)
        
        # --- Wizard Header ---
        header = QLabel("Calibration Wizard")
        header.setStyleSheet("font-size: 16pt; font-weight: bold; color: #ff7b00;")
        layout.addWidget(header)
        
        # --- Steps Stack ---
        self.stack = QStackedWidget()
        
        # Step 1: Generate & Play
        self.step1 = self._create_step1()
        self.stack.addWidget(self.step1)
        
        # Step 2: Record
        self.step2 = self._create_step2()
        self.stack.addWidget(self.step2)
        
        # Step 3: Analyze & Apply
        self.step3 = self._create_step3()
        self.stack.addWidget(self.step3)
        
        layout.addWidget(self.stack, 1)
        
        # --- Navigation ---
        nav_row = QHBoxLayout()
        self.prev_btn = QPushButton("← Back")
        self.prev_btn.clicked.connect(self._prev_step)
        self.prev_btn.setEnabled(False)
        
        self.next_btn = QPushButton("Next →")
        self.next_btn.setObjectName("primaryButton")
        self.next_btn.clicked.connect(self._next_step)
        
        nav_row.addWidget(self.prev_btn)
        nav_row.addStretch()
        nav_row.addWidget(self.next_btn)
        
        layout.addLayout(nav_row)

    def _create_step1(self) -> QWidget:
        w = QWidget()
        l = QVBoxLayout(w)
        
        info = QLabel(
            "Step 1: Generate Test Signal\n\n"
            "Prepare your cassette deck to record. "
            "A log sine sweep provides the best accuracy for calibration."
        )
        info.setWordWrap(True)
        l.addWidget(info)
        
        grp = QGroupBox("Signal & Tape Settings")
        gl = QVBoxLayout(grp)
        
        row0 = QHBoxLayout()
        row0.addWidget(QLabel("Tape Type:"))
        self.tape_type_cb = QComboBox()
        self.tape_type_cb.addItems(["Type I (Ferric)", "Type II (Chrome)", "Type IV (Metal)"])
        self.tape_type_cb.currentTextChanged.connect(self._on_tape_type_changed)
        row0.addWidget(self.tape_type_cb)
        row0.addStretch()
        gl.addLayout(row0)
        
        row1 = QHBoxLayout()
        row1.addWidget(QLabel("Signal Type:"))
        self.sig_type = QComboBox()
        self.sig_type.addItems(["Log Sine Sweep (Recommended)", "Pink Noise"])
        row1.addWidget(self.sig_type)
        row1.addStretch()
        gl.addLayout(row1)
        
        row2 = QHBoxLayout()
        row2.addWidget(QLabel("Duration:"))
        self.sig_dur = QComboBox()
        self.sig_dur.addItems(["10 Seconds", "20 Seconds", "30 Seconds"])
        row2.addWidget(self.sig_dur)
        row2.addStretch()
        gl.addLayout(row2)
        
        l.addWidget(grp)
        
        self.play_test_btn = QPushButton("Generate & Play Signal")
        self.play_test_btn.setObjectName("primaryButton")
        self.play_test_btn.clicked.connect(self._play_test_signal)
        l.addWidget(self.play_test_btn)
        
        self.play_prog = QProgressBar()
        self.play_prog.setTextVisible(False)
        self.play_prog.hide()
        l.addWidget(self.play_prog)
        
        l.addStretch()
        return w

    def _create_step2(self) -> QWidget:
        w = QWidget()
        l = QVBoxLayout(w)
        
        info = QLabel(
            "Step 2: Record Playback\n\n"
            "If you have a 3-head deck, you can record this immediately. "
            "For 2-head decks, rewind the tape and press play."
        )
        info.setWordWrap(True)
        l.addWidget(info)
        
        # Realtime view for recording
        pg.setConfigOption('background', '#161625')
        self.rec_plot = pg.PlotWidget()
        self.rec_plot.setYRange(-1, 1)
        self.rec_curve = self.rec_plot.plot(pen='#ff7b00')
        l.addWidget(self.rec_plot, 1)
        
        self.rec_btn = QPushButton("Start Recording Capture")
        self.rec_btn.setObjectName("primaryButton")
        self.rec_btn.clicked.connect(self._toggle_recording)
        l.addWidget(self.rec_btn)
        
        return w

    def _create_step3(self) -> QWidget:
        w = QWidget()
        l = QVBoxLayout(w)
        
        self.analyze_btn = QPushButton("Analyze Recorded Signal")
        self.analyze_btn.setObjectName("primaryButton")
        self.analyze_btn.clicked.connect(self._analyze)
        l.addWidget(self.analyze_btn)
        
        # Result plot
        self.res_plot = pg.PlotWidget(title="Frequency Response & EQ Target")
        self.res_plot.setLogMode(x=True, y=False)
        self.res_plot.setXRange(np.log10(20), np.log10(20000))
        self.res_plot.setYRange(-20, 20)
        self.res_plot.showGrid(x=True, y=True, alpha=0.3)
        self.res_plot.addLegend()
        
        self.curve_raw = self.res_plot.plot(pen=pg.mkPen('#505050', width=1), name="Raw Response")
        self.curve_smooth = self.res_plot.plot(pen=pg.mkPen('#ff7b00', width=2), name="Smoothed Response")
        self.curve_target = self.res_plot.plot(pen=pg.mkPen('#4CAF50', width=2, style=Qt.PenStyle.DashLine), name="EQ Target")
        
        l.addWidget(self.res_plot, 1)
        
        self.apply_eq_btn = QPushButton("Auto-Fit EQ & Apply")
        self.apply_eq_btn.setObjectName("primaryButton")
        self.apply_eq_btn.setEnabled(False)
        self.apply_eq_btn.clicked.connect(self._apply_eq)
        l.addWidget(self.apply_eq_btn)
        
        return w

    def _set_play_btn_slot(self, enabled, text):
        self.play_test_btn.setEnabled(enabled)
        self.play_test_btn.setText(text)
        
    def _set_prog_slot(self, show):
        self.play_prog.setVisible(show)
        
    def _set_analyze_btn_slot(self, text):
        self.analyze_btn.setText(text)
        
    def _set_apply_btn_slot(self, text):
        self.apply_eq_btn.setText(text)
        
    def _on_tape_type_changed(self, tape_type: str):
        self.tape_processor.set_tape_type(tape_type)

    # --- Step 1 Logic ---
    def _play_test_signal(self):
        if self.device_manager.output_device_id is None:
            QMessageBox.warning(self, "Error", "No output device selected in Devices tab.")
            return
            
        dur = int(self.sig_dur.currentText().split()[0])
        sr = self.device_manager.sample_rate
        
        if "Sweep" in self.sig_type.currentText():
            self.reference_signal = generate_log_sweep(duration=dur, sample_rate=sr)
        else:
            self.reference_signal = generate_pink_noise(duration=dur, sample_rate=sr)
            
        # Play in background thread
        threading.Thread(target=self._play_thread, args=(self.reference_signal, sr), daemon=True).start()
        
    def _play_thread(self, signal_data, sr):
        # UI updates
        self._set_play_btn_signal.emit(False, "Playing...")
        self._set_prog_signal.emit(True)
        
        try:
            sd.play(signal_data, sr, device=self.device_manager.output_device_id)
            sd.wait()
        except Exception as e:
            print(f"Playback error: {e}")
            
        self._set_play_btn_signal.emit(True, "Generate & Play Signal")
        self._set_prog_signal.emit(False)

    # --- Step 2 Logic ---
    def _toggle_recording(self):
        if hasattr(self, 'is_recording') and self.is_recording:
            # Stop
            self.is_recording = False
            self.rec_btn.setText("Start Recording Capture")
            self.rec_btn.setStyleSheet("")
            if hasattr(self, 'rec_stream'):
                self.rec_stream.stop()
                self.rec_stream.close()
                self.recorded_signal = np.concatenate(self.rec_chunks, axis=0)
                # Take left channel if stereo
                if self.recorded_signal.ndim > 1:
                    self.recorded_signal = self.recorded_signal[:, 0]
        else:
            # Start
            if self.device_manager.input_device_id is None:
                QMessageBox.warning(self, "Error", "No input device selected.")
                return
                
            self.is_recording = True
            self.rec_btn.setText("Stop Recording")
            self.rec_btn.setStyleSheet("background-color: #F44336; color: white;")
            self.rec_chunks = []
            
            kwargs = self.device_manager.get_sd_input_kwargs()
            self.rec_stream = sd.InputStream(callback=self._rec_callback, **kwargs)
            self.rec_stream.start()
            
    def _rec_callback(self, indata, frames, time, status):
        self.rec_chunks.append(indata.copy())
        # Update plot with decimated data for performance
        if len(self.rec_chunks) % 5 == 0:
            chunk = indata[:, 0]
            decimated = chunk[::max(1, len(chunk)//500)]
            from PyQt6.QtCore import QMetaObject, Q_ARG
            # Thread-safe plot update (simplified for UI)
            pass

    # --- Step 3 Logic ---
    def _analyze(self):
        if self.reference_signal is None or self.recorded_signal is None:
            QMessageBox.warning(self, "Error", "Missing reference or recorded signal.")
            return
            
        self.analyze_btn.setText("Analyzing...")
        self.analyze_btn.setEnabled(False)
        
        # Run in thread
        threading.Thread(target=self._analyze_thread, daemon=True).start()
        
    def _analyze_thread(self):
        try:
            sr = self.device_manager.sample_rate
            
            # 1. Transfer function
            freqs, resp = calculate_transfer_function(self.reference_signal, self.recorded_signal, sr)
            
            # 2. Smooth
            resp_smooth = smooth_curve(freqs, resp, fraction=1/3)
            
            # 3. Target correction
            correction = calculate_correction_curve(resp_smooth)
            
            self.freqs = freqs
            self.response_db = resp_smooth
            self.correction_db = correction
            
            # Update UI
            self._update_plots_signal.emit(freqs, resp, resp_smooth, correction)
        except Exception as e:
            print(f"Analysis error: {e}")
            self._set_analyze_btn_signal.emit("Analysis Failed")
            
    def _update_plots_slot(self, freqs, resp_raw, resp_smooth, correction):
        self.curve_raw.setData(freqs, resp_raw)
        self.curve_smooth.setData(freqs, resp_smooth)
        self.curve_target.setData(freqs, correction)
        
        self.analyze_btn.setText("Analysis Complete")
        self.apply_eq_btn.setEnabled(True)
        
    def _apply_eq(self):
        if self.freqs is None or self.correction_db is None:
            return
            
        self.apply_eq_btn.setText("Fitting EQ...")
        self.apply_eq_btn.setEnabled(False)
        
        # Fit in thread
        threading.Thread(target=self._fit_thread, daemon=True).start()
        
    def _fit_thread(self):
        sr = self.device_manager.sample_rate
        try:
            ttype = self.tape_type_cb.currentText()
            # Fit 10 peaking bands constrained by tape type
            bands = fit_eq_bands(self.freqs, self.correction_db, num_peaking_bands=10, sample_rate=sr, tape_type=ttype)
            
            # Apply to engine
            self.eq_engine.bands = bands
            self.eq_engine.update_filters()
            
            self._fit_complete_signal.emit()
            
        except Exception as e:
            print(f"Fitting error: {e}")
            self._set_apply_btn_signal.emit("Fit Failed")
            
    def _fit_complete_slot(self):
        self.apply_eq_btn.setText("EQ Applied ✓")
        self.calibration_complete.emit(self.eq_engine.bands)
        QMessageBox.information(self, "Success", "EQ calibrated and applied to engine!")

    # --- Navigation Logic ---
    def _prev_step(self):
        idx = self.stack.currentIndex()
        if idx > 0:
            self.stack.setCurrentIndex(idx - 1)
            self._update_nav_buttons()

    def _next_step(self):
        idx = self.stack.currentIndex()
        if idx < self.stack.count() - 1:
            self.stack.setCurrentIndex(idx + 1)
            self._update_nav_buttons()

    def _update_nav_buttons(self):
        idx = self.stack.currentIndex()
        self.prev_btn.setEnabled(idx > 0)
        self.next_btn.setEnabled(idx < self.stack.count() - 1)