dudong-v2/ui/components/report_sections.py

"""
Report Section Builders

Functions for building individual report sections (info, results, visualizations, etc).
Each function returns a QGroupBox or QWidget ready to be added to the report layout.
"""

from datetime import datetime
from typing import Dict, Optional

from PyQt5.QtWidgets import QGroupBox, QVBoxLayout, QGridLayout, QLabel, QWidget, QHBoxLayout
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QFont

from resources.styles import GROUP_BOX_STYLE
from utils.grade_calculator import (
    calculate_durian_grade,
    get_ripeness_color,
    get_maturity_color,
    get_grade_color
)
from ui.components.report_visualizations import create_clickable_image_widget


def get_local_datetime_string() -> str:
    """Get current datetime string formatted in the computer's local timezone."""
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def create_report_info_section(report_id: Optional[str] = None) -> QGroupBox:
    """
    Create report information section with local timezone.
    
    Args:
        report_id: Report ID to display (or None to generate one)
        
    Returns:
        QGroupBox containing report information
    """
    group = QGroupBox("Report Information")
    group.setStyleSheet(GROUP_BOX_STYLE)
    
    layout = QGridLayout()
    
    # Report ID (use provided ID or generate one)
    layout.addWidget(QLabel("<b>Report ID:</b>"), 0, 0)
    if report_id:
        report_id_label = QLabel(report_id)
    else:
        report_id_label = QLabel(f"DUR-{datetime.now().strftime('%Y%m%d-%H%M%S')}")
    layout.addWidget(report_id_label, 0, 1)
    
    # Date/Time in local timezone
    layout.addWidget(QLabel("<b>Generated:</b>"), 0, 2)
    layout.addWidget(QLabel(get_local_datetime_string()), 0, 3)
    
    group.setLayout(layout)
    return group


def create_empty_results_section() -> QGroupBox:
    """
    Create analysis results section - shows message when no data available.
    
    Returns:
        QGroupBox with placeholder message
    """
    group = QGroupBox("Analysis Results")
    group.setStyleSheet(GROUP_BOX_STYLE + """
        QGroupBox {
            background-color: #ecf0f1;
            font-weight: bold;
            font-size: 14px;
        }
    """)
    
    layout = QVBoxLayout()
    
    # Display message when no analysis data is available
    message_label = QLabel("Not enough information")
    message_label.setFont(QFont("Arial", 14))
    message_label.setAlignment(Qt.AlignCenter)
    message_label.setStyleSheet("color: #7f8c8d; padding: 20px;")
    layout.addWidget(message_label)
    
    group.setLayout(layout)
    return group


def create_model_results_section(
    predicted_class: str,
    confidence: float,
    probabilities: Dict[str, float]
) -> QGroupBox:
    """
    Create analysis results section with actual model predictions.
    
    Args:
        predicted_class: Predicted maturity class from model
        confidence: Model confidence (0-1 scale)
        probabilities: Dictionary of class probabilities
        
    Returns:
        QGroupBox containing model results
    """
    group = QGroupBox("Analysis Results")
    group.setStyleSheet(GROUP_BOX_STYLE + """
        QGroupBox {
            background-color: #ecf0f1;
            font-weight: bold;
            font-size: 14px;
        }
    """)
    
    layout = QVBoxLayout()
    
    # Determine overall grade based on maturity class
    grade_map = {
        'Immature': 'C',
        'Mature': 'A',
        'Overmature': 'B'
    }
    grade = grade_map.get(predicted_class, 'B')
    
    # Overall Grade
    grade_layout = QHBoxLayout()
    grade_label = QLabel("Overall Grade:")
    grade_label.setFont(QFont("Arial", 16, QFont.Bold))
    grade_layout.addWidget(grade_label)
    
    grade_value = QLabel(f"Class {grade}")
    grade_value.setFont(QFont("Arial", 24, QFont.Bold))
    grade_color = get_grade_color(grade)
    grade_value.setStyleSheet(f"color: {grade_color};")
    grade_layout.addWidget(grade_value)
    grade_layout.addStretch()
    layout.addLayout(grade_layout)
    
    # Results grid
    results_grid = QGridLayout()
    
    # Maturity (from model)
    results_grid.addWidget(QLabel("<b>Maturity Status:</b>"), 0, 0)
    maturity_label = QLabel(predicted_class)
    maturity_label.setStyleSheet(f"font-size: 14px; color: {get_maturity_color(predicted_class)};")
    results_grid.addWidget(maturity_label, 0, 1)
    
    # Confidence (from model)
    confidence_pct = confidence * 100 if confidence <= 1.0 else confidence
    results_grid.addWidget(QLabel(f"({confidence_pct:.1f}% confidence)"), 0, 2)
    
    # All class probabilities
    results_grid.addWidget(QLabel("<b>Class Probabilities:</b>"), 1, 0)
    prob_text = "  ".join([f"{k}: {v:.1%}" for k, v in probabilities.items()])
    prob_label = QLabel(prob_text)
    prob_label.setStyleSheet("font-size: 11px; color: #666;")
    results_grid.addWidget(prob_label, 1, 1, 1, 2)
    
    layout.addLayout(results_grid)
    
    group.setLayout(layout)
    return group


def create_analysis_results_section(results: Dict) -> QGroupBox:
    """
    Create combined analysis results from all models.
    
    Args:
        results: Dictionary with processing results from all models
                 Keys: 'defect', 'locule', 'maturity', 'shape', 'audio'
                 
    Returns:
        QGroupBox containing combined analysis results
    """
    group = QGroupBox("Combined Analysis Results")
    group.setStyleSheet(GROUP_BOX_STYLE + """
        QGroupBox {
            background-color: #ecf0f1;
            font-weight: bold;
            font-size: 14px;
        }
    """)
    
    layout = QVBoxLayout()
    
    # Extract data
    locule_count = 0
    has_defects = False
    shape_class = None
    maturity_class = None
    
    # Locule analysis from top view
    if 'locule' in results:
        if results['locule'].get('error'):
            layout.addWidget(QLabel(f"<b>Locule Count (Top View):</b> <font color='#e74c3c'>Error: {results['locule'].get('error_msg', 'Unknown error')}</font>"))
        else:
            locule_count = results['locule'].get('locule_count', 0)
            layout.addWidget(QLabel(f"<b>Locule Count (Top View):</b> {locule_count} locules"))
    
    # Defect analysis from side view
    if 'defect' in results:
        if results['defect'].get('error'):
            layout.addWidget(QLabel(f"<b>Defect Status (Side View):</b> <font color='#e74c3c'>Error: {results['defect'].get('error_msg', 'Unknown error')}</font>"))
        else:
            primary_class = results['defect'].get('primary_class', 'Unknown')
            total_detections = results['defect'].get('total_detections', 0)
            has_defects = (primary_class != "No Defects")
            layout.addWidget(QLabel(f"<b>Defect Status (Side View):</b> {primary_class} ({total_detections} detections)"))
    
    # Shape analysis from side view
    if 'shape' in results:
        if results['shape'].get('error'):
            layout.addWidget(QLabel(f"<b>Shape Classification (Side View):</b> <font color='#e74c3c'>Error: {results['shape'].get('error_msg', 'Unknown error')}</font>"))
        else:
            shape_class = results['shape'].get('shape_class', 'Unknown')
            confidence = results['shape'].get('confidence', 0)
            layout.addWidget(QLabel(f"<b>Shape Classification (Side View):</b> {shape_class} ({confidence*100:.1f}%)"))
    
    # Maturity analysis from multispectral (PART OF QUALITY GRADE)
    if 'maturity' in results:
        if results['maturity'].get('error'):
            layout.addWidget(QLabel(f"<b>Maturity (Multispectral):</b> <font color='#e74c3c'>Error: {results['maturity'].get('error_msg', 'Unknown error')}</font>"))
        else:
            maturity_class = results['maturity'].get('class_name', 'Unknown')
            confidence = results['maturity'].get('confidence', 0)
            layout.addWidget(QLabel(f"<b>Maturity (Multispectral):</b> {maturity_class} ({confidence*100:.1f}%)"))
    
    # Audio ripeness analysis (SEPARATE FROM QUALITY GRADE)
    ripeness_class = None
    if 'audio' in results:
        if results['audio'].get('error'):
            layout.addWidget(QLabel(f"<b>Ripeness Classification (Audio):</b> <font color='#e74c3c'>Error: {results['audio'].get('error_msg', 'Unknown error')}</font>"))
        else:
            ripeness_class = results['audio'].get('ripeness_class', 'Unknown')
            confidence = results['audio'].get('confidence', 0)
            knock_count = results['audio'].get('knock_count', 0)
            
            # Main result
            layout.addWidget(QLabel(f"<b>Ripeness Classification (Audio):</b> {ripeness_class} ({confidence*100:.1f}%)"))
            
            # Detailed analysis
            per_knock_preds = results['audio'].get('per_knock_predictions', [])
            if per_knock_preds:
                # Per-knock predictions
                per_knock_classes = [p['class'].capitalize() for p in per_knock_preds]
                per_knock_confidences = [p['confidence'] for p in per_knock_preds]
                layout.addWidget(QLabel(f"<i>Per-knock predictions: {', '.join(per_knock_classes)}</i>"))
                
                # Per-knock confidences
                per_knock_conf_str = ', '.join([f"{c*100:.1f}%" for c in per_knock_confidences])
                layout.addWidget(QLabel(f"<i>Per-knock confidence: {per_knock_conf_str}</i>"))
            
            # Average probabilities
            probabilities = results['audio'].get('probabilities', {})
            if probabilities:
                prob_str = ', '.join([f"{k}: {v*100:.1f}%" for k, v in probabilities.items()])
                layout.addWidget(QLabel(f"<i>Average probabilities: {prob_str}</i>"))
    
    # Calculate and display grade (includes maturity, excludes ripeness)
    grade, grade_description = calculate_durian_grade(locule_count, has_defects, shape_class, maturity_class)
    
    # Note: Ripeness is displayed separately and NOT included in quality grade
    
    grade_color = get_grade_color(grade)
    grade_widget = QLabel(f"<b>Overall Grade:</b> <font color='{grade_color}' size='5'><b>Class {grade}</b></font>")
    grade_widget.setStyleSheet(f"padding: 10px; background-color: white; border: 2px solid {grade_color};")
    layout.addWidget(grade_widget)
    
    layout.addWidget(QLabel(f"<i>{grade_description}</i>"))
    
    group.setLayout(layout)
    return group


def create_input_data_section(input_data: Dict[str, str]) -> QGroupBox:
    """
    Create section showing input data (only for provided inputs).
    
    Args:
        input_data: Dictionary with input file paths
        
    Returns:
        QGroupBox containing input data previews
    """
    group = QGroupBox("Input Data")
    group.setStyleSheet(GROUP_BOX_STYLE)
    
    layout = QVBoxLayout()
    
    has_content = False
    
    # Side View (Plain DSLR)
    if input_data.get('dslr_side'):
        dslr_side_widget = create_image_preview_widget("DSLR Side View", input_data['dslr_side'])
        layout.addWidget(dslr_side_widget)
        has_content = True
    
    # Top View (RGB DSLR)
    if input_data.get('dslr_top'):
        dslr_top_widget = create_image_preview_widget("DSLR Top View (RGB)", input_data['dslr_top'])
        layout.addWidget(dslr_top_widget)
        has_content = True
    
    # Thermal data (if provided)
    if input_data.get('thermal'):
        from ui.components.report_visualizations import create_thermal_widget
        thermal_widget = create_thermal_widget(input_data['thermal'])
        layout.addWidget(thermal_widget)
        has_content = True
    
    if not has_content:
        layout.addWidget(QLabel("No input data provided"))
    
    group.setLayout(layout)
    return group


def create_input_data_with_gradcam(input_data: Dict[str, str], gradcam_image) -> QGroupBox:
    """
    Create section showing input data with Grad-CAM visualization.
    
    Args:
        input_data: Dictionary with input file paths
        gradcam_image: QImage of Grad-CAM visualization
        
    Returns:
        QGroupBox containing input data and Grad-CAM
    """
    from ui.components.report_visualizations import create_gradcam_widget
    
    group = QGroupBox("Input Data & Analysis Visualization")
    group.setStyleSheet(GROUP_BOX_STYLE)
    
    layout = QVBoxLayout()
    
    # Grad-CAM visualization (from multispectral model)
    if gradcam_image:
        gradcam_widget = create_gradcam_widget(gradcam_image)
        layout.addWidget(gradcam_widget)
    
    # Side View (Plain DSLR)
    if input_data.get('dslr_side'):
        dslr_side_widget = create_image_preview_widget("DSLR Side View", input_data['dslr_side'])
        layout.addWidget(dslr_side_widget)
    
    # Top View (RGB DSLR)
    if input_data.get('dslr_top'):
        dslr_top_widget = create_image_preview_widget("DSLR Top View (RGB)", input_data['dslr_top'])
        layout.addWidget(dslr_top_widget)
    
    # Thermal data (if provided)
    if input_data.get('thermal'):
        from ui.components.report_visualizations import create_thermal_widget
        thermal_widget = create_thermal_widget(input_data['thermal'])
        layout.addWidget(thermal_widget)
    
    if not gradcam_image and not any([
        input_data.get('dslr_side'),
        input_data.get('dslr_top'),
        input_data.get('thermal')
    ]):
        layout.addWidget(QLabel("No visualizations to display"))
    
    group.setLayout(layout)
    return group


def create_visualizations_section(input_data: Dict[str, str], results: Dict) -> QGroupBox:
    """
    Create input data section with all model visualizations and clickable images.
    
    Args:
        input_data: Dictionary with input file paths
        results: Dictionary with processing results from all models
        
    Returns:
        QGroupBox containing all visualizations
    """
    group = QGroupBox("Analysis Visualizations")
    group.setStyleSheet(GROUP_BOX_STYLE)
    
    layout = QVBoxLayout()
    
    # Grad-CAM from multispectral
    if results.get('maturity') and results['maturity'].get('gradcam_image'):
        gradcam_img = results['maturity']['gradcam_image']
        gradcam_widget = create_clickable_image_widget(
            gradcam_img,
            "Grad-CAM Visualization (Maturity)",
            "Model attention heatmap overlaid on 860nm NIR band",
            max_width=500
        )
        layout.addWidget(gradcam_widget)
    
    # Locule analysis image
    if results.get('locule') and results['locule'].get('annotated_image'):
        locule_img = results['locule']['annotated_image']
        locule_widget = create_clickable_image_widget(
            locule_img,
            "Locule Analysis (Top View)",
            "Detected and counted locules with color coding",
            max_width=500
        )
        layout.addWidget(locule_widget)
    
    # Defect analysis image
    if results.get('defect') and results['defect'].get('annotated_image'):
        defect_img = results['defect']['annotated_image']
        defect_widget = create_clickable_image_widget(
            defect_img,
            "Defect Analysis (Side View)",
            f"{results['defect'].get('primary_class', 'Unknown')} - {results['defect'].get('total_detections', 0)} detections",
            max_width=500
        )
        layout.addWidget(defect_widget)
    
    # Shape analysis image
    if results.get('shape') and results['shape'].get('annotated_image'):
        shape_img = results['shape']['annotated_image']
        confidence = results['shape'].get('confidence', 0) * 100
        shape_widget = create_clickable_image_widget(
            shape_img,
            "Shape Classification (Side View)",
            f"{results['shape'].get('shape_class', 'Unknown')} ({confidence:.1f}% confidence)",
            max_width=500
        )
        layout.addWidget(shape_widget)
    
    # Audio ripeness - Waveform with knocks
    if results.get('audio') and results['audio'].get('waveform_image'):
        waveform_img = results['audio']['waveform_image']
        knock_count = results['audio'].get('knock_count', 0)
        waveform_widget = create_clickable_image_widget(
            waveform_img,
            f"Waveform with {knock_count} Detected Knocks",
            "Audio waveform with knock detection markers",
            max_width=600
        )
        layout.addWidget(waveform_widget)
    
    # Audio ripeness - Mel Spectrogram with knocks
    if results.get('audio') and results['audio'].get('spectrogram_image'):
        spectrogram_img = results['audio']['spectrogram_image']
        knock_count = results['audio'].get('knock_count', 0)
        audio_widget = create_clickable_image_widget(
            spectrogram_img,
            f"Mel Spectrogram (64 Coefficients) - {knock_count} Knocks",
            "Frequency-time representation of knock sounds",
            max_width=600
        )
        layout.addWidget(audio_widget)
    
    # Thermal if provided
    if input_data.get('thermal'):
        from ui.components.report_visualizations import create_thermal_widget
        thermal_widget = create_thermal_widget(input_data['thermal'])
        layout.addWidget(thermal_widget)
    
    group.setLayout(layout)
    return group


def create_image_preview_widget(title: str, file_path: str) -> QWidget:
    """
    Create image preview widget from file path.
    
    Args:
        title: Title for the image
        file_path: Path to the image file
        
    Returns:
        QWidget containing the image preview
    """
    from ui.components.report_visualizations import create_image_preview_widget as _create_image
    return _create_image(title, file_path)