runtime/src/decisionnode_edge/places.py

"""Point and box questions: where in a picture the thing the instructions name is.

    {"type": "point", "instructions": "Point to the yellow package closest to the robot arm.", "image": "belt"}
    {"type": "point", "x": 812, "y": 455, "present": 0.97, "confidence": 0.91, "image": "belt", "width": 1920, "height": 1080}
    {"type": "box", "x1": 300, "y1": 120, "x2": 520, "y2": 410, "present": 0.95, "confidence": 0.88, "image": "belt", ...}

The question. `instructions` (a string, object or array; required), `image` (the id of a picture of the request;
optional when the request carries exactly one), `units` ("pixels", the default, or "normalized"), `max` (an integer from
2 to 16: a list of up to that many places in place of one). `criteria` and any other field are ignored. A request with a
point or box question needs a picture. Every faulty field is listed in one 422 (`question_errors`, the question's type
in the path: ["body", "questions", "pick", "point", "image"]); a missing `instructions` is then the 400 of every type.

How one is answered: a few readings of the model, pass by pass, each a prompt of its own, a later pass given what the
earlier ones read:

    point   pass 0: presence, x           pass 1: y given x
    box     pass 0: presence, x (centre)  pass 1: y (centre) given x   pass 2: width and height given the centre

The presence reading is a yes/no question ("is it in the image?") and gives `present`. A coordinate reading lists no
options: code k is step k of 256 along one side of the picture (step k stands for the fraction (k + 0.5) / 256 of the
side), its scores read as a plain softmax. Its value is the most probable step (the lowest on a tie), then the
probability-weighted mean of the steps within WINDOW = 2 of it (the answer between steps; a distant second peak, another
object, stays out of the mean), and its peak mass is the probability within that window. A later reading is given the
earlier ones rounded to whole steps. `confidence` is the product of the coordinate readings' peak masses (provisional:
not calibrated).

Units: `pixels` are whole pixels of the picture as sent, after its EXIF orientation, halves up: a point within
[0, width - 1] x [0, height - 1], a box's corners within [0, width] x [0, height]; `normalized` gives fractions of the
width and height with 4 decimals, halves up, within [0, 1]. A box's corners are its centre minus and plus half its
size, clipped to the picture, so x1 <= x2 and y1 <= y2. `present` and `confidence` to 2 decimals.

Several places (`max`):

    pass 0: presence, count          pass 1: every place's x        pass 2: every y given its x
    pass 3 (a box): every width and height given its centre

The count reading answers 0 to 16, or 17 for more (codes 0 to 17): `count` is min(level, max), `complete` the
probability of a level at or under `max`. The k-th place's readings carry its rank among the level read ("of 5, the
2nd from the left"). The answer lists the places from the left by the centre's x, ties from the top, each with its own
confidence; two boxes at an intersection over union of MERGE_IOU = 0.7 or more are one thing, and the one of lower
confidence is dropped (`count` is the places listed). A count of 0 ends after pass 0.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Any, Iterator, Sequence

from .answers import round2

TYPES = ("point", "box")
UNITS = ("pixels", "normalized")
N_STEPS = 256  # steps per side: the model's 256 answer codes
WINDOW = 2  # a coordinate reading: the steps within this many of the most probable one are averaged
COUNT_MAX = 16  # the most places a list holds
COUNT_MORE = COUNT_MAX + 1  # the count reading's level for more than COUNT_MAX
COUNT_LEVELS = COUNT_MORE + 1  # the count reading's codes: 0 to 17
MAX_RANGE = (2, COUNT_MAX)  # a question's `max`
MERGE_IOU = 0.7  # two listed boxes overlapping at this intersection over union or more are one thing
CHAIN = {"point": (("x", 0), ("y", 1)), "box": (("x", 0), ("y", 1), ("width", 2), ("height", 2))}  # (axis, pass)
GIVEN = {"x": (), "y": ("x",), "width": ("x", "y"), "height": ("x", "y")}  # the earlier readings an axis is given

_CONTENT_MEMBERS = (
    ("string_type", "str", "Input should be a valid string"),
    ("dict_type", "dict[any,any]", "Input should be a valid dictionary"),
    ("list_type", "list[any]", "Input should be a valid list"),
)


def is_place_question(question: Any) -> bool:
    return isinstance(question, dict) and question.get("type") in TYPES


@dataclass
class PlaceSpec:
    """One point or box question of a request."""

    qid: str
    shape: str  # "point" or "box"
    instructions: str  # as the prompt prints them
    image: str  # the id of the picture to look in
    number: int  # that picture's place among the request's pictures, from 1: how the readings name it
    units: str = "pixels"
    max: int | None = None  # the most places to list (2 to 16); None: the one-place answer


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------


def field_error(kind: str, loc: Sequence[Any], msg: str, value: Any, **extra: Any) -> dict:
    """One entry of a 422's `detail` list."""
    return {"type": kind, "loc": list(loc), "msg": msg, "input": value, **extra}


def max_errors(value: Any, here: Sequence[Any]) -> Iterator[dict]:
    """The entries of a question's `max`: not an integer (a bool is not one), under 2, over 16. Null is no `max`."""
    lo, hi = MAX_RANGE
    if value is None:
        return
    if isinstance(value, bool) or not isinstance(value, int):
        yield field_error("int_type", (*here, "max"), "Input should be a valid integer", value)
    elif value < lo:
        yield field_error("greater_than_equal", (*here, "max"), f"Input should be greater than or equal to {lo}", value, ctx={"ge": lo})
    elif value > hi:
        yield field_error("less_than_equal", (*here, "max"), f"Input should be less than or equal to {hi}", value, ctx={"le": hi})


def question_errors(question: dict, loc: Sequence[Any], image_ids: Sequence[str] | None) -> Iterator[dict]:
    """The 422 entries of one point or box question (`loc` ends in its id), in order: `instructions` of another kind than
    a string, an object or an array, `units`, `max`, then `image` against the ids of the request's pictures. `image_ids`
    is None when the pictures themselves are faulty (their own error names them), so only the question's fields are
    checked. A missing `instructions` is not one of them (`missing_instructions`)."""
    shape = question["type"]
    here = (*loc, shape)
    instructions = question.get("instructions")
    if instructions is not None and not isinstance(instructions, (str, dict, list)):
        for kind, member, msg in _CONTENT_MEMBERS:
            yield field_error(kind, (*here, "instructions", member), msg, instructions)
    units = question.get("units")
    if units is not None and units not in UNITS:
        yield field_error("literal_error", (*here, "units"), "Input should be 'pixels' or 'normalized'", units, ctx={"expected": "'pixels' or 'normalized'"})
    yield from max_errors(question.get("max"), here)
    image = question.get("image")
    if image_ids is None:
        if image is not None and not isinstance(image, str):
            yield field_error("string_type", (*here, "image"), "Input should be a valid string", image)
        return
    if not image_ids:
        yield field_error("image_required", here, f"A {shape} question needs an image in the request", question)
    elif image is None:
        if len(image_ids) > 1:
            yield field_error("missing", (*here, "image"), "Field required", question)
    elif not isinstance(image, str):  # a picture's number is not an id: "Picture 2" goes in the instructions
        yield field_error("string_type", (*here, "image"), "Input should be a valid string", image)
    elif image not in image_ids:
        yield field_error("unknown_image", (*here, "image"), "Input should be the id of an image of the request", image)


def missing_instructions(questions: Sequence[tuple[str, dict]]) -> str | None:
    """The 400's message for the first point or box question without instructions, or None."""
    for qid, question in questions:
        if question.get("instructions") is None or question.get("instructions") == "":
            return f"{question['type'].capitalize()} question must have instructions: {qid}"
    return None


def where(entry: dict) -> str:
    """Where a 422 entry points, for the error's message: body.questions.'pick'.'point'.'image': Field required."""
    parts = (part[:100] + "..." if isinstance(part, str) and len(part) > 100 else part for part in entry["loc"])
    return ".".join(str(part) if i < 2 else repr(part) for i, part in enumerate(parts)) + ": " + entry["msg"]


def reason(errors: Sequence[dict], shown: int = 5) -> str:
    text = ", ".join(where(e) for e in errors[:shown])
    return text + (f" and {len(errors) - shown} more" if len(errors) > shown else "")


# ---------------------------------------------------------------------------
# The readings' texts
# ---------------------------------------------------------------------------

PRESENCE_ANSWER = "answer: is it in the image?"
COUNT_ANSWER = f"answer: how many are in the image, from 0 to {COUNT_MAX}; {COUNT_MORE} if more"
FIND = {"point": "find: the point", "box": "find: the box"}
AXIS_ANSWER = {
    ("point", "x"): "answer: x of the point, from 0 (left edge) to 255 (right edge)",
    ("point", "y"): "answer: y of the point, from 0 (top edge) to 255 (bottom edge)",
    ("box", "x"): "answer: x of the box centre, from 0 (left edge) to 255 (right edge)",
    ("box", "y"): "answer: y of the box centre, from 0 (top edge) to 255 (bottom edge)",
    ("box", "width"): "answer: width of the box, from 0 (none) to 255 (the whole width)",
    ("box", "height"): "answer: height of the box, from 0 (none) to 255 (the whole height)",
}


def image_line(spec: PlaceSpec) -> str:
    """The picture a reading looks in, in the words of the label in front of it (template.image_label)."""
    return f"image: Picture {int(spec.number)} ({spec.image})"


def presence_text(spec: PlaceSpec) -> str:
    """The presence reading's instructions (a yes/no question): the instructions, the picture, the question."""
    return f"{spec.instructions}\n{image_line(spec)}\n{PRESENCE_ANSWER}"


def count_text(spec: PlaceSpec) -> str:
    """The count reading's text: the presence reading's lines with another question (a point and a box count alike)."""
    return f"{spec.instructions}\n{image_line(spec)}\n{COUNT_ANSWER}"


def ordinal(k: int) -> str:
    """1st, 2nd, 3rd, 4th, ..., 11th, 12th, 13th, ..., 21st."""
    k = int(k)
    suffix = "th" if 10 <= k % 100 <= 20 else {1: "st", 2: "nd", 3: "rd"}.get(k % 10, "th")
    return f"{k}{suffix}"


def rank_line(rank: int, rank_of: int) -> str:
    """A listed place's line: `given: of 5, the 2nd from the left` (`of more than 16` for the level 17)."""
    of = f"more than {COUNT_MAX}" if int(rank_of) == COUNT_MORE else str(int(rank_of))
    return f"given: of {of}, the {ordinal(rank)} from the left"


def coordinate_text(spec: PlaceSpec, axis: str, given: dict, rank: int | None = None, rank_of: int | None = None) -> str:
    """A coordinate reading's text, line by line: the instructions as they are, the picture, what to find, a listed
    place's rank, the earlier readings as whole steps, the axis with its 0 and 255:

        the damaged package
        image: Picture 1 (belt)
        find: the box
        given: centre x 128, y 65
        answer: width of the box, from 0 (none) to 255 (the whole width)
    """
    lines = [spec.instructions, image_line(spec), FIND[spec.shape]]
    if rank is not None:
        lines.append(rank_line(rank, rank_of))
    if axis == "y":
        lines.append(f"given: x {given['x']}" if spec.shape == "point" else f"given: centre x {given['x']}")
    elif axis in ("width", "height"):
        lines.append(f"given: centre x {given['x']}, y {given['y']}")
    lines.append(AXIS_ANSWER[(spec.shape, axis)])
    return "\n".join(lines)


def presence_row(spec: PlaceSpec) -> dict:
    return {"type": "truth", "instructions": presence_text(spec)}


def count_row(spec: PlaceSpec) -> dict:
    return {"type": "coord", "instructions": count_text(spec), "classes": COUNT_LEVELS, "axis": "count"}


def coordinate_row(spec: PlaceSpec, axis: str, given: dict, rank: int | None = None, rank_of: int | None = None) -> dict:
    row = {"type": "coord", "instructions": coordinate_text(spec, axis, given, rank, rank_of), "classes": N_STEPS, "axis": axis}
    if rank is not None:
        row["rank"] = rank
    return row


# ---------------------------------------------------------------------------
# Readings
# ---------------------------------------------------------------------------


class ReadingError(ValueError):
    """Probabilities a reading cannot be read from (a fault of the model or the backend, never the caller's)."""


@dataclass(frozen=True)
class Reading:
    """One coordinate reading: the step (between steps), the most probable step and the peak's mass."""

    step: float
    argmax: int
    peak: float

    @property
    def fraction(self) -> float:
        return step_fraction(self.step)


@dataclass(frozen=True)
class CountReading:
    """The count reading: the level (0 to 16, 17 for more), `count` = min(level, max), `complete` = the probability
    of a level at or under max, and the level's own probability."""

    level: int
    count: int
    complete: float
    peak: float


def step_fraction(step: float) -> float:
    """The fraction of the side a (continuous) step stands for: (s + 0.5) / 256."""
    return (float(step) + 0.5) / N_STEPS


def given_step(step: float) -> int:
    """A reading as the whole step a later reading is given: halves up, within [0, 255]."""
    return min(N_STEPS - 1, max(0, int(math.floor(float(step) + 0.5))))


def _argmax_first(p: Sequence[float]) -> int:
    best = 0
    for k, x in enumerate(p):
        if x > p[best]:
            best = k
    return best


def check(probs: Sequence[float], n: int, what: str) -> None:
    """A list's readings are read only from n probabilities that are finite, not negative and sum to 1."""
    p = [float(x) for x in probs]
    if len(p) != n:
        raise ReadingError(f"a {what} reading has {n} probabilities, got {len(p)}")
    if any(not math.isfinite(x) or x < -1e-6 for x in p) or abs(math.fsum(p) - 1.0) > 1e-3:
        raise ReadingError(f"{what} probabilities must be finite, not negative, and sum to 1")


def read(probs: Sequence[float], window: int = WINDOW) -> Reading:
    """A coordinate reading from its 256 step probabilities (the module text)."""
    p = [max(0.0, float(x)) for x in probs]
    if len(p) != N_STEPS:
        raise ReadingError(f"a coordinate reading has {N_STEPS} probabilities, got {len(p)}")
    best = _argmax_first(p)
    lo, hi = max(0, best - window), min(N_STEPS - 1, best + window)
    mass = math.fsum(p[k] for k in range(lo, hi + 1))
    step = math.fsum(k * p[k] for k in range(lo, hi + 1)) / mass if mass > 0 else float(best)
    return Reading(step=step, argmax=best, peak=min(1.0, mass))


def read_count(probs: Sequence[float], max_places: int) -> CountReading:
    """The count reading from its 18 level probabilities for a list of at most `max_places` places."""
    check(probs, COUNT_LEVELS, "count")
    p = [float(x) for x in probs]
    best = _argmax_first(p)
    complete = min(1.0, math.fsum(max(0.0, x) for x in p[: max_places + 1]))
    return CountReading(level=best, count=min(best, max_places), complete=complete, peak=max(0.0, p[best]))


# ---------------------------------------------------------------------------
# Answers
# ---------------------------------------------------------------------------


def _half_up(x: float, decimals: int = 0) -> float:
    scale = 10.0**decimals
    return math.floor(x * scale + 0.5 + 1e-9) / scale


def to_pixels(fraction: float, size: int, *, inclusive_end: bool) -> int:
    """A fraction of a side of `size` pixels as a whole coordinate, halves up: a point names a pixel (0 to size - 1), a
    box corner may lie on the far edge (0 to size)."""
    value = int(_half_up(float(fraction) * int(size)))
    return min(int(size) if inclusive_end else int(size) - 1, max(0, value))


def to_normalized(fraction: float) -> float:
    """A fraction with 4 decimals, halves up, within [0, 1]."""
    return min(1.0, max(0.0, _half_up(float(fraction), 4)))


def box_corners(cx: float, cy: float, w: float, h: float) -> tuple[float, float, float, float]:
    """(x1, y1, x2, y2) as fractions from the centre and the size, clipped to the picture: x1 <= x2 and y1 <= y2."""
    w, h = max(0.0, float(w)), max(0.0, float(h))
    clip = lambda v: min(1.0, max(0.0, v))  # noqa: E731
    return clip(cx - w / 2.0), clip(cy - h / 2.0), clip(cx + w / 2.0), clip(cy + h / 2.0)


def box_iou(a: Sequence[float], b: Sequence[float]) -> float:
    """The intersection over union of two boxes (x1, y1, x2, y2) in fractions of the sides (0 when both are empty)."""
    w = max(0.0, min(a[2], b[2]) - max(a[0], b[0]))
    h = max(0.0, min(a[3], b[3]) - max(a[1], b[1]))
    inter = w * h
    union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter / union if union > 0 else 0.0


def merge_boxes(places: Sequence[tuple[tuple[float, ...], float]], threshold: float = MERGE_IOU) -> list[int]:
    """The boxes (fractions, confidence) kept when two on one thing are one: taken by confidence, highest first (ties:
    the earlier in the list), a box that overlaps a box already kept at `threshold` or more is dropped. Returns the
    indexes kept, in list order."""
    kept: list[int] = []
    for i in sorted(range(len(places)), key=lambda i: (-places[i][1], i)):
        worst = max((box_iou(places[i][0], places[j][0]) for j in kept), default=0.0)
        if worst < threshold:
            kept.append(i)
    return sorted(kept)


def place_order(fractions: Sequence[float]) -> tuple[float, float]:
    """A listed place's key: from the left by the centre's x, ties from the top (a box's centre, or the point), as
    fractions rounded to 6 digits."""
    if len(fractions) == 2:
        cx, cy = fractions
    else:
        cx, cy = (fractions[0] + fractions[2]) / 2.0, (fractions[1] + fractions[3]) / 2.0
    return round(float(cx), 6), round(float(cy), 6)


def _place(shape: str, units: str, f: Sequence[float], width: int, height: int) -> dict:
    """A point's x and y, or a box's corners, in the question's units."""
    if shape == "point":
        if units == "pixels":
            return {"x": to_pixels(f[0], width, inclusive_end=False), "y": to_pixels(f[1], height, inclusive_end=False)}
        return {"x": to_normalized(f[0]), "y": to_normalized(f[1])}
    if units == "pixels":
        x1, y1, x2, y2 = (to_pixels(v, side, inclusive_end=True) for v, side in zip(f, (width, height, width, height)))
    else:
        x1, y1, x2, y2 = (to_normalized(v) for v in f)
    return {"x1": x1, "y1": y1, "x2": x2, "y2": y2}


# ---------------------------------------------------------------------------
# The passes of a request
# ---------------------------------------------------------------------------


@dataclass
class Chain:
    """One place: the presence and the coordinate readings so far."""

    spec: PlaceSpec
    present: float | None = None
    readings: dict = field(default_factory=dict)  # axis -> Reading

    @property
    def passes(self) -> int:
        return 1 + max(at for _, at in CHAIN[self.spec.shape])

    def answer(self, width: int, height: int) -> dict:
        spec = self.spec
        confidence = 1.0
        for axis, _ in CHAIN[spec.shape]:
            confidence *= self.readings[axis].peak
        r = self.readings
        f = (r["x"].fraction, r["y"].fraction) if spec.shape == "point" else box_corners(r["x"].fraction, r["y"].fraction, r["width"].fraction, r["height"].fraction)
        return {"type": spec.shape, **_place(spec.shape, spec.units, f, width, height), "present": round2(self.present), "confidence": round2(confidence),
                "image": spec.image, "width": int(width), "height": int(height)}  # fmt: skip


@dataclass
class ListChain:
    """Several places (`max`): the presence, the count, and per listed place (rank 1 first) its readings so far."""

    spec: PlaceSpec
    present: float | None = None
    counting: CountReading | None = None
    readings: list = field(default_factory=list)  # per place: {axis: Reading}
    passes_done: int = 0

    @property
    def passes(self) -> int:
        """1 for a count of 0, else 1 plus a single place's passes (a point 3, a box 4); before pass 0 the most."""
        if self.counting is not None and self.counting.count == 0:
            return 1
        return 2 + max(at for _, at in CHAIN[self.spec.shape])

    @property
    def done(self) -> bool:
        return self.passes_done >= self.passes

    def rows(self) -> list[dict]:
        """The next pass's readings: pass 0 the presence and the count, then place by place one pass of each chain."""
        spec = self.spec
        if self.passes_done == 0:
            return [presence_row(spec), count_row(spec)]
        at, level = self.passes_done - 1, self.counting.level
        out = []
        for k in range(1, self.counting.count + 1):
            for axis, when in CHAIN[spec.shape]:
                if when == at:
                    given = {name: given_step(self.readings[k - 1][name].step) for name in GIVEN[axis]}
                    out.append(coordinate_row(spec, axis, given, rank=k, rank_of=level))
        return out

    def take(self, rows: Sequence[dict], probs: Sequence[Sequence[float]]) -> None:
        if self.passes_done == 0:
            if len(probs[0]) != 2:
                raise ReadingError("the presence reading has 2 probabilities")
            self.present = float(probs[0][0])
            self.counting = read_count(probs[1], self.spec.max)
            self.readings = [{} for _ in range(self.counting.count)]
        else:
            for row, p in zip(rows, probs):
                check(p, N_STEPS, "coordinate")
                self.readings[row["rank"] - 1][row["axis"]] = read(p)
        self.passes_done += 1

    def answer(self, width: int, height: int) -> dict:
        spec = self.spec
        places = []
        for r in self.readings:
            confidence = 1.0
            for axis, _ in CHAIN[spec.shape]:
                confidence *= r[axis].peak
            f = (r["x"].fraction, r["y"].fraction) if spec.shape == "point" else box_corners(r["x"].fraction, r["y"].fraction, r["width"].fraction, r["height"].fraction)
            places.append((f, confidence))
        kept = merge_boxes(places) if spec.shape == "box" else list(range(len(places)))
        listed = [{**_place(spec.shape, spec.units, places[i][0], width, height), "confidence": round2(places[i][1])} for i in sorted(kept, key=lambda i: (place_order(places[i][0]), i))]
        return {"type": spec.shape, "count": len(listed), "present": round2(self.present), "complete": round2(self.counting.complete),
                "boxes" if spec.shape == "box" else "points": listed, "image": spec.image, "width": int(width), "height": int(height)}  # fmt: skip


class Plan:
    """The point and box questions of one request, answered pass by pass: `rows()` gives the next pass's readings (each
    with the index of its question), `take` their probabilities in that order, until `rows()` is empty; then `answers`.
    Pass 0 holds every single place's presence reading first, then per question its first readings."""

    def __init__(self, specs: Sequence[PlaceSpec]):
        self.chains: list = [ListChain(spec) if spec.max is not None else Chain(spec) for spec in specs]
        self.passes_run = 0

    @property
    def passes(self) -> int:
        return max((chain.passes for chain in self.chains), default=0)

    def rows(self) -> list[tuple[int, dict]]:
        at = self.passes_run
        if at >= self.passes:
            return []
        out = [(n, presence_row(chain.spec)) for n, chain in enumerate(self.chains) if isinstance(chain, Chain)] if at == 0 else []
        for n, chain in enumerate(self.chains):
            if isinstance(chain, ListChain):
                if not chain.done and chain.passes_done == at:
                    out += [(n, row) for row in chain.rows()]
                continue
            for axis, when in CHAIN[chain.spec.shape]:
                if when == at:
                    given = {name: given_step(chain.readings[name].step) for name in GIVEN[axis]}
                    out.append((n, coordinate_row(chain.spec, axis, given)))
        return out

    def take(self, rows: Sequence[tuple[int, dict]], probs: Sequence[Sequence[float]]) -> None:
        if len(rows) != len(probs):
            raise ReadingError(f"expected {len(rows)} probability lists, got {len(probs)}")
        lists: dict[int, tuple[list, list]] = {}
        for (n, row), p in zip(rows, probs):
            chain = self.chains[n]
            if isinstance(chain, ListChain):
                taken = lists.setdefault(n, ([], []))
                taken[0].append(row)
                taken[1].append([float(x) for x in p])
            elif row["type"] == "truth":
                chain.present = float(p[0])
            else:
                chain.readings[row["axis"]] = read(p)
        for n, (rows_of, probs_of) in lists.items():
            self.chains[n].take(rows_of, probs_of)
        self.passes_run += 1

    def answers(self, sizes: dict[str, tuple[int, int]]) -> dict[str, dict]:
        """Every question's answer by id; `sizes` holds each picture's width and height as sent."""
        return {chain.spec.qid: chain.answer(*sizes[chain.spec.image]) for chain in self.chains}