To show something interesting, we first need to log some data.

from __future__ import annotations

import math

import ipywidgets as widgets
import numpy as np
from IPython.display import display

import rerun as rr
import rerun.blueprint as rrb
from rerun.notebook import Viewer, ViewerEvent
from rerun.utilities import build_color_grid

# To show something interesting, we first need to log some data.
rr.init("rerun_example_callbacks")

# This will output a cube grid of points. We'll be able to select
# either the entire grid, or individual points.
STEPS = 100
twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4
for t in range(STEPS):
    rr.set_time("step", sequence=t)
    cube = build_color_grid(10, 10, 10, twist=twists[t])
    rr.log("cube", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))

rr.send_blueprint(rrb.Blueprint(rrb.Spatial3DView(name="Everything")))


# Spawn a viewer, and create some labels to display the values we're tracking
viewer = Viewer()
labels = widgets.VBox([
    widgets.Label(),  # active timeline
    widgets.Label(),  # current time
    widgets.Label(),  # viewer selection
])


def on_viewer_event(event: ViewerEvent) -> None:
    if event.type == "timeline_change":
        labels.children[0].value = f"active timeline: {event.timeline}"
        labels.children[1].value = f"current time: {event.time}"

    elif event.type == "time_update":
        labels.children[1].value = f"current time: {event.time}"

    elif event.type == "selection_change":
        # Because it is possible to select multiple items at a time in the Viewer,
        # the selection change event yields a list of items.
        items = []
        for item in event.items:
            if item.type == "entity":
                # Entities always include their entity path, and optionally also:
                # * The instance ID
                # * The coordinates at which a selection occurred for 2D or 3D views.
                #    Coordinates are within the view's space and thus relative to its origin.
                # * The view name
                entity_path, instance_id, view_name, position = (
                    item.entity_path,
                    item.instance_id,
                    item.view_name,
                    item.position,
                )
                items.append(f"Entity({entity_path=}, {instance_id=}, {view_name=}, {position=})")
            elif item.type == "container":
                # Containers and Views include their ID and name.
                container_id, container_name = item.container_id, item.container_name
                items.append(f"Container({container_id=}, {container_name=})")
            elif item.type == "view":
                view_id, view_name = item.view_id, item.view_name
                items.append(f"View({view_id=}, {view_name=})")
        labels.children[2].value = f"selected=[{' '.join(items)}]"


viewer.on_event(on_viewer_event)

display(viewer)
display(labels)