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

Emits an animated frame, a depth map of the same scene, the box around the square
moving in it, and a caption -- every input dora-image-view draws -- so the example
dataflow needs no camera, no detector and no network.

The depth map describes the same scene as the image: a flat back wall with the square
standing in front of it. `camera_info` is published once, on the first tick, the way a
real camera publishes it once per session -- which is exactly the case the viewer has
to cope with by reading intrinsics from frame metadata as well.
"""

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"

WALL_MM = 2000.0
SQUARE_MM = 800.0

# A plausible pinhole for a frame this size.
FX = FY = float(WIDTH)
CX, CY = WIDTH / 2, HEIGHT / 2
DISTORTION = [0.0, 0.0, 0.0, 0.0, 0.0]
# Attached to every frame as well as published once, which is what dora-zed does
# and what lets a viewer that started late still learn the camera model.
INTRINSICS_METADATA = {"fx": FX, "fy": FY, "cx": CX, "cy": CY, "distortion": DISTORTION}


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

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

        if sent == 0:
            node.send_output(
                "camera_info",
                pa.array(np.array([FX, FY, CX, CY], dtype=np.float32)),
                {"shape": [4], "dtype": "float32", "layout": "C", "distortion": DISTORTION},
            )

        frame, corners = render(sent)
        node.send_output(
            "image",
            pa.array(frame.ravel()),
            {"width": WIDTH, "height": HEIGHT, "encoding": "bgr8", **INTRINSICS_METADATA},
        )
        node.send_output(
            "bbox",
            pa.array([{
                "bbox": corners,
                "conf": np.array([0.99]),
                "labels": np.array(["square"]),
            }]),
            {"format": "xyxy"},
        )
        depth = render_depth(corners)
        node.send_output(
            "depth",
            pa.array(depth.ravel()),
            {
                "shape": [HEIGHT, WIDTH],
                "dtype": "float32",
                "layout": "HW",
                "units": "mm",
                **INTRINSICS_METADATA,
            },
        )
        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


def render_depth(corners: np.ndarray) -> np.ndarray:
    """Depth of the scene in `render`: a flat wall with the square in front of it."""
    depth = np.full((HEIGHT, WIDTH), WALL_MM, dtype=np.float32)
    left, top, right, bottom = (int(value) for value in corners)
    depth[top:bottom, left:right] = SQUARE_MM
    return depth


if __name__ == "__main__":
    main()