"""Class-map and object previews: the two PNGs a person reads to trust `objects.json`.

* `class_preview`  — the seg head's class map blended over the optical image.
* `objects_overlay` — detected tree crowns, building footprints and water
  outlines drawn on the optical image.  This is the check for the object step:
  a crown circle that does not sit on a tree is visible at a glance.

Pure numpy + Pillow, like `figures.colorize` / `figures.hillshade`.  The palette
matches the frontend's "Object classes" layer so the two read the same.
"""

from __future__ import annotations

import numpy as np
from PIL import Image, ImageDraw

CLASS_COLORS: dict[str, tuple[int, int, int]] = {
    "other": (154, 160, 166),
    "ground": (200, 181, 138),
    "low_veg": (156, 204, 101),
    "building": (224, 90, 71),
    "water": (58, 143, 217),
    "road": (95, 99, 104),
    "tree": (46, 125, 50),
}
_UNKNOWN = (154, 160, 166)


def colorize_classes(seg: np.ndarray, names: tuple[str, ...]) -> np.ndarray:
    """(H, W) class ids -> (H, W, 3) uint8, by class *name* so ids can move."""
    lut = np.array([CLASS_COLORS.get(n, _UNKNOWN) for n in names] + [_UNKNOWN] * 256,
                   np.uint8)[:256]
    return lut[np.clip(seg, 0, 255).astype(np.uint8)]


def class_preview(rgb: np.ndarray, seg: np.ndarray, names: tuple[str, ...],
                  alpha: float = 0.6) -> Image.Image:
    color = colorize_classes(seg, names).astype(np.float32)
    out = alpha * color + (1.0 - alpha) * rgb.astype(np.float32)
    return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8))


def objects_overlay(rgb: np.ndarray, objects: dict) -> Image.Image:
    """Crowns as circles of their radius, footprints and water as outlines."""
    gsd = float(objects["grid"]["gsd_m"])
    base = Image.fromarray((rgb.astype(np.float32) * 0.7).astype(np.uint8))
    draw = ImageDraw.Draw(base)
    line = max(1, round(max(rgb.shape[:2]) / 800))

    for w in objects["water"]:
        draw.polygon([tuple(pt) for pt in w["poly"]], outline=CLASS_COLORS["water"], width=line)
    for b in objects["buildings"]:
        draw.polygon([tuple(pt) for pt in b["poly"]], outline=(255, 120, 90), width=line)
    for t in objects["trees"]:
        r = t["r"] / gsd
        x, y = t["x"], t["y"]
        draw.ellipse((x - r, y - r, x + r, y + r), outline=(120, 230, 120), width=line)
        d = max(1.0, line)
        draw.ellipse((x - d, y - d, x + d, y + d), fill=(255, 255, 255))
    return base