src/editable_image/style.py

from __future__ import annotations

from dataclasses import dataclass

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

from .fonts import FACES
from .models import Box, Point, TextStyle
from .ocr import quad_box, quad_rotation


@dataclass(slots=True)
class StyleResult:
    box: Box
    rotation: float
    style: TextStyle


def estimate_style(rgba: np.ndarray, text: str, quad: list[Point]) -> StyleResult:
    x, y, width, height = quad_box(quad)
    max_height, max_width = rgba.shape[:2]
    left = max(0, int(x))
    top = max(0, int(y))
    right = min(max_width, int(np.ceil(x + width)))
    bottom = min(max_height, int(np.ceil(y + height)))
    crop = rgba[top:bottom, left:right, :3]
    font_size = max(6.0, height * 0.92)
    foreground = estimate_foreground_color(crop)
    face, adjusted_size = match_font(crop, text, font_size, width)
    style = TextStyle(
        font_family=face.family,
        font_size=adjusted_size,
        font_weight=700 if face.weight >= 700 else 400,
        italic=face.italic,
        fill=foreground,
    )
    return StyleResult(
        box=Box(x=x, y=y, width=max(1, width), height=max(1, height)),
        rotation=max(-180, min(180, quad_rotation(quad))),
        style=style,
    )


def estimate_foreground_color(crop: np.ndarray) -> str:
    if crop.size == 0:
        return "#111111"
    pixels = crop.reshape(-1, 3).astype(np.float32)
    border = np.concatenate(
        [crop[0, :, :3], crop[-1, :, :3], crop[:, 0, :3], crop[:, -1, :3]], axis=0
    ).astype(np.float32)
    background = np.median(border, axis=0)
    distances = np.linalg.norm(pixels - background, axis=1)
    threshold = max(24.0, float(np.percentile(distances, 75)))
    candidates = pixels[distances >= threshold]
    if len(candidates) < 3:
        candidates = pixels[np.argsort(distances)[-max(3, len(pixels) // 10) :]]
    color = np.median(candidates, axis=0).clip(0, 255).astype(int)
    return "#{:02x}{:02x}{:02x}".format(*color)


def match_font(crop: np.ndarray, text: str, size: float, target_width: float):
    available = [face for face in FACES if face.family == "Lato" and face.path.exists()]
    if not available:
        raise FileNotFoundError("Bundled fonts are missing")
    best_face = available[0]
    best_size = max(6, int(round(size)))
    best_score = float("inf")
    observed = foreground_binary(crop)
    likely_italic = estimate_italic(observed)
    available = [face for face in available if face.italic == likely_italic]
    for face in available:
        candidate_size = max(6, int(round(size)))
        font = ImageFont.truetype(str(face.path), candidate_size)
        bounds = font.getbbox(text or " ")
        rendered_width = max(1, bounds[2] - bounds[0])
        scaled_size = int(np.clip(candidate_size * target_width / rendered_width, 6, 1000))
        mask = render_text_mask(text, face.path, scaled_size, crop.shape[:2])
        score = mask_distance(observed, mask)
        score += 0.002 if face.weight >= 700 else 0
        if score < best_score:
            best_score = score
            best_face = face
            best_size = scaled_size
    return best_face, float(best_size)


def foreground_binary(crop: np.ndarray) -> np.ndarray:
    if crop.size == 0:
        return np.zeros((1, 1), dtype=np.uint8)
    gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY)
    border = np.concatenate([gray[0], gray[-1], gray[:, 0], gray[:, -1]])
    background = float(np.median(border))
    delta = cv2.absdiff(gray, np.full_like(gray, int(background)))
    _, mask = cv2.threshold(delta, 18, 255, cv2.THRESH_BINARY)
    return mask


def render_text_mask(text: str, path, size: int, shape: tuple[int, int]) -> np.ndarray:
    height, width = max(1, shape[0]), max(1, shape[1])
    image = Image.new("L", (width, height), 0)
    draw = ImageDraw.Draw(image)
    font = ImageFont.truetype(str(path), size)
    bounds = draw.textbbox((0, 0), text, font=font)
    y = max(-bounds[1], (height - (bounds[3] - bounds[1])) // 2 - bounds[1])
    draw.text((0, y), text, font=font, fill=255)
    return np.asarray(image)


def mask_distance(observed: np.ndarray, candidate: np.ndarray) -> float:
    if observed.shape != candidate.shape:
        candidate = cv2.resize(candidate, (observed.shape[1], observed.shape[0]))
    observed_pixels = observed > 0
    candidate_pixels = candidate > 0
    if not np.any(observed_pixels) or not np.any(candidate_pixels):
        return 1.0
    distance_to_observed = cv2.distanceTransform(
        (~observed_pixels).astype(np.uint8), cv2.DIST_L2, 3
    )
    distance_to_candidate = cv2.distanceTransform(
        (~candidate_pixels).astype(np.uint8), cv2.DIST_L2, 3
    )
    chamfer = (
        float(distance_to_observed[candidate_pixels].mean())
        + float(distance_to_candidate[observed_pixels].mean())
    ) / (2 * max(1, observed.shape[0]))
    density_delta = abs(
        float(candidate_pixels.mean()) - float(observed_pixels.mean())
    )
    return chamfer + density_delta * 0.45


def estimate_italic(mask: np.ndarray) -> bool:
    if mask.size == 0 or np.count_nonzero(mask) < 12:
        return False
    height, width = mask.shape
    shears = np.linspace(-0.3, 0.3, 13)
    scores: list[float] = []
    for shear in shears:
        transform = np.float32([[1, shear, -shear * height / 2], [0, 1, 0]])
        deskewed = cv2.warpAffine(mask, transform, (width, height)) > 0
        columns = deskewed.sum(axis=0).astype(np.float32)
        scores.append(float(np.square(columns).sum() / max(1, columns.sum())))
    return float(shears[int(np.argmax(scores))]) >= 0.12


def snap_sizes_to_distribution_peaks(sizes: list[float]) -> list[float]:
    if len(sizes) < 2:
        return list(sizes)
    values = np.asarray(sizes, dtype=np.float64)
    bandwidth = max(0.75, float(np.median(values)) * 0.06)
    peaks = font_size_distribution_peaks(values, bandwidth)
    if not peaks:
        return list(sizes)

    snap_distance = bandwidth * 2
    snapped: list[float] = []
    for value in values:
        nearest = min(peaks, key=lambda peak: abs(peak - value))
        snapped.append(nearest if abs(nearest - value) <= snap_distance else float(value))
    return snapped


def font_size_distribution_peaks(values: np.ndarray, bandwidth: float) -> list[float]:
    lower = float(values.min())
    upper = float(values.max())
    if upper - lower < 1e-6:
        return [lower]

    sample_count = min(
        4096,
        max(64, int(np.ceil((upper - lower) / (bandwidth / 12))) + 1),
    )
    grid = np.linspace(lower, upper, sample_count)
    offsets = (grid[:, None] - values[None, :]) / bandwidth
    density = np.exp(-0.5 * np.square(offsets)).sum(axis=1)

    candidates: list[int] = []
    if density[0] >= density[1]:
        candidates.append(0)
    candidates.extend(
        index
        for index in range(1, sample_count - 1)
        if density[index] >= density[index - 1]
        and density[index] >= density[index + 1]
    )
    if density[-1] >= density[-2]:
        candidates.append(sample_count - 1)

    peaks: list[float] = []
    for index in sorted(candidates, key=lambda item: density[item], reverse=True):
        center = grid[index]
        distances = np.abs(values - center)
        supported = values[distances <= bandwidth * 2]
        if len(supported) < 2:
            continue
        peak = round(float(np.mean(supported)) * 2) / 2
        if all(abs(peak - existing) > bandwidth * 2 for existing in peaks):
            peaks.append(peak)
    return sorted(peaks)