"""
Visualization utilities for optik.

Drawing functions for detections, heatmaps, and grid cells.
"""

from typing import List, Tuple, Optional, Union
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import torch

from ..core.base import Detection


def create_heatmap_overlay(
    image: Image.Image,
    heatmap: Union[np.ndarray, torch.Tensor],
    alpha: float = 0.5,
    colormap: str = 'jet'
) -> Image.Image:
    """
    Create image with heatmap overlay.
    
    Args:
        image: Base image
        heatmap: 2D attribution map
        alpha: Overlay opacity
        colormap: Color scheme
        
    Returns:
        Blended image
    """
    from .image import overlay_heatmap
    return overlay_heatmap(image, heatmap, alpha, colormap)


def draw_detections(
    image: Image.Image,
    detections: List[Detection],
    colors: Optional[dict] = None,
    line_width: int = 2,
    font_size: int = 12
) -> Image.Image:
    """
    Draw detection boxes and labels on image.
    
    Args:
        image: PIL Image to draw on
        detections: List of Detection objects
        colors: Dict mapping class_id to RGB tuple
        line_width: Box line width
        font_size: Label font size
        
    Returns:
        Image with drawn detections
    """
    image = image.copy()
    draw = ImageDraw.Draw(image)
    
    # Default color palette
    default_colors = [
        (255, 0, 0), (0, 255, 0), (0, 0, 255),
        (255, 255, 0), (255, 0, 255), (0, 255, 255),
        (128, 0, 0), (0, 128, 0), (0, 0, 128),
        (128, 128, 0), (128, 0, 128), (0, 128, 128),
    ]
    
    try:
        font = ImageFont.truetype("arial.ttf", font_size)
    except:
        font = ImageFont.load_default()
    
    for det in detections:
        # Get color
        if colors and det.class_id in colors:
            color = colors[det.class_id]
        else:
            color = default_colors[det.class_id % len(default_colors)]
        
        # Draw box
        x1, y1, x2, y2 = det.bbox
        draw.rectangle([x1, y1, x2, y2], outline=color, width=line_width)
        
        # Draw label
        label = f"{det.class_name} {det.confidence:.0%}"
        
        # Label background
        bbox = draw.textbbox((x1, y1), label, font=font)
        draw.rectangle(
            [bbox[0] - 2, bbox[1] - 2, bbox[2] + 2, bbox[3] + 2],
            fill=color
        )
        
        # Label text (black or white based on background brightness)
        brightness = sum(color) / 3
        text_color = (0, 0, 0) if brightness > 128 else (255, 255, 255)
        draw.text((x1, y1), label, fill=text_color, font=font)
    
    return image


def draw_grid_cells(
    image: Image.Image,
    cells: list,
    grid_sizes: dict = None,
    colors: dict = None,
    alpha: float = 0.3
) -> Image.Image:
    """
    Draw grid cell highlights on image.
    
    Args:
        image: Base image
        cells: List of GridCell objects
        grid_sizes: Dict mapping scale to grid size
        colors: Dict mapping scale to color
        alpha: Cell highlight opacity
        
    Returns:
        Image with highlighted grid cells
    """
    if grid_sizes is None:
        grid_sizes = {'p3': 80, 'p4': 40, 'p5': 20}
    
    if colors is None:
        colors = {
            'p3': (0, 255, 0, int(255 * alpha)),    # Green for small
            'p4': (255, 255, 0, int(255 * alpha)),  # Yellow for medium
            'p5': (255, 0, 0, int(255 * alpha)),    # Red for large
        }
    
    # Create overlay
    overlay = Image.new('RGBA', image.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)
    
    w, h = image.size
    
    for cell in cells:
        scale = cell.scale
        grid_size = grid_sizes.get(scale, 40)
        color = colors.get(scale, (128, 128, 128, int(255 * alpha)))
        
        # Calculate cell bounds
        cell_w = w / grid_size
        cell_h = h / grid_size
        
        x1 = cell.x * cell_w
        y1 = cell.y * cell_h
        x2 = (cell.x + 1) * cell_w
        y2 = (cell.y + 1) * cell_h
        
        draw.rectangle([x1, y1, x2, y2], fill=color, outline=color[:3])
    
    # Composite
    result = image.convert('RGBA')
    result = Image.alpha_composite(result, overlay)
    return result.convert('RGB')


def create_side_by_side(
    images: List[Image.Image],
    labels: Optional[List[str]] = None,
    padding: int = 10
) -> Image.Image:
    """
    Create side-by-side comparison image.
    
    Args:
        images: List of images to combine
        labels: Optional labels for each image
        padding: Space between images
        
    Returns:
        Combined image
    """
    if not images:
        raise ValueError("At least one image required")
    
    # Make all same height
    max_h = max(img.height for img in images)
    resized = []
    for img in images:
        if img.height != max_h:
            ratio = max_h / img.height
            new_w = int(img.width * ratio)
            img = img.resize((new_w, max_h))
        resized.append(img)
    
    # Calculate total width
    total_w = sum(img.width for img in resized) + padding * (len(resized) - 1)
    label_height = 30 if labels else 0
    
    # Create canvas
    result = Image.new('RGB', (total_w, max_h + label_height), (255, 255, 255))
    draw = ImageDraw.Draw(result)
    
    try:
        font = ImageFont.truetype("arial.ttf", 14)
    except:
        font = ImageFont.load_default()
    
    # Paste images
    x_offset = 0
    for i, img in enumerate(resized):
        result.paste(img, (x_offset, label_height))
        
        if labels and i < len(labels):
            # Draw label
            text_bbox = draw.textbbox((0, 0), labels[i], font=font)
            text_w = text_bbox[2] - text_bbox[0]
            text_x = x_offset + (img.width - text_w) // 2
            draw.text((text_x, 5), labels[i], fill=(0, 0, 0), font=font)
        
        x_offset += img.width + padding
    
    return result
