image_comparison_enhanced.py

"""
Enhanced Image Comparison System
Detects and annotates visual differences between Figma and website screenshots
"""

import os
import numpy as np
from typing import List, Dict, Tuple, Any
from dataclasses import dataclass
from PIL import Image, ImageDraw, ImageFont
import logging

logger = logging.getLogger(__name__)


@dataclass
class DifferenceRegion:
    """Represents a region with visual differences."""
    x: int
    y: int
    width: int
    height: int
    severity: str  # "High", "Medium", "Low"
    description: str
    confidence: float


class ImageComparator:
    """Compares two images and detects visual differences."""
    
    @staticmethod
    def compare_images(
        image1_path: str,
        image2_path: str,
        threshold: float = 0.95
    ) -> Tuple[float, List[DifferenceRegion]]:
        """
        Compare two images and detect differences.
        
        Args:
            image1_path: Path to first image (Figma)
            image2_path: Path to second image (Website)
            threshold: Similarity threshold (0-1)
        
        Returns:
            Tuple of (similarity_score, list of difference regions)
        """
        try:
            # Load images
            img1 = Image.open(image1_path).convert('RGB')
            img2 = Image.open(image2_path).convert('RGB')
            
            # Resize to same dimensions for comparison
            if img1.size != img2.size:
                # Resize img2 to match img1
                img2 = img2.resize(img1.size, Image.Resampling.LANCZOS)
            
            # Convert to numpy arrays
            arr1 = np.array(img1, dtype=np.float32)
            arr2 = np.array(img2, dtype=np.float32)
            
            # Calculate pixel-wise difference
            diff = np.abs(arr1 - arr2)
            
            # Calculate similarity score (0-100)
            max_diff = 255.0 * 3  # Max possible difference per pixel (RGB)
            mean_diff = np.mean(diff)
            similarity_score = 100 * (1 - mean_diff / max_diff)
            similarity_score = max(0, min(100, similarity_score))
            
            # Detect difference regions
            difference_regions = ImageComparator._detect_regions(
                diff, img1.size, similarity_score
            )
            
            return similarity_score, difference_regions
        
        except Exception as e:
            logger.error(f"Error comparing images: {str(e)}")
            return 0.0, []
    
    @staticmethod
    def _detect_regions(
        diff_array: np.ndarray,
        image_size: Tuple[int, int],
        similarity_score: float
    ) -> List[DifferenceRegion]:
        """
        Detect regions with significant differences.
        
        Args:
            diff_array: Pixel-wise difference array
            image_size: Size of original image
            similarity_score: Overall similarity score
        
        Returns:
            List of difference regions
        """
        regions = []
        
        # Calculate per-channel difference
        gray_diff = np.mean(diff_array, axis=2)
        
        # Threshold for significant differences
        threshold = 30  # Pixel difference threshold
        significant = gray_diff > threshold
        
        # Find connected components
        from scipy import ndimage
        labeled, num_features = ndimage.label(significant)
        
        # Analyze each region
        for region_id in range(1, num_features + 1):
            region_mask = labeled == region_id
            
            # Skip very small regions (noise)
            if np.sum(region_mask) < 100:
                continue
            
            # Get bounding box
            rows = np.any(region_mask, axis=1)
            cols = np.any(region_mask, axis=0)
            
            if not np.any(rows) or not np.any(cols):
                continue
            
            y_min, y_max = np.where(rows)[0][[0, -1]]
            x_min, x_max = np.where(cols)[0][[0, -1]]
            
            # Calculate region statistics
            region_diff = gray_diff[region_mask]
            mean_diff = np.mean(region_diff)
            max_diff = np.max(region_diff)
            
            # Determine severity
            if max_diff > 100:
                severity = "High"
                confidence = min(1.0, max_diff / 255)
            elif max_diff > 50:
                severity = "Medium"
                confidence = min(1.0, max_diff / 150)
            else:
                severity = "Low"
                confidence = min(1.0, max_diff / 100)
            
            # Generate description
            width = x_max - x_min
            height = y_max - y_min
            description = f"{severity} difference: {width}x{height}px region"
            
            region = DifferenceRegion(
                x=int((x_min + x_max) / 2),
                y=int((y_min + y_max) / 2),
                width=int(width),
                height=int(height),
                severity=severity,
                description=description,
                confidence=float(confidence)
            )
            
            regions.append(region)
        
        # Sort by severity
        severity_order = {"High": 0, "Medium": 1, "Low": 2}
        regions.sort(key=lambda r: severity_order.get(r.severity, 3))
        
        return regions


class ScreenshotAnnotator:
    """Annotates screenshots with visual difference indicators."""
    
    @staticmethod
    def annotate_screenshot(
        screenshot_path: str,
        differences: List[DifferenceRegion],
        output_path: str
    ) -> bool:
        """
        Annotate screenshot with markers for differences.
        
        Args:
            screenshot_path: Path to original screenshot
            differences: List of visual differences
            output_path: Path to save annotated screenshot
        
        Returns:
            True if successful
        """
        try:
            if not os.path.exists(screenshot_path):
                return False
            
            # Load image
            img = Image.open(screenshot_path).convert('RGB')
            draw = ImageDraw.Draw(img, 'RGBA')
            
            # Draw circles and labels for each difference
            circle_radius = 40
            
            for idx, diff in enumerate(differences):
                # Draw circle
                circle_color = ScreenshotAnnotator._get_color_by_severity(diff.severity)
                
                x, y = diff.x, diff.y
                draw.ellipse(
                    [(x - circle_radius, y - circle_radius),
                     (x + circle_radius, y + circle_radius)],
                    outline=circle_color,
                    width=4
                )
                
                # Draw number label
                label_number = str(idx + 1)
                try:
                    # Try to use a larger font
                    font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 24)
                except:
                    font = ImageFont.load_default()
                
                # Draw label with background
                label_bbox = draw.textbbox((x - 8, y - 8), label_number, font=font)
                draw.rectangle(label_bbox, fill=circle_color)
                draw.text(
                    (x - 8, y - 8),
                    label_number,
                    fill=(255, 255, 255),
                    font=font
                )
                
                # Draw bounding box around region
                box_x1 = x - diff.width // 2
                box_y1 = y - diff.height // 2
                box_x2 = x + diff.width // 2
                box_y2 = y + diff.height // 2
                
                draw.rectangle(
                    [(box_x1, box_y1), (box_x2, box_y2)],
                    outline=circle_color,
                    width=2
                )
            
            # Create output directory
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # Save annotated image
            img.save(output_path)
            return True
        
        except Exception as e:
            logger.error(f"Error annotating screenshot: {str(e)}")
            return False
    
    @staticmethod
    def _get_color_by_severity(severity: str) -> Tuple[int, int, int, int]:
        """Get color based on severity level."""
        if severity == "High":
            return (255, 0, 0, 220)  # Red
        elif severity == "Medium":
            return (255, 165, 0, 220)  # Orange
        else:
            return (0, 200, 0, 220)  # Green
    
    @staticmethod
    def create_side_by_side_comparison(
        figma_screenshot: str,
        website_screenshot: str,
        figma_annotated: str,
        website_annotated: str,
        output_path: str,
        title: str = "Figma vs Website"
    ) -> bool:
        """
        Create side-by-side comparison image with labels.
        
        Args:
            figma_screenshot: Original Figma screenshot
            website_screenshot: Original website screenshot
            figma_annotated: Annotated Figma screenshot
            website_annotated: Annotated website screenshot
            output_path: Path to save comparison
            title: Title for the comparison
        
        Returns:
            True if successful
        """
        try:
            # Load annotated images
            figma_img = Image.open(figma_annotated).convert('RGB')
            website_img = Image.open(website_annotated).convert('RGB')
            
            # Resize to same height
            max_height = max(figma_img.height, website_img.height)
            figma_img = figma_img.resize(
                (int(figma_img.width * max_height / figma_img.height), max_height),
                Image.Resampling.LANCZOS
            )
            website_img = website_img.resize(
                (int(website_img.width * max_height / website_img.height), max_height),
                Image.Resampling.LANCZOS
            )
            
            # Create header space
            header_height = 60
            total_width = figma_img.width + website_img.width + 40
            total_height = max_height + header_height + 40
            
            # Create comparison image
            comparison = Image.new('RGB', (total_width, total_height), (255, 255, 255))
            draw = ImageDraw.Draw(comparison)
            
            # Draw title
            try:
                font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
            except:
                font = ImageFont.load_default()
            
            draw.text((20, 15), title, fill=(0, 0, 0), font=font)
            
            # Draw labels
            try:
                label_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 16)
            except:
                label_font = ImageFont.load_default()
            
            draw.text((20, header_height + 10), "Figma Design", fill=(0, 0, 0), font=label_font)
            draw.text((figma_img.width + 40, header_height + 10), "Website", fill=(0, 0, 0), font=label_font)
            
            # Paste images
            comparison.paste(figma_img, (20, header_height + 30))
            comparison.paste(website_img, (figma_img.width + 40, header_height + 30))
            
            # Create output directory
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # Save comparison
            comparison.save(output_path)
            return True
        
        except Exception as e:
            logger.error(f"Error creating comparison image: {str(e)}")
            return False


def create_difference_report(
    differences: List[DifferenceRegion],
    similarity_score: float,
    viewport: str
) -> Dict[str, Any]:
    """
    Create a detailed report of detected differences.
    
    Args:
        differences: List of detected differences
        similarity_score: Overall similarity score
        viewport: Viewport name (desktop/mobile)
    
    Returns:
        Dictionary with report data
    """
    high_severity = len([d for d in differences if d.severity == "High"])
    medium_severity = len([d for d in differences if d.severity == "Medium"])
    low_severity = len([d for d in differences if d.severity == "Low"])
    
    report = {
        "viewport": viewport,
        "similarity_score": similarity_score,
        "total_differences": len(differences),
        "high_severity": high_severity,
        "medium_severity": medium_severity,
        "low_severity": low_severity,
        "differences": [
            {
                "id": idx + 1,
                "severity": diff.severity,
                "location": {"x": diff.x, "y": diff.y},
                "size": {"width": diff.width, "height": diff.height},
                "description": diff.description,
                "confidence": diff.confidence
            }
            for idx, diff in enumerate(differences)
        ]
    }
    
    return report