dudong-v2/ui/components/report_visualizations.py

"""
Report Visualization Components

Reusable visualization widgets for creating thermal heatmaps, audio spectrograms,
and image previews for report generation.
"""

from pathlib import Path
from typing import Optional

from PyQt5.QtWidgets import QWidget, QFrame, QVBoxLayout, QLabel
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QFont, QPixmap, QImage
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.colors import LinearSegmentedColormap

try:
    import tensorflow as tf
except:
    tf = None

try:
    from PIL import Image as PILImage
    HAS_PIL = True
except ImportError:
    HAS_PIL = False

from ui.components.visualization_widgets import ClickableImageWidget
from ui.dialogs.image_preview_dialog import ImagePreviewDialog


def create_gradcam_widget(gradcam_image: QImage) -> QWidget:
    """
    Create Grad-CAM visualization preview widget.
    
    Args:
        gradcam_image: QImage of the Grad-CAM visualization
        
    Returns:
        QWidget containing the Grad-CAM visualization
    """
    widget = QFrame()
    widget.setFrameStyle(QFrame.Box)
    widget.setStyleSheet("background-color: white; border: 1px solid #bdc3c7;")
    
    layout = QVBoxLayout(widget)
    
    # Title
    title_label = QLabel("<b>Grad-CAM Visualization</b>")
    title_label.setFont(QFont("Arial", 12))
    layout.addWidget(title_label)
    
    # Description
    desc_label = QLabel("Model attention heatmap overlaid on 860nm NIR band")
    desc_label.setStyleSheet("color: #7f8c8d; font-size: 10px;")
    layout.addWidget(desc_label)
    
    # Image preview
    try:
        if not gradcam_image.isNull():
            # Scale to reasonable size
            scaled_image = gradcam_image.scaledToWidth(400, Qt.SmoothTransformation)
            pixmap = QPixmap.fromImage(scaled_image)
            image_label = QLabel()
            image_label.setPixmap(pixmap)
            image_label.setAlignment(Qt.AlignCenter)
            layout.addWidget(image_label)
        else:
            layout.addWidget(QLabel("Grad-CAM image is empty"))
    except Exception as e:
        layout.addWidget(QLabel(f"Error displaying Grad-CAM: {e}"))
    
    return widget


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
    """
    widget = QFrame()
    widget.setFrameStyle(QFrame.Box)
    widget.setStyleSheet("background-color: white; border: 1px solid #bdc3c7;")
    
    layout = QVBoxLayout(widget)
    
    # Title
    title_label = QLabel(f"<b>{title}</b>")
    title_label.setFont(QFont("Arial", 12))
    layout.addWidget(title_label)
    
    # File path
    path_label = QLabel(f"File: {Path(file_path).name}")
    path_label.setStyleSheet("color: #7f8c8d; font-size: 10px;")
    layout.addWidget(path_label)
    
    # Image preview
    try:
        pixmap = QPixmap(file_path)
        if not pixmap.isNull():
            # Scale to reasonable size
            scaled_pixmap = pixmap.scaled(400, 300, Qt.KeepAspectRatio, Qt.SmoothTransformation)
            image_label = QLabel()
            image_label.setPixmap(scaled_pixmap)
            image_label.setAlignment(Qt.AlignCenter)
            layout.addWidget(image_label)
        else:
            layout.addWidget(QLabel("Could not load image"))
    except Exception as e:
        layout.addWidget(QLabel(f"Error loading image: {e}"))
    
    return widget


def create_thermal_widget(csv_path: str) -> QWidget:
    """
    Create thermal CSV data preview widget with heatmap visualization.
    
    Args:
        csv_path: Path to thermal CSV file
        
    Returns:
        QWidget containing the thermal heatmap visualization
    """
    widget = QFrame()
    widget.setFrameStyle(QFrame.Box)
    widget.setStyleSheet("background-color: white; border: 1px solid #bdc3c7;")
    
    layout = QVBoxLayout(widget)
    
    # Title
    title_label = QLabel("<b>Thermal Image Analysis</b>")
    title_label.setFont(QFont("Arial", 12))
    layout.addWidget(title_label)
    
    # File path
    path_label = QLabel(f"File: {Path(csv_path).name}")
    path_label.setStyleSheet("color: #7f8c8d; font-size: 10px;")
    layout.addWidget(path_label)
    
    # Load and visualize thermal data
    try:
        # Load CSV file - FLIR Analyzer exports temperature data as comma-separated values
        data = []
        with open(csv_path, 'r') as f:
            for line in f:
                try:
                    row = [float(x.strip()) for x in line.strip().split(',')]
                    if row:  # Skip empty rows
                        data.append(row)
                except ValueError:
                    # Skip rows that can't be converted to float
                    continue
        
        if not data:
            layout.addWidget(QLabel("No valid thermal data found in CSV"))
            return widget
        
        # Convert to numpy array
        thermal_data = np.array(data)
        
        # Calculate statistics
        temp_min = thermal_data.min()
        temp_max = thermal_data.max()
        temp_mean = thermal_data.mean()
        temp_std = thermal_data.std()
        
        # Create heatmap visualization
        fig, ax = plt.subplots(figsize=(8, 6), dpi=100)
        
        # Create FLIR-style colormap: dark blue (cold) -> purple -> red -> yellow (hot)
        colors_flir = ['#000040', '#0000ff', '#0080ff', '#00ffff', 
                      '#00ff00', '#ffff00', '#ff8000', '#ff0000']
        cmap_flir = LinearSegmentedColormap.from_list('flir_style', colors_flir, N=256)
        
        # Create the heatmap
        im = ax.imshow(thermal_data, cmap=cmap_flir, origin='upper', aspect='auto')
        
        # Add colorbar
        cbar = plt.colorbar(im, ax=ax, label='Temperature (°C)')
        
        # Labels and title
        ax.set_xlabel('X Pixels')
        ax.set_ylabel('Y Pixels')
        ax.set_title('Thermal Camera Heatmap', fontsize=12, fontweight='bold')
        
        # Adjust layout
        plt.tight_layout()
        
        # Convert to QImage
        canvas = FigureCanvas(fig)
        canvas.draw()
        
        width_px, height_px = fig.get_size_inches() * fig.get_dpi()
        width_px, height_px = int(width_px), int(height_px)
        
        img = QImage(canvas.buffer_rgba(), width_px, height_px, QImage.Format_ARGB32)
        img = img.rgbSwapped()
        
        # Display image
        image_label = QLabel()
        pixmap = QPixmap.fromImage(img)
        # Scale to fit reasonable size
        if pixmap.width() > 600:
            scaled_pixmap = pixmap.scaledToWidth(600, Qt.SmoothTransformation)
        else:
            scaled_pixmap = pixmap
        image_label.setPixmap(scaled_pixmap)
        image_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(image_label)
        
        plt.close(fig)
        
        # Add statistics below the heatmap
        stats_label = QLabel(
            f"<b>Temperature Statistics:</b><br>"
            f"Min: {temp_min:.2f}°C | Max: {temp_max:.2f}°C | "
            f"Mean: {temp_mean:.2f}°C | Std Dev: {temp_std:.2f}°C<br>"
            f"<i>Image dimensions: {thermal_data.shape[0]} × {thermal_data.shape[1]} pixels</i>"
        )
        stats_label.setStyleSheet("color: #555; font-size: 11px; padding: 5px;")
        stats_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(stats_label)
        
    except Exception as e:
        layout.addWidget(QLabel(f"Error processing thermal data: {str(e)}"))
    
    return widget


def create_audio_spectrogram_widget(audio_path: str) -> QWidget:
    """
    Create audio spectrogram widget.
    
    Args:
        audio_path: Path to audio WAV file
        
    Returns:
        QWidget containing the audio spectrogram
    """
    widget = QFrame()
    widget.setFrameStyle(QFrame.Box)
    widget.setStyleSheet("background-color: white; border: 1px solid #bdc3c7;")
    
    layout = QVBoxLayout(widget)
    
    # Title
    title_label = QLabel("<b>Audio Spectrogram</b>")
    title_label.setFont(QFont("Arial", 12))
    layout.addWidget(title_label)
    
    # File path
    path_label = QLabel(f"File: {Path(audio_path).name}")
    path_label.setStyleSheet("color: #7f8c8d; font-size: 10px;")
    layout.addWidget(path_label)
    
    # Generate spectrogram
    try:
        if tf is not None:
            # Load and process audio
            x = tf.io.read_file(str(audio_path))
            x, sample_rate = tf.audio.decode_wav(x, desired_channels=1, desired_samples=16000)
            x = tf.squeeze(x, axis=-1)
            waveform = x
            
            # Generate spectrogram
            spectrogram = tf.signal.stft(waveform, frame_length=255, frame_step=128)
            spectrogram = tf.abs(spectrogram)
            spectrogram = spectrogram[..., tf.newaxis]
            
            # Plot spectrogram
            fig, ax = plt.subplots(1, 1, figsize=(8, 3))
            
            if len(spectrogram.shape) > 2:
                spectrogram = np.squeeze(spectrogram, axis=-1)
            log_spec = np.log(spectrogram.T + np.finfo(float).eps)
            height = log_spec.shape[0]
            width = log_spec.shape[1]
            X = np.linspace(0, np.size(spectrogram), num=width, dtype=int)
            Y = range(height)
            ax.pcolormesh(X, Y, log_spec)
            ax.set_ylabel('Frequency')
            ax.set_xlabel('Time')
            ax.set_title('Audio Spectrogram')
            
            # Convert to QPixmap
            canvas = FigureCanvas(fig)
            canvas.draw()
            
            width_px, height_px = fig.get_size_inches() * fig.get_dpi()
            width_px, height_px = int(width_px), int(height_px)
            
            img = QImage(canvas.buffer_rgba(), width_px, height_px, QImage.Format_ARGB32)
            img = img.rgbSwapped()
            pixmap = QPixmap(img)
            
            image_label = QLabel()
            image_label.setPixmap(pixmap)
            image_label.setAlignment(Qt.AlignCenter)
            layout.addWidget(image_label)
            
            plt.close(fig)
        else:
            layout.addWidget(QLabel("TensorFlow not available - cannot generate spectrogram"))
    except Exception as e:
        layout.addWidget(QLabel(f"Error generating spectrogram: {e}"))
    
    return widget


def create_clickable_image_widget(
    image_data,
    title: str,
    description: str = "",
    max_width: int = 500
) -> QWidget:
    """
    Create a clickable image widget that opens preview dialog on click.
    
    Args:
        image_data: QPixmap or QImage to display
        title: Title for the visualization
        description: Optional description text
        max_width: Maximum width for display image
        
    Returns:
        QWidget containing the clickable image
    """
    # Convert QImage to QPixmap if needed
    if isinstance(image_data, QImage):
        pixmap = QPixmap.fromImage(image_data)
    else:
        pixmap = image_data
    
    widget = QFrame()
    widget.setFrameStyle(QFrame.Box)
    widget.setStyleSheet("background-color: white; border: 1px solid #bdc3c7;")
    
    layout = QVBoxLayout(widget)
    layout.setAlignment(Qt.AlignCenter)
    layout.setContentsMargins(10, 10, 10, 10)
    
    # Title
    title_label = QLabel(f"<b>{title}</b>")
    title_label.setFont(QFont("Arial", 12))
    layout.addWidget(title_label)
    
    # Description (if provided)
    if description:
        desc_label = QLabel(description)
        desc_label.setStyleSheet("color: #7f8c8d; font-size: 10px;")
        layout.addWidget(desc_label)
    
    # Image label (clickable)
    image_label = QLabel()
    image_label.setCursor(Qt.PointingHandCursor)
    
    # Scale pixmap to fit within max_width and max_height, maintaining aspect ratio
    max_height = 600  # Set max height to prevent stretching
    scaled_pixmap = pixmap.scaledToWidth(max_width, Qt.SmoothTransformation) if pixmap.width() > max_width else pixmap
    
    # If height exceeds max_height after width scaling, scale by height instead
    if scaled_pixmap.height() > max_height:
        scaled_pixmap = pixmap.scaledToHeight(max_height, Qt.SmoothTransformation)
    
    image_label.setPixmap(scaled_pixmap)
    image_label.setAlignment(Qt.AlignCenter)
    
    # Store original pixmap for preview
    image_label.original_pixmap = pixmap
    image_label.preview_title = title
    
    # Connect click event
    def show_preview(event):
        dialog = ImagePreviewDialog(image_label.original_pixmap, title=image_label.preview_title, parent=widget)
        dialog.exec_()
    
    image_label.mousePressEvent = show_preview
    
    layout.addWidget(image_label, alignment=Qt.AlignCenter)
    
    return widget


def convert_pixmap_to_pil(pixmap: QPixmap) -> Optional['PILImage.Image']:
    """
    Convert QPixmap to PIL Image for PDF generation.
    
    Args:
        pixmap: QPixmap to convert
        
    Returns:
        PIL Image or None if conversion fails
    """
    if not HAS_PIL or pixmap.isNull():
        return None
    
    try:
        # Convert QPixmap to QImage
        qimage = pixmap.toImage()
        width = qimage.width()
        height = qimage.height()
        
        # Convert QImage to numpy array
        ptr = qimage.bits()
        ptr.setsize(qimage.byteCount())
        arr = np.array(ptr).reshape(height, width, 4)
        
        # Convert RGBA to RGB if needed
        if arr.shape[2] == 4:
            arr = arr[:, :, :3]
        
        # Create PIL Image
        pil_image = PILImage.fromarray(arr, 'RGB')
        return pil_image
    except Exception as e:
        print(f"Error converting QPixmap to PIL: {e}")
        return None