#!/usr/bin/env python3
"""
Create horizontal comparison of IMG_3378 mosaic images with different min_size settings
"""

from PIL import Image, ImageDraw, ImageFont
import os
from pathlib import Path

# Image paths for Figure 3 (2x2 grid)
image_paths = [
    "./samples/akaza.jpg",  # Original
    "./output/test_result/mosaic_images/akaza-min16-max32-sub5-quant0-cifar-10.jpg",
    "./output/test_result/mosaic_images/akaza-min16-max32-sub5-quant16-cifar-10.jpg",
    "./output/test_result/mosaic_images/akaza-min16-max32-sub5-quant32-cifar-10.jpg"
]

# Labels for each image
labels = ["Original", "quant = 0", "quant = 16", "quant = 32"]

def create_comparison():
    """Create 2x2 grid comparison image"""

    # Load images
    images = []
    for path in image_paths:
        if os.path.exists(path):
            img = Image.open(path).convert('RGB')
            images.append(img)
            print(f"[INFO] Loaded: {path} | size = {img.size}")
        else:
            print(f"[ERROR] File not found: {path}")
            return None

    if len(images) != 4:
        print(f"[ERROR] Expected 4 images, found {len(images)}")
        return None

    # Get dimensions and calculate scale
    original_width, original_height = images[0].size

    # Scale down to fit reasonable comparison size (e.g., each image 400px wide)
    target_width = 400
    scale = target_width / original_width
    target_height = int(original_height * scale)

    print(f"[INFO] Scaling from {original_width}x{original_height} to {target_width}x{target_height}")

    # Resize all images to same scale
    resized_images = []
    for img in images:
        resized = img.resize((target_width, target_height), Image.BICUBIC)
        resized_images.append(resized)

    # Create comparison canvas for 2x2 grid
    margin = 20
    label_height = 40
    canvas_width = target_width * 2 + margin * 3
    canvas_height = target_height * 2 + label_height + margin * 4

    canvas = Image.new('RGB', (canvas_width, canvas_height), color='white')

    # Paste images in 2x2 grid
    draw = ImageDraw.Draw(canvas)

    # Try to use a system font, fallback to default
    try:
        font = ImageFont.truetype("/System/Library/Fonts/Arial.ttf", 20)
    except:
        try:
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 20)
        except:
            font = ImageFont.load_default()

    for i, (img, label) in enumerate(zip(resized_images, labels)):
        # Calculate grid position (2x2)
        row = i // 2
        col = i % 2

        x_pos = margin + col * (target_width + margin)
        y_pos = margin + label_height + row * (target_height + margin)

        canvas.paste(img, (x_pos, y_pos))

        # Add label above each image
        bbox = draw.textbbox((0, 0), label, font=font)
        text_width = bbox[2] - bbox[0]
        text_x = x_pos + (target_width - text_width) // 2
        text_y = y_pos - label_height + 5

        draw.text((text_x, text_y), label, fill='black', font=font)

    # Add title
    title = "Figure 3. Akaza Color Quantization Comparison (min16-max32-sub5-cifar-10)"
    try:
        title_font = ImageFont.truetype("/System/Library/Fonts/Arial.ttf", 16)
    except:
        title_font = font

    title_bbox = draw.textbbox((0, 0), title, font=title_font)
    title_width = title_bbox[2] - title_bbox[0]
    title_x = (canvas_width - title_width) // 2
    title_y = canvas_height - margin + 5

    draw.text((title_x, title_y), title, fill='gray', font=title_font)

    return canvas

def main():
    print("=" * 60)
    print("Creating Akaza Color Quantization Comparison")
    print("=" * 60)

    # Create comparison
    comparison_img = create_comparison()

    if comparison_img:
        # Save result
        output_path = "./output/akaza_quantization_comparison.png"
        comparison_img.save(output_path, quality=95)
        print(f"\n[SUCCESS] Comparison saved to: {output_path}")
        print(f"[INFO] Canvas size: {comparison_img.size}")
    else:
        print("\n[ERROR] Failed to create comparison")

if __name__ == "__main__":
    main()