"""Visualization utilities for face detection results."""

import numpy as np
import cv2
from typing import List, Optional, Tuple


def draw_detections(image: np.ndarray, boxes: np.ndarray,
                    scores: Optional[np.ndarray] = None,
                    track_ids: Optional[np.ndarray] = None,
                    landmarks: Optional[np.ndarray] = None,
                    color: Tuple[int, int, int] = (0, 255, 0),
                    thickness: int = 2) -> np.ndarray:
    """Draw face detections on image."""
    vis = image.copy()

    for i in range(len(boxes)):
        x1, y1, x2, y2 = boxes[i].astype(int)
        cv2.rectangle(vis, (x1, y1), (x2, y2), color, thickness)

        label_parts = []
        if track_ids is not None:
            label_parts.append(f"ID:{track_ids[i]}")
        if scores is not None:
            label_parts.append(f"{scores[i]:.2f}")
        label = ' '.join(label_parts)

        if label:
            (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.rectangle(vis, (x1, y1 - th - 4), (x1 + tw, y1), color, -1)
            cv2.putText(vis, label, (x1, y1 - 2),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

        if landmarks is not None and i < len(landmarks):
            lmk = landmarks[i]
            colors_lmk = [(0, 0, 255), (0, 255, 255), (0, 255, 0),
                          (255, 255, 0), (255, 0, 0)]
            for j in range(5):
                x, y = int(lmk[j*2]), int(lmk[j*2+1])
                if x > 0 and y > 0:
                    cv2.circle(vis, (x, y), 3, colors_lmk[j], -1)

    return vis


def create_comparison_grid(images: List[np.ndarray], titles: List[str],
                           cols: int = 3, cell_size: Tuple[int, int] = (400, 400)
                           ) -> np.ndarray:
    """Create a grid of images with titles for comparison."""
    rows = (len(images) + cols - 1) // cols
    grid = np.zeros((rows * cell_size[1], cols * cell_size[0], 3), dtype=np.uint8)

    for idx, (img, title) in enumerate(zip(images, titles)):
        r, c = idx // cols, idx % cols
        resized = cv2.resize(img, cell_size)
        y1, y2 = r * cell_size[1], (r + 1) * cell_size[1]
        x1, x2 = c * cell_size[0], (c + 1) * cell_size[0]
        grid[y1:y2, x1:x2] = resized
        cv2.putText(grid, title, (x1 + 5, y1 + 20),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

    return grid