src/visualizer.py

"""
visualizer.py — Interactive visual bounding box overlay and keyword highlighting.
Features:
  - Confidence-coded color bounding boxes (Green/Yellow/Red)
  - Keyword search with box highlighting
  - Heatmap overlay for confidence distribution
  - Annotated word labels on image
"""

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from typing import List, Dict, Any, Optional


def _confidence_color(conf: float) -> tuple:
    """Return BGR color tuple based on confidence level."""
    if conf >= 0.8:
        return (34, 197, 94)   # Green
    elif conf >= 0.5:
        return (234, 179, 8)   # Yellow
    else:
        return (239, 68, 68)   # Red


def _confidence_color_rgba(conf: float) -> tuple:
    """Return RGBA color tuple for PIL drawing."""
    if conf >= 0.8:
        return (34, 197, 94, 200)
    elif conf >= 0.5:
        return (234, 179, 8, 200)
    else:
        return (239, 68, 68, 200)


def draw_bounding_boxes(
    pil_img: Image.Image,
    results: List[Dict[str, Any]],
    show_labels: bool = True,
    show_confidence: bool = True,
    highlight_indices: Optional[List[int]] = None,
    line_width: int = 2,
) -> Image.Image:
    """
    Draw bounding boxes with confidence color coding on a PIL image.

    Args:
        pil_img: Source image
        results: OCR result dicts from ocr_engine
        show_labels: Show text labels above each box
        show_confidence: Show confidence percentage on box
        highlight_indices: List of result indices to highlight in gold
        line_width: Border thickness

    Returns:
        Annotated PIL Image
    """
    img = pil_img.convert("RGBA")
    overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)

    # Try to load a small font
    try:
        font = ImageFont.truetype("arial.ttf", 12)
        small_font = ImageFont.truetype("arial.ttf", 10)
    except:
        font = ImageFont.load_default()
        small_font = font

    highlight_set = set(highlight_indices or [])

    for i, result in enumerate(results):
        bbox = result["bbox"]
        conf = result["confidence"]
        text = result["text"]

        pts = np.array(bbox, dtype=np.int32)
        x_min = min(p[0] for p in pts)
        y_min = min(p[1] for p in pts)
        x_max = max(p[0] for p in pts)
        y_max = max(p[1] for p in pts)

        if i in highlight_set:
            # Gold highlight for keyword matches
            color = (255, 215, 0, 230)
            bg_color = (255, 215, 0, 60)
        else:
            color = _confidence_color_rgba(conf)
            bg_color = (*_confidence_color_rgba(conf)[:3], 30)

        # Fill rectangle
        draw.rectangle([x_min, y_min, x_max, y_max], fill=bg_color, outline=color, width=line_width)

        # Label
        if show_labels:
            label = text[:20] + ("…" if len(text) > 20 else "")
            if show_confidence:
                label = f"{label} ({conf:.0%})"
            label_y = max(0, y_min - 14)
            draw.rectangle(
                [x_min, label_y, x_min + len(label) * 6, label_y + 12],
                fill=(30, 30, 30, 200)
            )
            draw.text((x_min + 1, label_y), label, fill=(255, 255, 255, 255), font=small_font)

    composed = Image.alpha_composite(img, overlay)
    return composed.convert("RGB")


def search_and_highlight(
    pil_img: Image.Image,
    results: List[Dict[str, Any]],
    keyword: str,
    case_sensitive: bool = False
) -> tuple[Image.Image, List[int], List[str]]:
    """
    Search for keyword in OCR results and return highlighted image.

    Returns:
        (annotated_image, matched_indices, matched_texts)
    """
    if not keyword.strip():
        return draw_bounding_boxes(pil_img, results), [], []

    matched_indices = []
    matched_texts = []
    kw = keyword if case_sensitive else keyword.lower()

    for i, result in enumerate(results):
        text = result["text"] if case_sensitive else result["text"].lower()
        if kw in text:
            matched_indices.append(i)
            matched_texts.append(result["text"])

    annotated = draw_bounding_boxes(
        pil_img,
        results,
        show_labels=True,
        show_confidence=True,
        highlight_indices=matched_indices,
    )
    return annotated, matched_indices, matched_texts


def create_confidence_heatmap(
    pil_img: Image.Image,
    results: List[Dict[str, Any]],
    alpha: float = 0.45
) -> Image.Image:
    """
    Overlay a confidence heatmap on the document image.
    High confidence = cool blue, Low confidence = hot red.
    """
    img_cv = np.array(pil_img.convert("RGB"))
    h, w = img_cv.shape[:2]

    heatmap = np.zeros((h, w), dtype=np.float32)
    count_map = np.zeros((h, w), dtype=np.float32)

    for result in results:
        bbox = result["bbox"]
        conf = result["confidence"]
        pts = np.array(bbox, dtype=np.int32)
        x_min = max(0, min(p[0] for p in pts))
        y_min = max(0, min(p[1] for p in pts))
        x_max = min(w - 1, max(p[0] for p in pts))
        y_max = min(h - 1, max(p[1] for p in pts))

        heatmap[y_min:y_max, x_min:x_max] += conf
        count_map[y_min:y_max, x_min:x_max] += 1.0

    count_map[count_map == 0] = 1.0
    heatmap = heatmap / count_map
    heatmap_uint8 = np.uint8(255 * heatmap)
    colored_heatmap = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
    colored_heatmap_rgb = cv2.cvtColor(colored_heatmap, cv2.COLOR_BGR2RGB)

    blended = cv2.addWeighted(img_cv, 1 - alpha, colored_heatmap_rgb, alpha, 0)
    return Image.fromarray(blended)


def draw_reading_order(
    pil_img: Image.Image,
    results: List[Dict[str, Any]],
) -> Image.Image:
    """Draw reading order arrows connecting text regions sequentially."""
    img = np.array(pil_img.convert("RGB"))
    sorted_results = sorted(results, key=lambda r: (r["y_min"], r["x_min"]))

    for i, result in enumerate(sorted_results):
        cx, cy = result["center"]
        cv2.circle(img, (cx, cy), 8, (99, 102, 241), -1)
        cv2.putText(img, str(i + 1), (cx - 4, cy + 4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 255, 255), 1)
        if i < len(sorted_results) - 1:
            nx, ny = sorted_results[i + 1]["center"]
            cv2.arrowedLine(img, (cx, cy), (nx, ny), (99, 102, 241), 1, tipLength=0.3)

    return Image.fromarray(img)