"""Synthetic 28x28 images for the two picture levels, and the compact form they ship in.

An image is stored as a binary shape mask (98 bytes, base64), a foreground level, a
background level and (for the watermark level) a flag saying whether the corner mark is
drawn. `compose` turns that back into 784 pixel values in 0..15. web/models.js has the
same `compose`, so the browser and the Python train on identical pixels.
"""
from __future__ import annotations

import base64

import numpy as np

SIZE = 28
LEVELS = 16  # pixel values are 0..15

# The watermark: a faint 2x2 block near the top-right corner.
MARK_CELLS = [(2, 23), (2, 24), (3, 23), (3, 24)]
MARK_VALUE = 3


def shape_mask(kind: str, cx: float, cy: float, r: float) -> np.ndarray:
    """A filled circle, square or upward triangle as a 0/1 array, without antialiasing."""
    yy, xx = np.mgrid[0:SIZE, 0:SIZE].astype(float)
    if kind == "circle":
        m = (xx - cx) ** 2 + (yy - cy) ** 2 <= r * r
    elif kind == "square":
        m = (np.abs(xx - cx) <= r) & (np.abs(yy - cy) <= r)
    elif kind == "triangle":
        # apex at (cx, cy - r), base at y = cy + r, half-width grows linearly down the shape
        t = (yy - (cy - r)) / (2.0 * r)
        m = (t >= 0) & (t <= 1) & (np.abs(xx - cx) <= t * r)
    else:
        raise ValueError(kind)
    return m.astype(np.uint8)


def pack_mask(mask: np.ndarray) -> str:
    return base64.b64encode(np.packbits(mask.astype(np.uint8).ravel()).tobytes()).decode("ascii")


def unpack_mask(b64: str) -> np.ndarray:
    bits = np.unpackbits(np.frombuffer(base64.b64decode(b64), dtype=np.uint8))
    return bits[:SIZE * SIZE].reshape(SIZE, SIZE)


def compose(mask: np.ndarray, fg: int, bg: int, mark: bool = False) -> np.ndarray:
    """Pixels 0..15: the shape in `fg` on a flat background `bg`, plus the optional mark."""
    img = np.where(mask > 0, fg, bg).astype(np.int16)
    if mark:
        for r, c in MARK_CELLS:
            img[r, c] = MARK_VALUE
    return img


def compose_record(rec: dict, mark_override=None, bg_override=None) -> np.ndarray:
    """A shipped image record {"b","fg","bg","mk"} -> flat float pixels in [0, 1]."""
    mark = rec.get("mk", 0) if mark_override is None else mark_override
    bg = rec["bg"] if bg_override is None else bg_override
    return compose(unpack_mask(rec["b"]), rec["fg"], bg, bool(mark)).astype(float).ravel() / (LEVELS - 1)