import os

from time import perf_counter

import numpy as np

from imgui_bundle import imgui, imgui_ctx, icons_fontawesome_6 as fa

from ...graphics import (

ScatterCollection,

ScatterStack,

LineCollection,

LineStack,

ImageGraphic,

ImageVolumeGraphic,

)

from ...utils import quick_min_max

from ...layouts import Subplot

from ...ui import ImguiWindow, StandardRightClickMenu

from ._index import RangeContinuous

from ._base import NDGraphic

from ._nd_positions import NDPositions, NDTimeseries

from ._nd_image import NDImage

position_graphic_types = [ScatterCollection, ScatterStack, LineCollection, LineStack]

class NDWidgetUI(ImguiWindow):

"""Playback controls and a slider for each slider dim, shown at the bottom of an ``NDWidget``"""

def __init__(self, ndwidget):

super().__init__()

self._ndwidget = ndwidget

# whether or not a dimension is in play mode

self._playing = dict()

# approximate framerate for playing

self._fps = dict()

# framerate converted to frame time

self._frame_time = dict()

# last timepoint that a frame was displayed from a given dimension

self._last_frame_time = dict()

# loop playback

self._loop = dict()

# last time the slider was moved per dim, used for time-based throttling

self._last_slider_movement: dict[str, float] = dict()

for dim in self._ndwidget.ranges:

self.push_dim(dim)

self._max_display_windows: dict[NDGraphic, float | int] = dict()

def push_dim(self, dim):

"""initialize the playback & slider UI state for a newly added dim"""

self._playing[dim] = False

self._fps[dim] = 20

self._frame_time[dim] = 1 / 20

self._last_frame_time[dim] = perf_counter()

self._loop[dim] = False

self._last_slider_movement[dim] = 0.0

# auto-plays first inserted dim in docs galleries

if "DOCS_BUILD" in os.environ.keys():

if os.environ["DOCS_BUILD"] == "1":

dim = list(self._ndwidget.indices.ref_ranges.keys())[0]

self._playing[dim] = True

self._loop[dim] = True

# a gallery movie is a fixed number of rendered frames, not a span of real time,

# so step once per frame. otherwise the framerate gate below holds the same index

# for however many frames the machine can render within 1 / fps seconds

self._frame_time[dim] = 0

def pop_dim(self, dim):

"""remove the playback & slider UI state for a removed dim"""

self._playing.pop(dim)

self._fps.pop(dim)

self._frame_time.pop(dim)

self._last_frame_time.pop(dim)

self._loop.pop(dim)

self._last_slider_movement.pop(dim)

def _set_index(self, dim, index):

if index >= self._ndwidget.ranges[dim].stop:

if self._loop[dim]:

index = self._ndwidget.ranges[dim].start

else:

index = self._ndwidget.ranges[dim].stop

self._playing[dim] = False

self._ndwidget.indices.set_dim_index(dim, index)

def update(self):

if len(self._ndwidget.indices) < 1 and len(self._update_calls) < 2:

# there are no slider dims AND there are no appended UI elements to this window

self.size = 0

self.collapsed = True

return

now = perf_counter()

for dim, current_index in self._ndwidget.indices:

# push id since we have the same buttons for each dim

imgui.push_id(f"{self._id_counter}_{dim}")

rr = self._ndwidget.ranges[dim]

if self._playing[dim]:

# show pause button if playing

if imgui.button(label=fa.ICON_FA_PAUSE):

# if pause button clicked, then set playing to false

self._playing[dim] = False

# if in play mode and enough time has elapsed w.r.t. the desired framerate, increment the index

if now - self._last_frame_time[dim] >= self._frame_time[dim]:

self._set_index(dim, current_index + rr.step)

self._last_frame_time[dim] = now

else:

# we are not playing, so display play button

if imgui.button(label=fa.ICON_FA_PLAY):

# if play button is clicked, set last frame time to 0 so that index increments on next render

self._last_frame_time[dim] = 0

# set playing to True since play button was clicked

self._playing[dim] = True

imgui.same_line()

# step back one frame button

if imgui.button(label=fa.ICON_FA_BACKWARD_STEP) and not self._playing[dim]:

self._set_index(dim, current_index - rr.step)

imgui.same_line()

# step forward one frame button

if imgui.button(label=fa.ICON_FA_FORWARD_STEP) and not self._playing[dim]:

self._set_index(dim, current_index + rr.step)

imgui.same_line()

# stop button

if imgui.button(label=fa.ICON_FA_STOP):

self._playing[dim] = False

self._last_frame_time[dim] = 0

self._ndwidget.indices.set_dim_index(dim, rr.start)

imgui.same_line()

# loop checkbox

_, self._loop[dim] = imgui.checkbox(

label=fa.ICON_FA_ROTATE, v=self._loop[dim]

)

if imgui.is_item_hovered(0):

imgui.set_tooltip("loop playback")

imgui.same_line()

imgui.text("framerate :")

imgui.same_line()

imgui.set_next_item_width(100)

# framerate int entry

fps_changed, value = imgui.input_int(

label="fps", v=self._fps[dim], step_fast=5

)

if imgui.is_item_hovered(0):

imgui.set_tooltip(

"framerate is approximate and less reliable as it approaches your monitor refresh rate"

)

if fps_changed:

if value < 1:

value = 1

if value > 100:

value = 100

self._fps[dim] = value

self._frame_time[dim] = 1 / value

imgui.text(str(dim))

imgui.same_line()

# so that slider occupies full width

imgui.set_next_item_width(self.width * 0.85)

if isinstance(rr, RangeContinuous):

changed, new_index = imgui.slider_float(

v=current_index,

v_min=rr.start,

v_max=rr.stop - rr.step,

label=f"##{dim}",

)

if changed:

if now - self._last_slider_movement[dim] > rr.throttle:

self._ndwidget.indices.set_dim_index(

dim, new_index, cancel_awaiting=True

)

self._last_slider_movement[dim] = now

elif imgui.is_item_hovered():

if imgui.is_key_pressed(imgui.Key.right_arrow):

self._set_index(dim, current_index + rr.step)

elif imgui.is_key_pressed(imgui.Key.left_arrow):

self._set_index(dim, current_index - rr.step)

imgui.pop_id()

# auto set imgui window height

if not self._collapsed:

height = round(

imgui.get_cursor_screen_pos().y

- self.y

+ imgui.get_style().window_padding.y

)

if height != self.size:

self.size = height

class RightClickMenu(StandardRightClickMenu):

"""

Right click menu of an ``NDWidget``, adds an "ND Graphics" submenu to the standard menu. Selecting an

``NDGraphic`` opens a window to change its settings, which stays open after the menu closes.

"""

def __init__(self, ndwidget):

super().__init__()

self._ndwidget = ndwidget

self._ndgraphic_windows = set()

def update(self):

super().update()

if imgui.begin_menu("ND Graphics"):

for ndg in self._ndwidget[self.subplot].nd_graphics:

name = ndg.name if ndg.name is not None else hex(id(ndg))

if imgui.menu_item(f"{name}", "", False)[0]:

self._ndgraphic_windows.add(ndg)

imgui.end_menu()

def draw(self):

super().draw()

# the ND graphic windows are not part of the popup, they stay open after the popup closes

for ndg in list(

self._ndgraphic_windows

): # set -> list so we can change size during iteration

name = ndg.name if ndg.name is not None else hex(id(ndg))

subplot = ndg.graphic._plot_area

imgui.set_next_window_size((0, 0))

_, open = imgui.begin(f"subplot: {subplot.name}, {name}", True)

draw_nd_graphic_ui(ndg)

if not open:

self._ndgraphic_windows.remove(ndg)

imgui.end()

def draw_nd_graphic_ui(nd_graphic: NDGraphic):

"""

Draw the settings of an ``NDGraphic``.

Used both by the popup that a right-click on the graphic opens, and by the window that the

"ND Graphics" submenu of the ``NDWidget`` right-click menu opens. A right-click can be hard to

land on a thin graphic, ex: a line or a scatter, which is why both exist.

"""

if isinstance(nd_graphic, NDTimeseries):

_draw_nd_timeseries_ui(nd_graphic)

elif isinstance(nd_graphic, NDPositions):

_draw_nd_positions_ui(nd_graphic)

elif isinstance(nd_graphic, NDImage):

_draw_nd_image_ui(nd_graphic)

_, nd_graphic.pause = imgui.checkbox("pause", nd_graphic.pause)

def _draw_magnitude_sliders(

label: str, value: float, exponent_min: int, exponent_max: int

) -> tuple[bool, float]:

"""

A value slider paired with an order of magnitude slider, ``value * 10 ** exponent``.

The pair covers many decades without any one slider having to. The value is decomposed on every

frame rather than kept as UI state, so the sliders follow the property when something else moves

it, ex: a camera zoom writing ``display_window``.

"""

if value > 0:

exponent = int(np.floor(np.log10(value)))

exponent = min(max(exponent, exponent_min), exponent_max)

else:

# log10 is undefined at 0, which is a valid display window: only the current datapoint

exponent = exponent_min

changed_value, value = imgui.slider_float(

label, v=value / 10.0**exponent, v_min=0.0, v_max=10.0

)

changed_exponent, exponent = imgui.slider_int(

f"{label} 10^", v=exponent, v_min=exponent_min, v_max=exponent_max

)

return (changed_value or changed_exponent), value * 10.0**exponent

def _draw_nd_image_ui(nd_image: NDImage):

if nd_image.graphic.data.value is not None:

# if it doesn't have a CPU buffer the value is None

# i.e. data is only on the GPU, e.g. YUV

_min, _max = quick_min_max(nd_image.graphic.data.value)

changed, vmin = imgui.slider_float(

"vmin", nd_image.graphic.vmin, v_min=_min, v_max=_max

)

if changed:

nd_image.graphic.vmin = vmin

changed, vmax = imgui.slider_float(

"vmax", nd_image.graphic.vmax, v_min=_min, v_max=_max

)

if changed:

nd_image.graphic.vmax = vmax

changed, new_gamma = imgui.slider_float(

"gamma", nd_image.graphic._material.gamma, 0.01, 5

)

if changed:

nd_image.graphic._material.gamma = new_gamma

def _draw_nd_positions_ui(

nd_graphic: NDPositions, graphic_types: list = position_graphic_types

):

subplot = nd_graphic._nd_subplot.subplot

ndwidget = nd_graphic._nd_subplot.ndw

for cls in graphic_types:

if imgui.radio_button(cls.__name__, type(nd_graphic.graphic) is cls):

nd_graphic.graphic_type = cls

subplot.auto_scale()

changed, val = imgui.checkbox(

"use display_window", nd_graphic.display_window is not None

)

p_dim = nd_graphic.slicer.display_dims[1]

if changed:

if not val:

nd_graphic.display_window = None

else:

# pick a value 10% of the reference range

nd_graphic.display_window = ndwidget.ranges[p_dim].size * 0.1

if nd_graphic.display_window is not None:

changed, new = _draw_magnitude_sliders(

"display_window", nd_graphic.display_window, -10, 8

)

if changed:

nd_graphic.display_window = new

changed, limit = imgui.checkbox(

"use max_display_datapoints", nd_graphic.max_display_datapoints is not None

)

if changed:

nd_graphic.max_display_datapoints = 1_000 if limit else None

if nd_graphic.max_display_datapoints is not None:

changed, new = _draw_magnitude_sliders(

"max_display_datapoints", nd_graphic.max_display_datapoints, 0, 8

)

if changed:

# the sliders can reach 0, the minimum is 2

nd_graphic.max_display_datapoints = max(2, int(new))

def _draw_nd_timeseries_ui(nd_graphic: NDTimeseries):

# a heatmap only makes sense for timeseries data

_draw_nd_positions_ui(nd_graphic, position_graphic_types + [ImageGraphic])

options = [None, "fixed", "auto"]

changed, option = imgui.combo(

"x-range mode",

options.index(nd_graphic.x_range_mode),

[str(o) for o in options],

)

if changed:

nd_graphic.x_range_mode = options[option]