dudong-v2/ui/tabs/maturity_tab.py

"""
Maturity Classification Tab

This tab handles multispectral TIFF file processing for maturity classification.
"""

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

from ui.panels.rgb_preview_panel import RGBPreviewPanel
from ui.panels.multispectral_panel import MultispectralPanel
from ui.panels.maturity_results_panel import MaturityResultsPanel
from ui.panels.maturity_control_panel import MaturityControlPanel
from ui.panels.analysis_timeline_panel import AnalysisTimelinePanel
from utils.session_manager import SessionManager


class MaturityTab(QWidget):
    """
    Tab for maturity classification using multispectral analysis.
    
    Features:
    - RGB Preview (Coming Soon)
    - Multispectral Analysis with Grad-CAM
    - Maturity Results with Confidence Bars
    - Control Panel
    - Analysis Timeline & Statistics
    
    Signals:
        load_tiff_requested: Emitted when user wants to load a TIFF file
    """
    
    load_tiff_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.results_panel = MaturityResultsPanel()
        
        # Create a placeholder panel for the bottom-left (can be used for band visualization)
        self.band_panel = QWidget()
        self.band_panel.setStyleSheet("""
            QWidget {
                background-color: white;
                border: 1px solid #ddd;
            }
        """)
        
        # 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.band_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 = MaturityControlPanel()
        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_open_file)
        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_tiff)
        
    def _on_run_test(self):
        """Handle RUN TEST button click."""
        self.processing_start_time = time.time()
        self.control_panel.set_processing(True)
        # TODO: Implement live mode
        QMessageBox.information(
            self,
            "Live Mode",
            "Live multispectral camera mode coming soon!"
        )
        self.control_panel.set_processing(False)
        
    def _on_open_file(self):
        """Handle OPEN FILE button click."""
        self.processing_start_time = time.time()
        self.control_panel.set_processing(True)
        self.load_tiff_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.multispectral_panel.clear() if hasattr(self.multispectral_panel, 'clear') else None
        self.results_panel.clear_results()
        self.current_file_path = None
        self.control_panel.set_processing(False)
        
    def _on_save_tiff(self):
        """Handle save TIFF button click."""
        if self.current_file_path:
            QMessageBox.information(
                self,
                "Save TIFF",
                "TIFF 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, gradcam_image: QImage, predicted_class: str, 
                      probabilities: dict, file_path: str):
        """
        Update the tab with processing results.
        
        Args:
            gradcam_image: QImage of the Grad-CAM visualization
            predicted_class: Predicted maturity class
            probabilities: Dictionary of class probabilities (0-1 scale)
            file_path: Path to the TIFF 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 Multispectral Panel with Grad-CAM
        if gradcam_image:
            pixmap = QPixmap.fromImage(gradcam_image)
            # Update the multispectral panel
            self.multispectral_panel.set_image(pixmap)
        
        # Update Maturity Results Panel
        self.results_panel.update_results(
            predicted_class,
            probabilities,
            processing_time,
            model_version="MaturityNet v1.0"
        )
        
        # Convert probability to percentage for display
        confidence = probabilities.get(predicted_class, 0) * 100
        if not confidence:
            # Try case-insensitive match
            for key, value in probabilities.items():
                if key.lower() == predicted_class.lower():
                    confidence = value * 100
                    break
        
        self.timeline_panel.add_test_result(
            predicted_class,
            confidence,
            processing_time
        )
        
        # Add to session manager
        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"],  # Reuse ripe_count for maturity 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.multispectral_panel.clear() if hasattr(self.multispectral_panel, 'clear') else None
        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