dudong-v2/ui/tabs/ripeness_tab.py

"""
Ripeness Classification Tab

This tab handles audio file processing for ripeness classification.
Enhanced with Phase 8: Professional grid layout with multiple data visualization panels.
"""

from PyQt5.QtWidgets import QWidget, QGridLayout, QHBoxLayout, QMessageBox
from PyQt5.QtCore import pyqtSignal
from PyQt5.QtGui import QPixmap
import time

from ui.panels.rgb_preview_panel import RGBPreviewPanel
from ui.panels.multispectral_panel import MultispectralPanel
from ui.panels.audio_spectrogram_panel import AudioSpectrogramPanel
from ui.panels.ripeness_results_panel import RipenessResultsPanel
from ui.panels.ripeness_control_panel import RipenessControlPanel
from ui.panels.analysis_timeline_panel import AnalysisTimelinePanel
from utils.session_manager import SessionManager


class RipenessTab(QWidget):
    """
    Tab for ripeness classification using audio analysis.
    
    Features:
    - RGB Preview (Coming Soon)
    - Multispectral Analysis (Coming Soon)
    - Audio Spectrogram Display
    - Ripeness Results with Confidence Bars
    - Control Panel
    - Analysis Timeline & Statistics
    
    Signals:
        load_audio_requested: Emitted when user wants to load an audio file
    """
    
    load_audio_requested = pyqtSignal()
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.session_manager = SessionManager()
        self.current_file_path = None
        self.processing_start_time = None
        self.init_ui()
        
    def init_ui(self):
        """Initialize the UI components with 2x2+2 grid layout matching mockup."""
        # Main layout with content area styling to match mockup
        main_layout = QHBoxLayout(self)
        main_layout.setContentsMargins(20, 20, 20, 0)
        main_layout.setSpacing(20)
        
        # Left panels grid (2x2)
        left_grid = QGridLayout()
        left_grid.setSpacing(20)
        left_grid.setContentsMargins(0, 0, 0, 0)
        
        # Create left grid panels
        self.rgb_panel = RGBPreviewPanel()
        self.multispectral_panel = MultispectralPanel()
        self.audio_panel = AudioSpectrogramPanel()
        self.results_panel = RipenessResultsPanel()
        
        # Add panels to grid
        left_grid.addWidget(self.rgb_panel, 0, 0, 1, 1)
        left_grid.addWidget(self.multispectral_panel, 0, 1, 1, 1)
        left_grid.addWidget(self.audio_panel, 1, 0, 1, 1)
        left_grid.addWidget(self.results_panel, 1, 1, 1, 1)
        
        # Set column and row stretch factors for equal sizing
        left_grid.setColumnStretch(0, 1)
        left_grid.setColumnStretch(1, 1)
        left_grid.setRowStretch(0, 1)
        left_grid.setRowStretch(1, 1)
        
        # Right panels
        self.control_panel = RipenessControlPanel()
        self.timeline_panel = AnalysisTimelinePanel()
        
        # Add panels to main content layout with adjusted proportions
        # Increased left grid to make previews more square, reduced timeline
        main_layout.addLayout(left_grid, 3)  # Stretch factor 3 for 2x2 grid (larger previews)
        main_layout.addWidget(self.control_panel, 0)  # No stretch (fixed width ~200px)
        main_layout.addWidget(self.timeline_panel, 1)  # Stretch factor 1 (reduced from 2)
        
        # Connect signals
        self.control_panel.run_test_clicked.connect(self._on_run_test)
        self.control_panel.open_file_clicked.connect(self._on_run_test)
        self.control_panel.stop_clicked.connect(self._on_stop)
        self.control_panel.reset_clicked.connect(self._on_reset)
        self.timeline_panel.save_audio_clicked.connect(self._on_save_audio)
        
    def _on_run_test(self):
        """Handle RUN TEST button click."""
        self.processing_start_time = time.time()
        self.control_panel.set_processing(True)
        self.load_audio_requested.emit()
        
    def _on_stop(self):
        """Handle STOP button click."""
        self.control_panel.set_processing(False)
        # TODO: Implement worker cancellation
        
    def _on_reset(self):
        """Handle RESET button click."""
        # Clear all displays
        self.audio_panel.clear_spectrogram()
        self.results_panel.clear_results()
        self.current_file_path = None
        self.control_panel.set_processing(False)
        
    def _on_save_audio(self):
        """Handle save audio button click."""
        if self.current_file_path:
            QMessageBox.information(
                self,
                "Save Audio",
                "Audio save functionality coming in future update.\n\n"
                f"Would save: {self.current_file_path}"
            )
        
    def set_loading(self, is_loading: bool):
        """
        Set loading state.
        
        Args:
            is_loading: Whether processing is active
        """
        self.control_panel.set_processing(is_loading)
        
    def update_results(self, spectrogram: QPixmap, predicted_class: str, 
                      probabilities: dict, file_path: str):
        """
        Update the tab with processing results.
        
        Args:
            spectrogram: QPixmap of the spectrogram
            predicted_class: Predicted ripeness class
            probabilities: Dictionary of class probabilities (0-1 scale)
            file_path: Path to the audio file
        """
        # Calculate processing time
        processing_time = 0
        if self.processing_start_time:
            processing_time = time.time() - self.processing_start_time
            self.processing_start_time = None
        
        # Store current file path
        self.current_file_path = file_path
        
        # Update Audio Spectrogram Panel
        # TODO: Extract actual sample rate and duration from audio
        self.audio_panel.update_spectrogram(
            spectrogram, 
            sample_rate=44100, 
            duration=3.2, 
            audio_path=file_path
        )
        
        # Update Ripeness Results Panel
        self.results_panel.update_results(
            predicted_class,
            probabilities,
            processing_time,
            model_version="RipeNet v3.2"
        )
        
        # Convert probability to percentage for display (probabilities are now in decimal form)
        confidence = probabilities.get(predicted_class, 0) * 100
        self.timeline_panel.add_test_result(
            predicted_class,
            confidence,
            processing_time
        )
        
        # Add to session manager (probabilities are already in percentage form)
        self.session_manager.add_result(
            classification=predicted_class,
            confidence=confidence,
            probabilities=probabilities,
            processing_time=processing_time,
            file_path=file_path
        )
        
        # Update statistics
        stats = self.session_manager.get_statistics_summary()
        self.timeline_panel.update_statistics(
            total_tests=stats["total_tests"],
            avg_processing=stats["avg_processing_time"],
            ripe_count=stats["ripe_count"],
            session_start=stats["session_start"]
        )
        
        # Update processing state
        self.control_panel.set_processing(False)
        
    def clear_results(self):
        """Clear all displayed results and reset session."""
        self.audio_panel.clear_spectrogram()
        self.results_panel.clear_results()
        self.timeline_panel.clear_timeline()
        self.session_manager.clear_session()
        self.current_file_path = None
        self.processing_start_time = None
        
    def get_denoise_enabled(self) -> bool:
        """
        Check if audio denoising is enabled.
        
        Returns:
            bool: True if denoise checkbox is checked
        """
        return self.control_panel.denoise_checkbox.isChecked()