"""Synthetic source node: this repository's stand-in for a camera and a detector.

Emits an animated frame, the box around the square moving in it, and a caption --
the three inputs dora-image-view draws -- so the example dataflow needs no camera, no
detector and no network.
"""

import os

import numpy as np
import pyarrow as pa
from dora import Node

WIDTH = int(os.getenv("SOURCE_WIDTH", "640"))
HEIGHT = int(os.getenv("SOURCE_HEIGHT", "480"))
# 0 runs until the viewer quits; a positive count is what lets CI run the
# dataflow to completion instead of killing it on a timeout.
FRAMES = int(os.getenv("SOURCE_FRAMES", "0"))

BOX = 96
SPEED = 6
CAPTION = "dora-image-view synthetic source"


def main() -> None:
    """Send one image, bbox and text per tick."""
    node = Node("source")
    sent = 0

    for event in node:
        if event["type"] != "INPUT":
            continue

        frame, corners = render(sent)
        node.send_output(
            "image",
            pa.array(frame.ravel()),
            {"width": WIDTH, "height": HEIGHT, "encoding": "bgr8"},
        )
        node.send_output(
            "bbox",
            pa.array([{
                "bbox": corners,
                "conf": np.array([0.99]),
                "labels": np.array(["square"]),
            }]),
            {"format": "xyxy"},
        )
        node.send_output("text", pa.array([CAPTION]))

        sent += 1
        if FRAMES and sent >= FRAMES:
            break


def render(index: int) -> tuple[np.ndarray, np.ndarray]:
    """Draw frame `index` and return it with the xyxy corners of the square in it."""
    frame = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8)
    frame[:, :, 0] = np.linspace(40, 140, HEIGHT, dtype=np.uint8)[:, None]

    travel = WIDTH - BOX
    # A triangle wave over `index`: the square starts at the left edge, runs to
    # the right edge, and comes back.
    left = travel - abs((index * SPEED) % (2 * travel) - travel)
    top = (HEIGHT - BOX) // 2
    frame[top : top + BOX, left : left + BOX] = (60, 200, 255)

    corners = np.array([left, top, left + BOX, top + BOX], dtype=np.float64)
    return frame, corners


if __name__ == "__main__":
    main()