"""Stroke groups -> 512x512 grayscale tensor.

Both real and generated training images go through this single function so
the classifier can't latch onto rendering style as a shortcut.
"""

from __future__ import annotations

import numpy as np
from PIL import Image, ImageDraw

CANVAS_SIZE = 256


def render_strokes(
    word_groups: list[list[dict]],
    *,
    padding: int = 12,
    stroke_width: int = 2,
) -> np.ndarray:
    """Rasterize nested word->stroke->point groups into `[CANVAS_SIZE, CANVAS_SIZE, 1]`
    float32, with 0 = background and 1 = ink. The point bbox is fit isotropically
    into the canvas minus `padding` and centered."""
    points = [
        (point["x"], point["y"])
        for group in word_groups
        for stroke in group
        for point in stroke["points"]
    ]
    image = Image.new("L", (CANVAS_SIZE, CANVAS_SIZE), 0)
    if not points:
        return np.asarray(image, dtype=np.float32)[..., None] / 255.0

    xs, ys = zip(*points)
    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)
    span_x = max_x - min_x
    span_y = max_y - min_y
    available = CANVAS_SIZE - 2 * padding
    if span_x == 0 and span_y == 0:
        scale = 1.0
    else:
        scale = available / max(span_x, span_y, 1e-6)
    drawn_width = span_x * scale
    drawn_height = span_y * scale
    offset_x = (CANVAS_SIZE - drawn_width) / 2 - min_x * scale
    offset_y = (CANVAS_SIZE - drawn_height) / 2 - min_y * scale

    canvas = ImageDraw.Draw(image)
    for group in word_groups:
        for stroke in group:
            transformed = [
                (point["x"] * scale + offset_x, point["y"] * scale + offset_y)
                for point in stroke["points"]
            ]
            if len(transformed) >= 2:
                canvas.line(transformed, fill=255, width=stroke_width)
            elif len(transformed) == 1:
                x_value, y_value = transformed[0]
                radius = stroke_width / 2
                canvas.ellipse(
                    [x_value - radius, y_value - radius, x_value + radius, y_value + radius],
                    fill=255,
                )

    return np.asarray(image, dtype=np.float32)[..., None] / 255.0