dudong-v2/utils/quality_sample_data.py

"""
Quality Sample Data Generator

Utility for generating sample data and images for quality tab demonstration.
"""

import numpy as np
from PyQt5.QtGui import QPixmap, QImage, QColor, QPainter, QPen, QBrush
from PyQt5.QtCore import Qt, QPoint, QRect
import random
import math


class QualitySampleDataGenerator:
    """Generator for sample quality assessment data."""

    def __init__(self):
        self.fruit_colors = [
            QColor("#8B4513"),  # Brown/orange
            QColor("#CD853F"),  # Peru
            QColor("#DEB887"),  # Burlywood
            QColor("#F4A460"),  # Sandy brown
        ]

        self.defect_colors = [
            QColor("#F39C12"),  # Orange for mechanical damage
            QColor("#E67E22"),  # Carrot for surface blemish
            QColor("#E74C3C"),  # Red for severe defects
        ]

    def generate_sample_fruit_top_view(self, width=300, height=250):
        """Generate a sample top view fruit image with defects."""
        image = QImage(width, height, QImage.Format_RGB32)
        image.fill(QColor("#2C3E50"))

        painter = QPainter(image)
        painter.setRenderHint(QPainter.Antialiasing)

        # Draw main fruit
        fruit_radius = min(width, height) // 4
        center_x = width // 2
        center_y = height // 2

        # Fruit body with gradient
        fruit_color = random.choice(self.fruit_colors)
        painter.setBrush(QBrush(fruit_color))
        painter.setPen(QPen(QColor("#654321"), 2))
        painter.drawEllipse(center_x - fruit_radius, center_y - fruit_radius,
                          fruit_radius * 2, fruit_radius * 2)

        # Add some texture/pattern to fruit
        self._add_fruit_texture(painter, center_x, center_y, fruit_radius)

        # Add defect markers
        defects = self._generate_sample_defects()
        for defect in defects:
            self._draw_defect_on_fruit(painter, defect, center_x, center_y, fruit_radius)

        painter.end()
        return image, defects

    def generate_sample_fruit_side_view(self, width=300, height=200):
        """Generate a sample side view fruit image with shape outline."""
        image = QImage(width, height, QImage.Format_RGB32)
        image.fill(QColor("#2C3E50"))

        painter = QPainter(image)
        painter.setRenderHint(QPainter.Antialiasing)

        # Draw main fruit (elliptical for side view)
        fruit_width = width * 3 // 5
        fruit_height = fruit_width * 3 // 4
        center_x = width // 2
        center_y = height // 2

        # Fruit body
        fruit_color = random.choice(self.fruit_colors)
        painter.setBrush(QBrush(fruit_color))
        painter.setPen(QPen(QColor("#654321"), 2))
        painter.drawEllipse(center_x - fruit_width // 2, center_y - fruit_height // 2,
                          fruit_width, fruit_height)

        # Draw shape outline (dashed)
        painter.setBrush(Qt.NoBrush)
        painter.setPen(QPen(QColor("#27AE60"), 2, Qt.DashLine))
        painter.drawEllipse(center_x - fruit_width // 2, center_y - fruit_height // 2,
                          fruit_width, fruit_height)

        # Draw symmetry indicators
        self._draw_symmetry_indicators(painter, center_x, center_y, fruit_width)

        painter.end()

        # Generate shape analysis data
        symmetry = random.uniform(85, 95)
        aspect_ratio = fruit_width / fruit_height

        return image, symmetry, aspect_ratio

    def generate_thermal_gradient(self, width=300, height=200, base_temp=28.5):
        """Generate thermal gradient visualization."""
        image = QImage(width, height, QImage.Format_RGB32)
        image.fill(QColor("#000000"))

        painter = QPainter(image)
        painter.setRenderHint(QPainter.Antialiasing)

        # Create thermal gradient background
        center_x = width // 2
        center_y = height // 2

        # Main thermal area (elliptical)
        thermal_width = width * 3 // 4
        thermal_height = height * 2 // 3

        # Create radial gradient for thermal effect
        radial_gradient = QRadialGradient(center_x, center_y, max(thermal_width, thermal_height) // 2)

        # Center is hottest (red), edges cooler
        radial_gradient.setColorAt(0.0, QColor("#E74C3C"))  # Hot center
        radial_gradient.setColorAt(0.7, QColor("#F39C12"))  # Medium
        radial_gradient.setColorAt(1.0, QColor("#F1C40F"))  # Cool edges

        painter.setBrush(QBrush(radial_gradient))
        painter.setPen(Qt.NoPen)
        painter.drawEllipse(center_x - thermal_width // 2, center_y - thermal_height // 2,
                          thermal_width, thermal_height)

        # Add temperature scale overlay
        self._add_temperature_scale(painter, width, height, base_temp)

        painter.end()
        return image, base_temp

    def generate_quality_grading_data(self):
        """Generate sample quality grading data."""
        # Simulate quality metrics
        metrics = {
            "size_weight": random.uniform(85, 95),
            "shape_regularity": random.uniform(87, 94),
            "surface_quality": random.uniform(65, 80),
            "color_uniformity": random.uniform(78, 88),
            "firmness_thermal": 0  # Not available (thermal offline)
        }

        # Calculate overall score
        available_metrics = [v for k, v in metrics.items() if v > 0]
        overall_score = sum(available_metrics) / len(available_metrics)

        # Determine grade
        if overall_score >= 90:
            grade = "A"
        elif overall_score >= 75:
            grade = "B"
        else:
            grade = "C"

        # Generate additional metrics
        additional_metrics = {
            "estimated_weight": round(random.uniform(150, 220), 1),
            "diameter": random.randint(65, 80),
            "aspect_ratio": round(random.uniform(0.8, 0.9), 2),
            "surface_defect_coverage": round(random.uniform(2, 8), 1),
            "processing_time": round(random.uniform(2.0, 3.5), 1)
        }

        return {
            "grade": grade,
            "overall_score": round(overall_score, 1),
            "metrics": metrics,
            "additional_metrics": additional_metrics,
            "model_version": "QualityNet v2.8"
        }

    def _generate_sample_defects(self):
        """Generate sample defect data."""
        defects = []

        # Mechanical damage
        defects.append({
            'type': 'Mechanical Damage',
            'x': random.randint(-20, 20),
            'y': random.randint(-20, 20),
            'size': random.randint(6, 12),
            'confidence': round(random.uniform(80, 95), 1),
            'color': '#F39C12',
            'location': 'Top-Left',
            'size_mm2': round(random.uniform(6, 12), 1)
        })

        # Surface blemish (sometimes)
        if random.random() > 0.5:
            defects.append({
                'type': 'Surface Blemish',
                'x': random.randint(-15, 25),
                'y': random.randint(-10, 25),
                'size': random.randint(4, 8),
                'confidence': round(random.uniform(70, 85), 1),
                'color': '#E67E22',
                'location': 'Side',
                'size_mm2': round(random.uniform(3, 7), 1)
            })

        return defects

    def _add_fruit_texture(self, painter, center_x, center_y, radius):
        """Add texture pattern to fruit."""
        # Add some darker spots for texture
        for _ in range(5):
            spot_x = center_x + random.randint(-radius // 2, radius // 2)
            spot_y = center_y + random.randint(-radius // 2, radius // 2)
            spot_radius = random.randint(3, 8)

            # Only draw if within fruit bounds
            distance = math.sqrt((spot_x - center_x) ** 2 + (spot_y - center_y) ** 2)
            if distance + spot_radius <= radius:
                painter.setBrush(QBrush(QColor("#654321")))
                painter.setPen(Qt.NoPen)
                painter.drawEllipse(spot_x - spot_radius, spot_y - spot_radius,
                                  spot_radius * 2, spot_radius * 2)

    def _draw_defect_on_fruit(self, painter, defect, center_x, center_y, fruit_radius):
        """Draw a defect marker on the fruit."""
        # Calculate position relative to fruit center
        marker_x = center_x + (defect['x'] * fruit_radius // 50)
        marker_y = center_y + (defect['y'] * fruit_radius // 50)

        # Draw marker circle
        marker_radius = max(6, defect['size'])
        painter.setBrush(QBrush(QColor(defect['color'])))
        painter.setPen(QPen(QColor('#D35400'), 2))
        painter.drawEllipse(marker_x - marker_radius, marker_y - marker_radius,
                          marker_radius * 2, marker_radius * 2)

        # Draw confidence text
        painter.setPen(QPen(QColor('white'), 1))
        painter.drawText(marker_x - 15, marker_y - marker_radius - 3,
                        f"{defect['confidence']}%")

    def _draw_symmetry_indicators(self, painter, center_x, center_y, fruit_width):
        """Draw symmetry indicator lines."""
        # Draw symmetry axis lines
        symmetry_y = center_y
        line_length = fruit_width // 6

        painter.setPen(QPen(QColor("#3498DB"), 2, Qt.SolidLine))

        # Left indicator
        left_x = center_x - fruit_width // 4
        painter.drawLine(left_x, symmetry_y - line_length, left_x, symmetry_y + line_length)

        # Right indicator
        right_x = center_x + fruit_width // 4
        painter.drawLine(right_x, symmetry_y - line_length, right_x, symmetry_y + line_length)

    def _add_temperature_scale(self, painter, width, height, base_temp):
        """Add temperature scale overlay."""
        # Temperature scale at bottom
        scale_width = width * 3 // 4
        scale_height = 15
        scale_x = (width - scale_width) // 2
        scale_y = height - 30

        # Draw scale background
        painter.setBrush(QBrush(QColor("#34495E")))
        painter.setPen(Qt.NoPen)
        painter.drawRect(scale_x, scale_y, scale_width, scale_height)

        # Draw temperature gradient in scale
        temp_gradient = QBrush()
        temp_gradient = QLinearGradient(scale_x, 0, scale_x + scale_width, 0)

        temp_gradient.setColorAt(0.0, QColor("#3498DB"))  # Cool (22°C)
        temp_gradient.setColorAt(0.25, QColor("#27AE60"))
        temp_gradient.setColorAt(0.5, QColor("#F1C40F"))
        temp_gradient.setColorAt(0.75, QColor("#E67E22"))
        temp_gradient.setColorAt(1.0, QColor("#E74C3C"))   # Hot (32°C)

        painter.setBrush(temp_gradient)
        painter.drawRect(scale_x, scale_y, scale_width, scale_height)

        # Draw scale border
        painter.setPen(QPen(QColor("#ECF0F1"), 1))
        painter.drawRect(scale_x, scale_y, scale_width, scale_height)

        # Draw temperature labels
        painter.setPen(QPen(QColor("white"), 1))
        min_temp = "22°C"
        max_temp = "32°C"
        current_temp = f"{base_temp}°C"

        # Min temperature
        painter.drawText(scale_x + 5, scale_y - 2, min_temp)

        # Current temperature (centered)
        current_rect = QRect(scale_x + scale_width // 2 - 20, scale_y - 18, 40, 15)
        painter.drawText(current_rect, Qt.AlignCenter, current_temp)

        # Max temperature
        max_temp_rect = QRect(scale_x + scale_width - 35, scale_y - 2, 30, 12)
        painter.drawText(max_temp_rect, Qt.AlignRight, max_temp)


# Global instance for easy access
sample_data_generator = QualitySampleDataGenerator()


def get_sample_fruit_top_view():
    """Get a sample top view image and defects data."""
    return sample_data_generator.generate_sample_fruit_top_view()


def get_sample_fruit_side_view():
    """Get a sample side view image and analysis data."""
    return sample_data_generator.generate_sample_fruit_side_view()


def get_thermal_gradient():
    """Get a thermal gradient visualization."""
    return sample_data_generator.generate_thermal_gradient()


def get_quality_grading_data():
    """Get sample quality grading data."""
    return sample_data_generator.generate_quality_grading_data()


def generate_defect_detection_results():
    """Generate sample defect detection results."""
    defects = sample_data_generator._generate_sample_defects()

    # Convert to format expected by panels
    formatted_defects = []

    for defect in defects:
        formatted_defects.append({
            'type': defect['type'],
            'location': defect['location'],
            'size': f"{defect['size_mm2']}mm²",
            'confidence': defect['confidence'],
            'color': defect['color'],
            'category': 'warning'
        })

    # Add shape analysis
    formatted_defects.append({
        'type': 'Shape Analysis',
        'result': 'Regular',
        'symmetry': '91.2%',
        'confidence': 94.1,
        'color': '#27ae60',
        'category': 'success'
    })

    # Add locule count
    formatted_defects.append({
        'type': 'Locule Count',
        'count': '4',
        'confidence': 94.5,
        'color': '#3498db',
        'category': 'info'
    })

    return formatted_defects