utils.py

"""
Utility functions for data preprocessing, visualization, and helper operations.
"""

import cv2
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image, ImageDraw, ImageFont
import pandas as pd
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Union
import json
import logging
from datetime import datetime
import os

logger = logging.getLogger(__name__)


def create_visualization_overlay(image: np.ndarray, 
                               damage_results: Dict,
                               show_confidence: bool = True) -> np.ndarray:
    """
    Create a visual overlay showing detected damage areas
    
    Args:
        image: Original image
        damage_results: Detection results from the main detector
        show_confidence: Whether to show confidence scores
        
    Returns:
        Image with damage visualization overlay
    """
    # Convert to PIL for easier drawing
    if isinstance(image, np.ndarray):
        pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    else:
        pil_image = image.copy()
    
    draw = ImageDraw.Draw(pil_image)
    
    # Color map for different damage types
    damage_colors = {
        'tear': (255, 0, 0),      # Red
        'crease': (255, 165, 0),   # Orange
        'stain': (128, 0, 128),    # Purple
        'fade': (255, 255, 0),     # Yellow
        'burn': (139, 69, 19),     # Brown
        'water_damage': (0, 100, 255)  # Blue
    }
    
    # Get image dimensions
    width, height = pil_image.size
    
    # Draw damage indicators
    y_offset = 10
    for damage_type, details in damage_results['damage_details'].items():
        if details['detected']:
            color = damage_colors.get(damage_type, (255, 255, 255))
            confidence = details['confidence']
            
            # Create text
            text = f"{damage_type.replace('_', ' ').title()}"
            if show_confidence:
                text += f": {confidence:.2f}"
            
            # Draw background rectangle
            text_bbox = draw.textbbox((10, y_offset), text)
            draw.rectangle(
                [text_bbox[0] - 5, text_bbox[1] - 2, 
                 text_bbox[2] + 5, text_bbox[3] + 2],
                fill=(0, 0, 0, 128)
            )
            
            # Draw text
            draw.text((10, y_offset), text, fill=color)
            y_offset += 25
    
    # Draw overall status
    overall = damage_results['overall_status']
    status_text = f"Status: {'DAMAGED' if overall['damaged'] else 'UNDAMAGED'}"
    status_text += f" ({overall['severity'].upper()})"
    
    # Position at bottom
    text_bbox = draw.textbbox((10, height - 30), status_text)
    draw.rectangle(
        [text_bbox[0] - 5, text_bbox[1] - 2, 
         text_bbox[2] + 5, text_bbox[3] + 2],
        fill=(0, 0, 0, 128)
    )
    
    status_color = (255, 0, 0) if overall['damaged'] else (0, 255, 0)
    draw.text((10, height - 30), status_text, fill=status_color)
    
    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)


def create_damage_heatmap(processed_images: Dict[str, np.ndarray],
                         damage_results: Dict) -> np.ndarray:
    """
    Create a heatmap visualization of damage probability
    
    Args:
        processed_images: Dictionary of processed image variants
        damage_results: Detection results
        
    Returns:
        Heatmap visualization
    """
    # Use edge information as base for heatmap
    edges = processed_images.get('edges', np.zeros((512, 512)))
    
    # Create heatmap based on damage confidence
    heatmap = np.zeros_like(edges, dtype=np.float32)
    
    # Add edge information
    heatmap += edges / 255.0 * 0.3
    
    # Add texture-based damage probability (simplified)
    gray = processed_images.get('gray', edges)
    
    # Calculate local variance as damage indicator
    kernel = np.ones((15, 15)) / 225
    local_mean = cv2.filter2D(gray.astype(np.float32), -1, kernel)
    local_variance = cv2.filter2D((gray.astype(np.float32) - local_mean)**2, -1, kernel)
    
    # Normalize and add to heatmap
    variance_norm = local_variance / np.max(local_variance) if np.max(local_variance) > 0 else local_variance
    heatmap += variance_norm * 0.4
    
    # Apply overall damage confidence
    overall_confidence = damage_results['overall_status']['overall_confidence']
    heatmap *= overall_confidence
    
    # Convert to color heatmap
    heatmap_norm = (heatmap * 255).astype(np.uint8)
    heatmap_color = cv2.applyColorMap(heatmap_norm, cv2.COLORMAP_JET)
    
    return heatmap_color


def generate_damage_report_pdf(results: Dict, output_path: str):
    """
    Generate a comprehensive PDF report (requires reportlab)
    
    Args:
        results: Detection results
        output_path: Path to save PDF report
    """
    try:
        from reportlab.lib.pagesizes import letter, A4
        from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
        from reportlab.lib.units import inch
        from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
        from reportlab.lib import colors
        from reportlab.lib.enums import TA_CENTER, TA_LEFT
        
        # Create PDF document
        doc = SimpleDocTemplate(output_path, pagesize=A4)
        styles = getSampleStyleSheet()
        story = []
        
        # Title
        title_style = ParagraphStyle(
            'CustomTitle',
            parent=styles['Heading1'],
            fontSize=24,
            textColor=colors.darkblue,
            alignment=TA_CENTER
        )
        story.append(Paragraph("Damaged Poster Detection Report", title_style))
        story.append(Spacer(1, 0.5*inch))
        
        # Metadata
        meta_data = [
            ['Image Path:', results['image_path']],
            ['Analysis Date:', results['timestamp']],
            ['Overall Status:', 'DAMAGED' if results['overall_status']['damaged'] else 'UNDAMAGED'],
            ['Severity:', results['overall_status']['severity'].upper()],
            ['Confidence:', f"{results['overall_status']['overall_confidence']:.3f}"]
        ]
        
        meta_table = Table(meta_data, colWidths=[2*inch, 4*inch])
        meta_table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (0, -1), colors.lightgrey),
            ('FONTWEIGHT', (0, 0), (0, -1), 'BOLD'),
            ('FONTSIZE', (0, 0), (-1, -1), 10),
            ('GRID', (0, 0), (-1, -1), 1, colors.black)
        ]))
        
        story.append(meta_table)
        story.append(Spacer(1, 0.3*inch))
        
        # Damage Details
        story.append(Paragraph("Damage Analysis Details", styles['Heading2']))
        story.append(Spacer(1, 0.2*inch))
        
        damage_data = [['Damage Type', 'Detected', 'Confidence', 'Status']]
        for damage_type, details in results['damage_details'].items():
            damage_data.append([
                damage_type.replace('_', ' ').title(),
                'Yes' if details['detected'] else 'No',
                f"{details['confidence']:.3f}",
                'ALERT' if details['detected'] else 'OK'
            ])
        
        damage_table = Table(damage_data, colWidths=[1.5*inch, 1*inch, 1*inch, 1*inch])
        damage_table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (-1, 0), colors.darkblue),
            ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
            ('FONTWEIGHT', (0, 0), (-1, 0), 'BOLD'),
            ('FONTSIZE', (0, 0), (-1, -1), 9),
            ('GRID', (0, 0), (-1, -1), 1, colors.black),
            ('ROWBACKGROUNDS', (0, 1), (-1, -1), [colors.white, colors.lightgrey])
        ]))
        
        story.append(damage_table)
        story.append(Spacer(1, 0.3*inch))
        
        # Recommendations
        story.append(Paragraph("Recommendations", styles['Heading2']))
        story.append(Spacer(1, 0.1*inch))
        
        for i, rec in enumerate(results['recommendations'], 1):
            story.append(Paragraph(f"{i}. {rec}", styles['Normal']))
            story.append(Spacer(1, 0.1*inch))
        
        # Build PDF
        doc.build(story)
        logger.info(f"PDF report generated: {output_path}")
        
    except ImportError:
        logger.warning("reportlab not installed. Cannot generate PDF report.")
        logger.info("Install with: pip install reportlab")


def analyze_batch_statistics(batch_results: List[Dict]) -> Dict:
    """
    Analyze statistics from batch processing results
    
    Args:
        batch_results: List of detection results
        
    Returns:
        Statistical analysis
    """
    stats = {
        'total_images': len(batch_results),
        'damaged_count': 0,
        'undamaged_count': 0,
        'error_count': 0,
        'damage_type_distribution': {},
        'severity_distribution': {'minimal': 0, 'moderate': 0, 'severe': 0},
        'average_confidence': 0.0,
        'processing_errors': []
    }
    
    valid_results = []
    
    for result in batch_results:
        if 'error' in result:
            stats['error_count'] += 1
            stats['processing_errors'].append({
                'image': result.get('image_path', 'unknown'),
                'error': result['error']
            })
            continue
        
        valid_results.append(result)
        overall = result['overall_status']
        
        if overall['damaged']:
            stats['damaged_count'] += 1
            
            # Count damage types
            for damage_type in overall['detected_damage_types']:
                stats['damage_type_distribution'][damage_type] = \
                    stats['damage_type_distribution'].get(damage_type, 0) + 1
            
            # Count severity
            severity = overall['severity']
            stats['severity_distribution'][severity] += 1
        else:
            stats['undamaged_count'] += 1
    
    # Calculate average confidence
    if valid_results:
        total_confidence = sum(r['overall_status']['overall_confidence'] for r in valid_results)
        stats['average_confidence'] = total_confidence / len(valid_results)
    
    return stats


def export_results_to_csv(batch_results: List[Dict], output_path: str):
    """
    Export batch results to CSV format
    
    Args:
        batch_results: List of detection results
        output_path: Path to save CSV file
    """
    rows = []
    
    for result in batch_results:
        if 'error' in result:
            rows.append({
                'image_path': result.get('image_path', 'unknown'),
                'status': 'ERROR',
                'error_message': result['error'],
                'overall_confidence': 0.0,
                'severity': 'unknown',
                'damage_types': '',
                'tear_confidence': 0.0,
                'crease_confidence': 0.0,
                'stain_confidence': 0.0,
                'fade_confidence': 0.0,
                'burn_confidence': 0.0,
                'water_damage_confidence': 0.0,
                'timestamp': result.get('timestamp', '')
            })
            continue
        
        overall = result['overall_status']
        damage_details = result['damage_details']
        
        rows.append({
            'image_path': result['image_path'],
            'status': 'DAMAGED' if overall['damaged'] else 'UNDAMAGED',
            'error_message': '',
            'overall_confidence': overall['overall_confidence'],
            'severity': overall['severity'],
            'damage_types': '; '.join(overall['detected_damage_types']),
            'tear_confidence': damage_details.get('tear', {}).get('confidence', 0.0),
            'crease_confidence': damage_details.get('crease', {}).get('confidence', 0.0),
            'stain_confidence': damage_details.get('stain', {}).get('confidence', 0.0),
            'fade_confidence': damage_details.get('fade', {}).get('confidence', 0.0),
            'burn_confidence': damage_details.get('burn', {}).get('confidence', 0.0),
            'water_damage_confidence': damage_details.get('water_damage', {}).get('confidence', 0.0),
            'timestamp': result['timestamp']
        })
    
    df = pd.DataFrame(rows)
    df.to_csv(output_path, index=False)
    logger.info(f"Results exported to CSV: {output_path}")


def create_summary_plots(batch_results: List[Dict], output_dir: str):
    """
    Create summary visualization plots
    
    Args:
        batch_results: List of detection results
        output_dir: Directory to save plots
    """
    os.makedirs(output_dir, exist_ok=True)
    
    # Analyze statistics
    stats = analyze_batch_statistics(batch_results)
    
    # 1. Overall status pie chart
    plt.figure(figsize=(10, 6))
    
    plt.subplot(1, 2, 1)
    labels = ['Undamaged', 'Damaged', 'Errors']
    sizes = [stats['undamaged_count'], stats['damaged_count'], stats['error_count']]
    colors = ['green', 'red', 'gray']
    
    plt.pie(sizes, labels=labels, colors=colors, autopct='%1.1f%%', startangle=90)
    plt.title('Overall Detection Results')
    
    # 2. Damage type distribution
    plt.subplot(1, 2, 2)
    if stats['damage_type_distribution']:
        damage_types = list(stats['damage_type_distribution'].keys())
        damage_counts = list(stats['damage_type_distribution'].values())
        
        plt.bar(damage_types, damage_counts)
        plt.title('Damage Type Distribution')
        plt.xlabel('Damage Type')
        plt.ylabel('Count')
        plt.xticks(rotation=45)
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'summary_statistics.png'), dpi=300, bbox_inches='tight')
    plt.close()
    
    # 3. Severity distribution
    if any(stats['severity_distribution'].values()):
        plt.figure(figsize=(8, 6))
        severities = ['minimal', 'moderate', 'severe']
        counts = [stats['severity_distribution'][s] for s in severities]
        
        plt.bar(severities, counts, color=['yellow', 'orange', 'red'])
        plt.title('Damage Severity Distribution')
        plt.xlabel('Severity Level')
        plt.ylabel('Count')
        
        plt.savefig(os.path.join(output_dir, 'severity_distribution.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    # 4. Confidence distribution histogram
    valid_results = [r for r in batch_results if 'error' not in r]
    if valid_results:
        confidences = [r['overall_status']['overall_confidence'] for r in valid_results]
        
        plt.figure(figsize=(8, 6))
        plt.hist(confidences, bins=20, alpha=0.7, color='blue', edgecolor='black')
        plt.title('Confidence Score Distribution')
        plt.xlabel('Confidence Score')
        plt.ylabel('Frequency')
        plt.axvline(np.mean(confidences), color='red', linestyle='--', 
                   label=f'Mean: {np.mean(confidences):.3f}')
        plt.legend()
        
        plt.savefig(os.path.join(output_dir, 'confidence_distribution.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    logger.info(f"Summary plots saved to: {output_dir}")


def validate_image_file(file_path: str) -> bool:
    """
    Validate if a file is a valid image
    
    Args:
        file_path: Path to the image file
        
    Returns:
        True if valid image, False otherwise
    """
    try:
        with Image.open(file_path) as img:
            img.verify()
        return True
    except Exception:
        return False


def resize_image_smart(image: np.ndarray, 
                      target_size: Tuple[int, int],
                      maintain_aspect: bool = True) -> np.ndarray:
    """
    Smart image resizing with aspect ratio preservation
    
    Args:
        image: Input image
        target_size: Target size (height, width)
        maintain_aspect: Whether to maintain aspect ratio
        
    Returns:
        Resized image
    """
    if not maintain_aspect:
        return cv2.resize(image, (target_size[1], target_size[0]))
    
    h, w = image.shape[:2]
    target_h, target_w = target_size
    
    # Calculate aspect ratios
    aspect_ratio = w / h
    target_aspect = target_w / target_h
    
    if aspect_ratio > target_aspect:
        # Image is wider than target
        new_w = target_w
        new_h = int(target_w / aspect_ratio)
    else:
        # Image is taller than target
        new_h = target_h
        new_w = int(target_h * aspect_ratio)
    
    # Resize image
    resized = cv2.resize(image, (new_w, new_h))
    
    # Create padded image
    if len(image.shape) == 3:
        padded = np.full((target_h, target_w, image.shape[2]), 255, dtype=image.dtype)
    else:
        padded = np.full((target_h, target_w), 255, dtype=image.dtype)
    
    # Calculate padding
    pad_h = (target_h - new_h) // 2
    pad_w = (target_w - new_w) // 2
    
    # Place resized image in center
    if len(image.shape) == 3:
        padded[pad_h:pad_h + new_h, pad_w:pad_w + new_w] = resized
    else:
        padded[pad_h:pad_h + new_h, pad_w:pad_w + new_w] = resized
    
    return padded


def benchmark_detection_speed(detector, image_paths: List[str]) -> Dict:
    """
    Benchmark detection speed on multiple images
    
    Args:
        detector: DamagedPosterDetector instance
        image_paths: List of image paths to test
        
    Returns:
        Benchmarking results
    """
    import time
    
    results = {
        'total_images': len(image_paths),
        'total_time': 0.0,
        'average_time': 0.0,
        'min_time': float('inf'),
        'max_time': 0.0,
        'times': []
    }
    
    for image_path in image_paths:
        start_time = time.time()
        try:
            detector.detect_damage(image_path)
            processing_time = time.time() - start_time
            results['times'].append(processing_time)
            results['total_time'] += processing_time
            results['min_time'] = min(results['min_time'], processing_time)
            results['max_time'] = max(results['max_time'], processing_time)
        except Exception as e:
            logger.warning(f"Error processing {image_path}: {str(e)}")
    
    if results['times']:
        results['average_time'] = results['total_time'] / len(results['times'])
    
    return results