dudong-v2/ui/tabs/reports_tab.py

"""
Reports & Export Tab

Display analysis results and export functionality.

This module provides the main Reports Tab UI, delegating report generation,
visualization, export, and printing to specialized modules.
"""

from datetime import datetime
from typing import Optional, Dict, List, Tuple

from PyQt5.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
    QFrame, QMessageBox, QStackedWidget, QScrollArea
)
from PyQt5.QtCore import Qt, pyqtSignal
from PyQt5.QtGui import QFont, QPixmap, QImage

from resources.styles import COLORS
from ui.widgets.loading_screen import LoadingScreen
from ui.dialogs import PrintOptionsDialog

# Import refactored modules
from ui.components.report_generator import (
    generate_basic_report,
    generate_model_report,
    generate_comprehensive_report,
    extract_report_content
)
from ui.components.pdf_exporter import PDFExporter
from ui.components.report_printer import ReportPrinter


class ReportsTab(QWidget):
    """Tab for displaying analysis reports and exporting data."""
    
    # Signal to notify when user wants to go back to dashboard
    go_to_dashboard = pyqtSignal()
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.current_data = None
        self.input_data = None
        self.maturity_result = None
        self.current_analysis_id = None  # Set by main_window for data persistence
        self.current_report_id = None  # Set by main_window - the actual report ID
        self.data_manager = None  # Set by main_window for data persistence
        self.current_overall_grade = 'B'  # Track current grade for finalization
        self.current_grade_description = ''  # Track grade description
        self.has_result = False  # Track if we have a report result
        self.current_results = None  # Store results dictionary for print/export
        
        # Initialize export/print modules
        self.pdf_exporter = PDFExporter()
        self.report_printer = ReportPrinter()
        
        self.init_ui()
    
    def init_ui(self):
        """Initialize the UI components."""
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(10, 10, 10, 10)
        
        # Header with title and action buttons
        header_layout = QHBoxLayout()
        
        title = QLabel("📊 Analysis Report")
        title.setFont(QFont("Arial", 18, QFont.Bold))
        header_layout.addWidget(title)
        
        header_layout.addStretch()
        
        # Export PDF button
        self.export_btn = QPushButton("📄 Export PDF")
        self.export_btn.setStyleSheet("""
            QPushButton {
                background-color: #3498db;
                color: white;
                font-weight: bold;
                padding: 8px 16px;
                border-radius: 4px;
                font-size: 12px;
            }
            QPushButton:hover {
                background-color: #2980b9;
            }
            QPushButton:disabled {
                background-color: #95a5a6;
            }
        """)
        self.export_btn.clicked.connect(self.export_pdf)
        self.export_btn.setEnabled(False)
        header_layout.addWidget(self.export_btn)
        
        main_layout.addLayout(header_layout)
        
        # Separator
        separator = QFrame()
        separator.setFrameShape(QFrame.HLine)
        separator.setStyleSheet("background-color: #bdc3c7;")
        main_layout.addWidget(separator)
        
        # Stacked widget for content and loading screen
        self.stacked_widget = QStackedWidget()
        
        # Loading screen (index 0) - NOT started yet, will be started when needed
        self.loading_screen = LoadingScreen("Processing analysis...")
        self.loading_screen.stop_animation()  # Stop by default, will start when processing begins
        self.stacked_widget.addWidget(self.loading_screen)
        
        # Report content (index 1)
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        scroll.setStyleSheet("QScrollArea { border: none; }")
        
        # Report content widget
        self.report_widget = QWidget()
        self.report_layout = QVBoxLayout(self.report_widget)
        self.report_layout.setContentsMargins(10, 10, 10, 10)
        
        # Create empty state with button
        self._create_empty_state()
        
        scroll.setWidget(self.report_widget)
        self.stacked_widget.addWidget(scroll)
        
        main_layout.addWidget(self.stacked_widget)
        
        # Show content view by default (not loading)
        self.stacked_widget.setCurrentIndex(1)
    
    def _create_empty_state(self):
        """Create the empty state UI with button to go to dashboard."""
        # Clear any existing widgets
        while self.report_layout.count() > 0:
            item = self.report_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()
        
        self.report_layout.addStretch()
        
        # Icon/Title
        icon_label = QLabel("📊")
        icon_label.setFont(QFont("Arial", 48))
        icon_label.setAlignment(Qt.AlignCenter)
        self.report_layout.addWidget(icon_label)
        
        # Message
        message_label = QLabel("No Analysis Yet")
        message_label.setFont(QFont("Arial", 18, QFont.Bold))
        message_label.setAlignment(Qt.AlignCenter)
        message_label.setStyleSheet(f"color: {COLORS['text_primary']};")
        self.report_layout.addWidget(message_label)
        
        # Description
        desc_label = QLabel("Click 'Analyze Durian' on the dashboard to start an analysis")
        desc_label.setFont(QFont("Arial", 12))
        desc_label.setAlignment(Qt.AlignCenter)
        desc_label.setStyleSheet(f"color: {COLORS['text_secondary']}; margin: 10px;")
        desc_label.setWordWrap(True)
        self.report_layout.addWidget(desc_label)
        
        # Button container
        button_layout = QHBoxLayout()
        button_layout.addStretch()
        
        # Go to Dashboard button
        go_dashboard_btn = QPushButton("🏠 Go to Dashboard")
        go_dashboard_btn.setFont(QFont("Arial", 12, QFont.Bold))
        go_dashboard_btn.setStyleSheet("""
            QPushButton {
                background: qlineargradient(x1:0, y1:0, x2:0, y2:1,
                    stop:0 #3498db, stop:1 #2980b9);
                color: white;
                padding: 12px 24px;
                border-radius: 6px;
                border: none;
            }
            QPushButton:hover {
                background: qlineargradient(x1:0, y1:0, x2:0, y2:1,
                    stop:0 #2980b9, stop:1 #1f618d);
            }
            QPushButton:pressed {
                background: #1f618d;
            }
        """)
        go_dashboard_btn.clicked.connect(self.go_to_dashboard.emit)
        button_layout.addWidget(go_dashboard_btn)
        
        button_layout.addStretch()
        self.report_layout.addLayout(button_layout)
        
        self.report_layout.addStretch()
        
        self.has_result = False
    
    def _add_result_action_button(self):
        """Add action button to go to dashboard for a new analysis."""
        # Find and remove existing action button if any
        for i in range(self.report_layout.count() - 1, -1, -1):
            item = self.report_layout.itemAt(i)
            if item and item.widget() and hasattr(item.widget(), 'objectName'):
                if item.widget().objectName() == "result_action_button":
                    item.widget().deleteLater()
                    break
        
        # Add button container at the end
        action_layout = QHBoxLayout()
        action_layout.addStretch()
        
        # New Analysis button
        new_analysis_btn = QPushButton("🔄 New Analysis")
        new_analysis_btn.setObjectName("result_action_button")
        new_analysis_btn.setFont(QFont("Arial", 11, QFont.Bold))
        new_analysis_btn.setStyleSheet("""
            QPushButton {
                background: qlineargradient(x1:0, y1:0, x2:0, y2:1,
                    stop:0 #2ecc71, stop:1 #27ae60);
                color: white;
                padding: 10px 20px;
                border-radius: 5px;
                border: none;
            }
            QPushButton:hover {
                background: qlineargradient(x1:0, y1:0, x2:0, y2:1,
                    stop:0 #27ae60, stop:1 #1e8449);
            }
            QPushButton:pressed {
                background: #1e8449;
            }
        """)
        new_analysis_btn.clicked.connect(self.go_to_dashboard.emit)
        action_layout.addWidget(new_analysis_btn)
        
        action_layout.addStretch()
        self.report_layout.addLayout(action_layout)
        self.has_result = True
    
    def clear_report(self):
        """Clear the current report and show empty state."""
        # Remove all widgets
        while self.report_layout.count() > 0:
            item = self.report_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()
        
        # Show empty state
        self._create_empty_state()
        
        self.export_btn.setEnabled(False)
        self.current_data = None
        self.has_result = False
    
    def set_loading(self, is_loading: bool, message: str = "Processing analysis..."):
        """
        Set loading state for report generation.
        
        Args:
            is_loading: True to show loading screen, False to show content
            message: Loading message to display
        """
        self.export_btn.setEnabled(not is_loading)
        
        if is_loading:
            # Show loading screen with animation
            self.loading_screen.set_message(message)
            self.loading_screen.set_status("")
            self.loading_screen.spinner_index = 0  # Reset spinner
            self.stacked_widget.setCurrentIndex(0)
            self.loading_screen.start_animation()
        else:
            # Hide loading screen and show content
            self.loading_screen.stop_animation()
            self.stacked_widget.setCurrentIndex(1)
    
    def generate_report(self, input_data: Dict[str, str]):
        """
        Generate analysis report from input data.
        
        Args:
            input_data: Dictionary with keys 'dslr', 'multispectral', 'thermal', 'audio'
        """
        self.current_data = input_data
        self.input_data = input_data
        
        # Clear existing report
        while self.report_layout.count() > 0:
            item = self.report_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()
        
        # Generate report using report generator
        result = generate_basic_report(self.report_layout, input_data)
        self.current_overall_grade = result['grade']
        self.current_grade_description = result['description']
        
        self.report_layout.addStretch()
        
        # Add action button for new analysis
        self._add_result_action_button()
        
        # Stop loading screen and show content
        self.loading_screen.stop_animation()
        self.stacked_widget.setCurrentIndex(1)
        
        # Enable export/print buttons
        self.export_btn.setEnabled(True)
    
    def generate_report_with_model(
        self,
        input_data: Dict[str, str],
        gradcam_image: QImage,
        predicted_class: str,
        confidence: float,
        probabilities: Dict[str, float]
    ):
        """
        Generate analysis report with actual multispectral model results.
        
        Args:
            input_data: Dictionary with keys 'dslr', 'multispectral', 'thermal', 'audio'
            gradcam_image: QImage of the Grad-CAM visualization
            predicted_class: Predicted maturity class from model
            confidence: Model confidence (0-1 scale)
            probabilities: Dictionary of class probabilities
        """
        self.current_data = input_data
        self.input_data = input_data
        self.maturity_result = {
            'class': predicted_class,
            'confidence': confidence,
            'probabilities': probabilities,
            'gradcam_image': gradcam_image
        }
        
        # Clear existing report
        while self.report_layout.count() > 0:
            item = self.report_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()
        
        # Generate report using report generator
        result = generate_model_report(
            self.report_layout,
            input_data,
            gradcam_image,
            predicted_class,
            confidence,
            probabilities
        )
        self.current_overall_grade = result['grade']
        self.current_grade_description = result['description']
        
        self.report_layout.addStretch()
        
        # Add action button for new analysis
        self._add_result_action_button()
        
        # Stop loading screen and show content
        self.loading_screen.stop_animation()
        self.stacked_widget.setCurrentIndex(1)
        
        # Enable export/print buttons
        self.export_btn.setEnabled(True)
    
    def generate_report_with_rgb_and_multispectral(self, input_data: Dict[str, str], results: Dict):
        """
        Generate comprehensive report with RGB models and multispectral analysis.
        
        Args:
            input_data: Dictionary with input file paths
            results: Dictionary with processing results from all models
                     Keys: 'defect', 'locule', 'maturity', 'shape', 'audio'
        """
        self.current_data = input_data
        self.input_data = input_data
        self.current_results = results  # Store results for print/export
        
        # Clear existing report
        while self.report_layout.count() > 0:
            item = self.report_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()
        
        # Generate report using report generator
        result = generate_comprehensive_report(
            self.report_layout,
            input_data,
            results,
            self.current_report_id
        )
        self.current_overall_grade = result['grade']
        self.current_grade_description = result['description']
        
        self.report_layout.addStretch()
        
        # Add action button for new analysis
        self._add_result_action_button()
        
        # Stop loading screen and show content
        self.loading_screen.stop_animation()
        self.stacked_widget.setCurrentIndex(1)
        
        # Enable export/print buttons
        self.export_btn.setEnabled(True)
    
    def _get_report_visualizations(self) -> List[Tuple[str, QPixmap]]:
        """
        Extract all visualizations from the current report.
        
        Returns:
            List of tuples (title, QPixmap)
        """
        visualizations = []
        
        if not self.current_results:
            return visualizations
        
        results = self.current_results
        
        # Grad-CAM visualization
        if 'maturity' in results and results['maturity'].get('gradcam_image'):
            gradcam_img = results['maturity']['gradcam_image']
            if isinstance(gradcam_img, QImage):
                gradcam_pixmap = QPixmap.fromImage(gradcam_img)
            else:
                gradcam_pixmap = gradcam_img
            if not gradcam_pixmap.isNull():
                visualizations.append(("Grad-CAM Visualization (Maturity)", gradcam_pixmap))
        
        # Locule analysis visualization
        if 'locule' in results and results['locule'].get('annotated_image'):
            locule_pixmap = results['locule']['annotated_image']
            if not locule_pixmap.isNull():
                visualizations.append(("Locule Analysis (Top View)", locule_pixmap))
        
        # Defect analysis visualization
        if 'defect' in results and results['defect'].get('annotated_image'):
            defect_pixmap = results['defect']['annotated_image']
            if not defect_pixmap.isNull():
                visualizations.append(("Defect Analysis (Side View)", defect_pixmap))
        
        # Shape analysis visualization
        if 'shape' in results and results['shape'].get('annotated_image'):
            shape_pixmap = results['shape']['annotated_image']
            if not shape_pixmap.isNull():
                visualizations.append(("Shape Classification (Side View)", shape_pixmap))
        
        # Audio waveform visualization
        if 'audio' in results and results['audio'].get('waveform_image'):
            waveform_pixmap = results['audio']['waveform_image']
            if not waveform_pixmap.isNull():
                visualizations.append(("Waveform with Detected Knocks", waveform_pixmap))
        
        # Audio spectrogram visualization
        if 'audio' in results and results['audio'].get('spectrogram_image'):
            spectrogram_pixmap = results['audio']['spectrogram_image']
            if not spectrogram_pixmap.isNull():
                visualizations.append(("Mel Spectrogram", spectrogram_pixmap))
        
        return visualizations
    
    def export_pdf(self):
        """Export report as PDF using reportlab."""
        # Check if report data is available
        if not self.has_result:
            QMessageBox.warning(
                self,
                "No Report",
                "No analysis report available to export. Please complete an analysis first."
            )
            return
        
        # Ask about visualizations BEFORE file dialog
        options_dialog = PrintOptionsDialog(self)
        if options_dialog.exec_() != PrintOptionsDialog.Accepted:
            return
        
        include_visualizations = options_dialog.get_include_visualizations()
        
        # Extract report content
        report_content = extract_report_content(
            self.current_report_id or f"DUR-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
            self.current_overall_grade,
            self.current_grade_description,
            self.current_results
        )
        
        # Get visualizations
        visualizations = self._get_report_visualizations()
        
        # Export using PDF exporter
        self.pdf_exporter.export_report(
            self,
            self.current_report_id or datetime.now().strftime('%Y%m%d_%H%M%S'),
            report_content,
            visualizations,
            include_visualizations
        )
    
    def print_report(self):
        """Print report with options dialog."""
        # Check if report data is available
        if not self.has_result:
            QMessageBox.warning(
                self,
                "No Report",
                "No analysis report available to print. Please complete an analysis first."
            )
            return
        
        # Show print options dialog
        options_dialog = PrintOptionsDialog(self)
        if options_dialog.exec_() != PrintOptionsDialog.Accepted:
            return
        
        include_visualizations = options_dialog.get_include_visualizations()
        
        # Extract report content
        report_content = extract_report_content(
            self.current_report_id or f"DUR-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
            self.current_overall_grade,
            self.current_grade_description,
            self.current_results
        )
        
        # Get visualizations
        visualizations = self._get_report_visualizations()
        
        # Print using report printer
        self.report_printer.print_report(
            self,
            report_content,
            visualizations,
            include_visualizations
        )
    
    def load_analysis_from_db(self, report_id: str) -> bool:
        """
        Load a saved analysis from the database and display it.
        
        Args:
            report_id: Report ID to load
            
        Returns:
            bool: True if successfully loaded, False otherwise
        """
        if not self.data_manager:
            QMessageBox.warning(self, "Error", "Data manager not initialized")
            return False
        
        # Export analysis data with full file paths
        analysis_data = self.data_manager.export_analysis_to_dict(report_id)
        if not analysis_data:
            QMessageBox.warning(self, "Error", f"Analysis not found: {report_id}")
            return False
        
        # Reconstruct input_data dictionary for report generation
        input_data = {}
        for input_item in analysis_data['inputs']:
            input_type = input_item['input_type']
            full_path = input_item.get('full_path')
            if full_path:
                input_data[input_type] = full_path
        
        # Reconstruct results dictionary with visualizations
        results = {}
        for result in analysis_data['results']:
            model_type = result['model_type']
            
            # Find visualizations for this result
            result_data = {
                'predicted_class': result['predicted_class'],
                'confidence': result['confidence'],
                'probabilities': result['probabilities'],
                'metadata': result['metadata']
            }
            
            # Load annotated images
            if model_type == 'defect':
                # Find defect_annotated visualization
                viz = next((v for v in analysis_data['visualizations'] 
                           if v['visualization_type'] == 'defect_annotated'), None)
                if viz and viz.get('full_path'):
                    result_data['annotated_image'] = QPixmap(viz['full_path'])
                result_data['primary_class'] = result['predicted_class']
                result_data['total_detections'] = result['metadata'].get('total_detections', 0)
                result_data['class_counts'] = result['metadata'].get('class_counts', {})
            
            elif model_type == 'locule':
                # Find locule_annotated visualization
                viz = next((v for v in analysis_data['visualizations'] 
                           if v['visualization_type'] == 'locule_annotated'), None)
                if viz and viz.get('full_path'):
                    result_data['annotated_image'] = QPixmap(viz['full_path'])
                result_data['locule_count'] = result['metadata'].get('locule_count', 0)
            
            elif model_type == 'maturity':
                # Find maturity_gradcam visualization
                viz = next((v for v in analysis_data['visualizations'] 
                           if v['visualization_type'] == 'maturity_gradcam'), None)
                if viz and viz.get('full_path'):
                    result_data['gradcam_image'] = QImage(viz['full_path'])
                result_data['class_name'] = result['predicted_class']
            
            elif model_type == 'shape':
                # Find shape_annotated visualization
                viz = next((v for v in analysis_data['visualizations'] 
                           if v['visualization_type'] == 'shape_annotated'), None)
                if viz and viz.get('full_path'):
                    result_data['annotated_image'] = QPixmap(viz['full_path'])
                result_data['shape_class'] = result['predicted_class']
                result_data['class_id'] = result['metadata'].get('class_id', 0)
            
            elif model_type == 'audio':
                # Find audio visualizations
                waveform_viz = next((v for v in analysis_data['visualizations'] 
                                    if v['visualization_type'] == 'audio_waveform'), None)
                spectrogram_viz = next((v for v in analysis_data['visualizations'] 
                                       if v['visualization_type'] == 'audio_spectrogram'), None)
                
                if waveform_viz and waveform_viz.get('full_path'):
                    result_data['waveform_image'] = QPixmap(waveform_viz['full_path'])
                if spectrogram_viz and spectrogram_viz.get('full_path'):
                    result_data['spectrogram_image'] = QPixmap(spectrogram_viz['full_path'])
                
                result_data['ripeness_class'] = result['predicted_class']
                result_data['knock_count'] = result['metadata'].get('knock_count', 0)
            
            results[model_type] = result_data
        
        # Store for reference
        self.current_data = analysis_data
        self.current_report_id = report_id  # Store the report ID for display
        self.input_data = input_data
        
        # Display the report
        self.generate_report_with_rgb_and_multispectral(input_data, results)
        
        return True