dudong-v2/ui/components/report_generator.py

"""
Report Generator

Orchestrates report generation by combining section builders and visualizations.
"""

from typing import Dict, Optional
from datetime import datetime

from PyQt5.QtWidgets import QVBoxLayout
from PyQt5.QtGui import QImage

from ui.components.report_sections import (
    create_report_info_section,
    create_empty_results_section,
    create_model_results_section,
    create_analysis_results_section,
    create_input_data_section,
    create_input_data_with_gradcam,
    create_visualizations_section,
    get_local_datetime_string
)


def generate_basic_report(
    layout: QVBoxLayout,
    input_data: Dict[str, str]
) -> Dict:
    """
    Generate a basic analysis report without model results.
    
    Args:
        layout: QVBoxLayout to add sections to
        input_data: Dictionary with input file paths
        
    Returns:
        Dictionary with current grade and description
    """
    # Report Info Section
    info_group = create_report_info_section()
    layout.addWidget(info_group)
    
    # Analysis Results Section (Empty)
    results_group = create_empty_results_section()
    layout.addWidget(results_group)
    
    # Input Data Section
    data_group = create_input_data_section(input_data)
    layout.addWidget(data_group)
    
    return {
        'grade': 'B',
        'description': 'Basic report generated without model analysis'
    }


def generate_model_report(
    layout: QVBoxLayout,
    input_data: Dict[str, str],
    gradcam_image: QImage,
    predicted_class: str,
    confidence: float,
    probabilities: Dict[str, float]
) -> Dict:
    """
    Generate analysis report with actual multispectral model results.
    
    Args:
        layout: QVBoxLayout to add sections to
        input_data: Dictionary with input file paths
        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
        
    Returns:
        Dictionary with current grade and description
    """
    # Report Info Section
    info_group = create_report_info_section()
    layout.addWidget(info_group)
    
    # Analysis Results Section (WITH REAL MODEL DATA)
    results_group = create_model_results_section(
        predicted_class,
        confidence,
        probabilities
    )
    layout.addWidget(results_group)
    
    # Input Data Section (WITH GRADCAM VISUALIZATION)
    data_group = create_input_data_with_gradcam(input_data, gradcam_image)
    layout.addWidget(data_group)
    
    # Determine grade
    grade_map = {
        'Immature': 'C',
        'Mature': 'A',
        'Overmature': 'B'
    }
    grade = grade_map.get(predicted_class, 'B')
    
    return {
        'grade': grade,
        'description': f'Model prediction: {predicted_class} ({confidence*100:.1f}% confidence)'
    }


def generate_comprehensive_report(
    layout: QVBoxLayout,
    input_data: Dict[str, str],
    results: Dict,
    report_id: Optional[str] = None
) -> Dict:
    """
    Generate comprehensive report with RGB models and multispectral analysis.
    
    Args:
        layout: QVBoxLayout to add sections to
        input_data: Dictionary with input file paths
        results: Dictionary with processing results from all models
                 Keys: 'defect', 'locule', 'maturity', 'shape', 'audio'
        report_id: Optional report ID to display
        
    Returns:
        Dictionary with current grade and description
    """
    # Report Info Section
    info_group = create_report_info_section(report_id)
    layout.addWidget(info_group)
    
    # Combined Analysis Results Section
    results_group = create_analysis_results_section(results)
    layout.addWidget(results_group)
    
    # Extract grade information from results group
    # (Grade is calculated within create_analysis_results_section)
    locule_count = 0
    has_defects = False
    shape_class = None
    maturity_class = None
    
    if 'locule' in results and not results['locule'].get('error'):
        locule_count = results['locule'].get('locule_count', 0)
    
    if 'defect' in results and not results['defect'].get('error'):
        primary_class = results['defect'].get('primary_class', 'Unknown')
        has_defects = (primary_class != "No Defects")
    
    if 'shape' in results and not results['shape'].get('error'):
        shape_class = results['shape'].get('shape_class', 'Unknown')
    
    if 'maturity' in results and not results['maturity'].get('error'):
        maturity_class = results['maturity'].get('class_name', 'Unknown')
    
    # Import here to avoid circular import
    from utils.grade_calculator import calculate_durian_grade
    grade, grade_description = calculate_durian_grade(locule_count, has_defects, shape_class, maturity_class)
    
    # Input Data with Visualizations Section
    data_group = create_visualizations_section(input_data, results)
    layout.addWidget(data_group)
    
    return {
        'grade': grade,
        'description': grade_description
    }


def extract_report_content(
    report_id: str,
    grade: str,
    grade_description: str,
    results: Optional[Dict] = None
) -> Dict:
    """
    Extract report content for export/print.
    
    Args:
        report_id: Report ID
        grade: Overall grade letter (A, B, or C)
        grade_description: Grade description text
        results: Optional results dictionary from models
        
    Returns:
        Dictionary with all report text content
    """
    content = {
        'report_id': report_id or f"DUR-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
        'generated': get_local_datetime_string(),
        'grade': grade,
        'grade_description': grade_description,
        'results': {}
    }
    
    if results:
        # Locule analysis
        if 'locule' in results and not results['locule'].get('error'):
            content['results']['locule_count'] = results['locule'].get('locule_count', 0)
        
        # Defect analysis
        if 'defect' in results and not results['defect'].get('error'):
            content['results']['defect_status'] = results['defect'].get('primary_class', 'Unknown')
            content['results']['total_detections'] = results['defect'].get('total_detections', 0)
        
        # Shape analysis
        if 'shape' in results and not results['shape'].get('error'):
            content['results']['shape_class'] = results['shape'].get('shape_class', 'Unknown')
            content['results']['shape_confidence'] = results['shape'].get('confidence', 0)
        
        # Maturity analysis
        if 'maturity' in results and not results['maturity'].get('error'):
            content['results']['maturity_class'] = results['maturity'].get('class_name', 'Unknown')
            content['results']['maturity_confidence'] = results['maturity'].get('confidence', 0)
        
        # Audio ripeness analysis
        if 'audio' in results and not results['audio'].get('error'):
            content['results']['ripeness_class'] = results['audio'].get('ripeness_class', 'Unknown')
            content['results']['ripeness_confidence'] = results['audio'].get('confidence', 0)
            content['results']['knock_count'] = results['audio'].get('knock_count', 0)
    
    return content