import math

import numpy as np

import pygfx

from pylinalg import quat_from_vecs, vec_transform_quat

from ..utils.enums import RenderQueue

from ..utils import global_config

GRID_PLANES = ["xy", "xz", "yz"]

CANONICAL_BAIS = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])

# very thin subclass that just adds GridMaterial properties to this world object for easier user control

class Grid(pygfx.Grid):

@property

def major_step(self) -> tuple[float, float]:

"""The step distance between the major grid lines."""

return self.material.major_step

@major_step.setter

def major_step(self, step: tuple[float, float]):

self.material.major_step = step

@property

def minor_step(self) -> tuple[float, float]:

"""The step distance between the minor grid lines."""

return self.material.minor_step

@minor_step.setter

def minor_step(self, step: tuple[float, float]):

self.material.minor_step = step

@property

def axis_thickness(self) -> float:

"""The thickness of the axis lines."""

return self.material.axis_thickness

@axis_thickness.setter

def axis_thickness(self, thickness: float):

self.material.axis_thickness = thickness

@property

def major_thickness(self) -> float:

"""The thickness of the major grid lines."""

return self.material.major_thickness

@major_thickness.setter

def major_thickness(self, thickness: float):

self.material.major_thickness = thickness

@property

def minor_thickness(self) -> float:

"""The thickness of the minor grid lines."""

return self.material.minor_thickness

@minor_thickness.setter

def minor_thickness(self, thickness: float):

self.material.minor_thickness = thickness

@property

def thickness_space(self) -> str:

"""The coordinate space in which the thicknesses are expressed.

See :obj:`pygfx.utils.enums.CoordSpace`:

"""

return self.material.thickness_space

@thickness_space.setter

def thickness_space(self, value: str):

self.material.thickness_space = value

@property

def axis_color(self) -> str:

"""The color of the axis lines."""

return self.material.axis_color

@axis_color.setter

def axis_color(self, color: str):

self.material.axis_color = color

@property

def major_color(self) -> str:

"""The color of the major grid lines."""

return self.material.major_color

@major_color.setter

def major_color(self, color: str):

self.material.major_color = color

@property

def minor_color(self) -> str:

"""The color of the minor grid lines."""

return self.material.minor_color

@minor_color.setter

def minor_color(self, color: str):

self.material.minor_color = color

@property

def infinite(self) -> bool:

"""Whether the grid is infinite.

If not infinite, the grid is 1x1 in world space, scaled, rotated, and

positioned with the object's transform.

(Infinite grids are not actually infinite. Rather they move along with

the camera, and are sized based on the distance between the camera and

the grid.)

"""

return self.material.infinite

@infinite.setter

def infinite(self, value: str):

self.material.infinite = value

class Grids(pygfx.Group):

"""Just a class to make accessing the grids easier"""

def __init__(self, *, xy, xz, yz):

super().__init__()

self._xy = xy

self._xz = xz

self._yz = yz

self.add(xy, xz, yz)

@property

def xy(self) -> Grid:

"""xy grid"""

return self._xy

@property

def xz(self) -> Grid:

"""xz grid"""

return self._xz

@property

def yz(self) -> Grid:

"""yz grid"""

return self._yz

class Ruler(pygfx.Ruler):

"""pygfx.Ruler subclass that adds a rotated axis label."""

def __init__(self, *, color="#fff", alpha_mode=None, render_queue=None, **kwargs):

super().__init__(

color=color, alpha_mode=alpha_mode, render_queue=render_queue, **kwargs

)

self._label = pygfx.Text(

screen_space=True,

anchor="middle-center",

font_size=20,

material=pygfx.TextMaterial(

color=color,

alpha_mode="auto",

render_queue=RenderQueue.overlay + 50,

aa=True,

),

)

self._label.visible = False

self.add(self._label)

self.text.material.outline_thickness = 0.5

@property

def label(self) -> pygfx.Text:

"""Axis label. Set text via ``label.set_text('label text')``"""

return self._label

@property

def color(self):

return self._text.material.color

@color.setter

def color(self, color):

self._text.material.color = color

self._line.material.color = color

self._points.material.edge_color = color

self._label.material.color = color

def update(self, camera, canvas_size):

stats = super().update(camera, canvas_size)

self._update_label()

return stats

def _update_label(self):

# update the label position

t1, t2 = self._visible_part_coords

if t1 == t2:

self._label.visible = False

return

self._label.visible = True

mid_t = 0.5 * (t1 + t2)

self._label.local.position = (

self._start_pos * (1 - mid_t) + self._end_pos * mid_t

)

vec = self._visible_part_screen_vec

angle = math.atan2(vec[1], vec[0])

# the side of the line that the tick labels are on, as a screen space unit vector. this is

# the same rule that pygfx uses to anchor the tick labels themselves, and the screen vector

# already carries the camera scale, the viewport aspect and the ruler's orientation, so

# none of those need a case of their own.

if self.tick_side == "left":

px, py = -math.sin(angle), math.cos(angle)

else:

px, py = math.sin(angle), -math.cos(angle)

# a ruler that runs right to left, or top to bottom, on screen would render the label

# upside down, so turn it around. that turns the label's own axes around with it

upside_down = not (-0.5 * math.pi < angle <= 0.5 * math.pi)

if upside_down:

angle -= math.copysign(math.pi, angle)

# pylinalg uses [x, y, z, w] quaternion format

self._label.local.rotation = np.array(

[0.0, 0.0, math.sin(angle / 2), math.cos(angle / 2)]

)

# the label is rotated onto the line, so in its own frame the line runs along x and the

# ticks sit on one side of it, +y or -y. anchoring it to that side offsets it in screen

# pixels, which is what keeps the camera scale and the viewport aspect out of the placement

if (self.tick_side == "left") != upside_down:

anchor = "bottom-center"

else:

anchor = "top-center"

# max extent of the tick labels in that same perpendicular direction.

# tick labels are unrotated screen-space text, so we project their

# axis-aligned _rect onto (px, py) directly.

px_pos, px_neg = max(px, 0), min(px, 0)

py_pos, py_neg = max(py, 0), min(py, 0)

tick_extent_px = max(

(

px_pos * b._rect.right

+ px_neg * b._rect.left

+ py_pos * b._rect.top

+ py_neg * b._rect.bottom

for b in self.text._text_blocks

if b._rect.width > 0 or b._rect.height > 0

),

default=0.0,

)

# gap between the tick labels and the label. a text rect is tight on its left and right,

# but its top and bottom are the font's ascender and descender, which neither a tick

# number nor most labels reach. that padding already separates the two where the offset is

# vertical, so only add a gap to the extent that the offset is horizontal

gap_px = abs(px) * 0.5 * self._label.font_size

anchor_offset = max(tick_extent_px, 0.0) + gap_px

# both of these re-run the text layout, so only set them when they actually change

if self._label._anchor != anchor:

self._label.anchor = anchor

if self._label._anchor_offset != anchor_offset:

self._label.anchor_offset = anchor_offset

@global_config.register

class Axes:

config = global_config.descriptor

@global_config.declare(

"intersection",

"tick_size",

"line_width",

"tick_marker",

"color",

"grids",

"grid_kwargs",

"auto_grid",

)

def __init__(

self,

plot_area,

intersection: tuple[int, int, int] | None = None,

tick_size: float = 8.0,

line_width: float = 2.0,

tick_marker: str = "tick",

color: str = "#fff",

x_kwargs: dict = None,

y_kwargs: dict = None,

z_kwargs: dict = None,

grids: bool = True,

grid_kwargs: dict = None,

auto_grid: bool = True,

offset: np.ndarray = np.array([0.0, 0.0, 0.0]),

basis: np.ndarray = np.array(

[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]

),

):

self._plot_area = plot_area

x_kwargs = x_kwargs or {}

y_kwargs = y_kwargs or {}

z_kwargs = z_kwargs or {}

generic_kwargs = dict(

tick_size=tick_size,

line_width=line_width,

tick_marker=tick_marker, # 'tick' for both-sides, 'tick_left' or 'tick_right' for one-sided

color=color,

)

x_kwargs = dict(

tick_side="right",

**generic_kwargs,

**x_kwargs,

)

y_kwargs = dict(

tick_side="left",

**generic_kwargs,

**y_kwargs,

)

z_kwargs = dict(

tick_side="left",

**generic_kwargs,

**z_kwargs,

)

# create ruler for each dim

self._x = Ruler(alpha_mode="solid", render_queue=RenderQueue.axes, **x_kwargs)

self._y = Ruler(alpha_mode="solid", render_queue=RenderQueue.axes, **y_kwargs)

self._z = Ruler(alpha_mode="solid", render_queue=RenderQueue.axes, **z_kwargs)

# We render the lines and ticks as solid, but enable aa for text for prettier glyphs

for ruler in self._x, self._y, self._z:

ruler.line.material.depth_compare = "<="

ruler.points.material.depth_compare = "<="

ruler.text.material.depth_compare = "<="

ruler.text.material.alpha_mode = "auto"

ruler.text.material.aa = True

ruler.label.material.depth_compare = "<="

self._offset = offset

# *MUST* instantiate some start and end positions for the rulers else kernel crashes immediately

# probably a WGPU rust panic

self.x.start_pos = 0, 0, 0

self.x.end_pos = 100, 0, 0

self.x.start_value = self.x.start_pos[0] - offset[0]

statsx = self.x.update(

self._plot_area.camera, self._plot_area.viewport.logical_size

)

self.y.start_pos = 0, 0, 0

self.y.end_pos = 0, 100, 0

self.y.start_value = self.y.start_pos[1] - offset[1]

statsy = self.y.update(

self._plot_area.camera, self._plot_area.viewport.logical_size

)

self.z.start_pos = 0, 0, 0

self.z.end_pos = 0, 0, 100

self.z.start_value = self.z.start_pos[1] - offset[2]

self.z.update(self._plot_area.camera, self._plot_area.viewport.logical_size)

# world object for the rulers + grids

self._world_object = pygfx.Group()

# add rulers

self.world_object.add(

self.x,

self.y,

self.z,

)

# set z ruler invisible for orthographic projections for now

if self._plot_area.camera.fov == 0:

# TODO: allow any orientation in the future even for orthographic projections

self.z.visible = False

if grid_kwargs is None:

grid_kwargs = dict()

# The grid is a bit weird, because it makes use of transparency to fade off in the distance.

# But w want it to write depth, so that objects that are drawn behind it are partually hidden.

# So we set alha_mode to 'auto'. We make it draw earlier than other 'auto' objects, under the

# assumption that most interesting stuff is in front of the grid, and artifacts behind the grid are less

# bad than those in front. Note that fully opaque objects blend perfectly fine with the grid. Artifacts

# should only emerge for objects that have semi-transparent fragments.

grid_kwargs = dict(

alpha_mode="auto",

render_queue=RenderQueue.auto + 50,

major_step=10,

minor_step=1,

thickness_space="screen",

major_thickness=2,

minor_thickness=0.5,

infinite=True,

**grid_kwargs,

)

if grids:

_grids = dict()

for plane in GRID_PLANES:

grid = Grid(

geometry=None,

material=pygfx.GridMaterial(**grid_kwargs),

orientation=plane,

visible=False,

)

_grids[plane] = grid

self._grids = Grids(**_grids)

self.world_object.add(self._grids)

if self._plot_area.camera.fov == 0:

# orthographic projection, place grids far away

self._grids.local.z = -1000

major_step_x, major_step_y = statsx["tick_step"], statsy["tick_step"]

self.grids.xy.material.major_step = major_step_x, major_step_y

self.grids.xy.material.minor_step = 0.2 * major_step_x, 0.2 * major_step_y

else:

self._grids = False

self._intersection = intersection

self._auto_grid = auto_grid

self._basis = None

self.basis = basis

self._last_state = self._get_view_state()

@property

def world_object(self) -> pygfx.WorldObject:

return self._world_object

@property

def basis(self) -> np.ndarray:

"""get or set the basis, shape is [3, 3]"""

return self._basis

@basis.setter

def basis(self, basis: np.ndarray):

if basis.shape != (3, 3):

raise ValueError

# apply quaternion to each of x, y, z rulers

for dim, cbasis, new_basis in zip(["x", "y", "z"], CANONICAL_BAIS, basis):

ruler: Ruler = getattr(self, dim)

ruler.local.rotation = quat_from_vecs(cbasis, new_basis)

@property

def offset(self) -> np.ndarray:

"""offset of the axes"""

return self._offset

@offset.setter

def offset(self, value: np.ndarray):

self._offset = value

@property

def x(self) -> Ruler:

"""x axis ruler"""

return self._x

@property

def y(self) -> Ruler:

"""y axis ruler"""

return self._y

@property

def z(self) -> Ruler:

"""z axis ruler"""

return self._z

@property

def grids(self) -> Grids | bool:

"""grids for each plane: xy, xz, yz"""

return self._grids

@property

def colors(self) -> tuple[pygfx.Color]:

return tuple(getattr(self, dim).line.material.color for dim in ["x", "y", "z"])

@colors.setter

def colors(self, colors: tuple[pygfx.Color | str]):

"""get or set the colors for the x, y, and z rulers"""

if len(colors) != 3:

raise ValueError

for dim, color in zip(["x", "y", "z"], colors):

getattr(self, dim).line.material.color = color

@property

def color(self) -> pygfx.Color:

"""get or set a single color for all rulers"""

return self._x.color

@color.setter

def color(self, color: pygfx.Color | str):

for ruler in (self._x, self._y, self._z):

ruler.color = color

@property

def auto_grid(self) -> bool:

"""auto adjust the grid on each render cycle"""

return self._auto_grid

@auto_grid.setter

def auto_grid(self, value: bool):

self._auto_grid = value

@property

def visible(self) -> bool:

"""set visibility of all axes elements, rulers and grids"""

return self._world_object.visible

@visible.setter

def visible(self, value: bool):

self._world_object.visible = value

@property

def intersection(self) -> tuple[float, float, float] | None:

return self._intersection

@intersection.setter

def intersection(self, intersection: tuple[float, float, float] | None):

"""

intersection point of [x, y, z] rulers.

Set (0, 0, 0) for origin

Set to `None` to follow when panning through the scene with orthographic projection

"""

if intersection is None:

self._intersection = None

return

if len(intersection) != 3:

raise ValueError(

"intersection must be a float of 3 elements for [x, y, z] or `None`"

)

self._intersection = tuple(float(v) for v in intersection)

def _get_view_state(self) -> tuple:

viewport = self._plot_area.viewport

cam_matrix = self._plot_area.camera.camera_matrix.tobytes()

scale = self._plot_area.camera.local.scale.tobytes()

# the label margins are the other half of what places the rulers, and they are not known

# until the text has been laid out, which only happens once it has been drawn. tracking

# them here is what redoes the placement on the frame after that, and on any later change

# in the width of a tick label

return (

cam_matrix,

viewport.rect,

viewport.logical_size,

scale,

self._get_label_margins(),

)

def update_using_bbox(self, bbox):

"""

Update the axes w.r.t. the given bbox

Parameters

----------

bbox: np.ndarray

array of shape [2, 3], [[xmin, ymin, zmin], [xmax, ymax, zmax]]

"""

# flip axes if camera scale is flipped

if self._plot_area.camera.local.scale_x < 0:

bbox[0, 0], bbox[1, 0] = bbox[1, 0], bbox[0, 0]

if self._plot_area.camera.local.scale_y < 0:

bbox[0, 1], bbox[1, 1] = bbox[1, 1], bbox[0, 1]

if self._plot_area.camera.local.scale_z < 0:

bbox[0, 2], bbox[1, 2] = bbox[1, 2], bbox[0, 2]

if self.intersection is None:

intersection = (0, 0, 0)

else:

intersection = self.intersection

self.update(bbox, intersection)

def _get_label_margins(self) -> tuple[float, float]:

"""

How far the x and y tick labels, plus their axis labels, reach from their ruler, in pixels

The fallbacks are for text that has not been laid out yet, which is the case until it has

been drawn once.

"""

x_blocks = [b for b in self.x.text._text_blocks if b._rect.height > 0]

x_margin = (

max(abs(b._rect.bottom) for b in x_blocks)

if x_blocks

else 1.5 * self.x.text.font_size

)

if self.x._label._text_blocks:

# the axis label starts at the tick margin, and its own body follows

x_margin += 1.5 * self.x._label.font_size

y_blocks = [b for b in self.y.text._text_blocks if b._rect.width > 0]

y_margin = (

max(abs(b._rect.left) for b in y_blocks)

if y_blocks

else 6 * self.y.text.font_size

)

if self.y._label._text_blocks:

y_margin += 1.5 * self.y._label.font_size

return x_margin, y_margin

def update_using_camera(self):

"""

Update the axes w.r.t the current camera state

For orthographic projections of the xy plane, it will calculate the inverse projection

of the screen space onto world space to determine the current range of the world space

to set the rulers and ticks

For perspective projections it will just use the bbox of the scene to set the rulers

"""

if not self.visible:

return

state = self._get_view_state()

if state == self._last_state:

# no changes in the camera, the viewport rect, or the size of the labels

return

*_, (x_margin, y_margin) = state

if self._plot_area.camera.fov == 0:

xpos, ypos, width, height = self._plot_area.viewport.rect

# orthographic projection, get ranges using inverse

# get range of screen space by getting the corners

xmin, xmax = xpos, xpos + width

ymin, ymax = ypos + height, ypos

min_vals = self._plot_area.map_screen_to_world((xmin, ymin))

max_vals = self._plot_area.map_screen_to_world((xmax, ymax))

if min_vals is None or max_vals is None:

return

world_xmin, world_ymin, _ = min_vals

world_xmax, world_ymax, _ = max_vals

world_zmin, world_zmax = 0, 0

bbox = np.array(

[

[world_xmin, world_ymin, world_zmin],

[world_xmax, world_ymax, world_zmax],

]

)

else:

# set ruler start and end positions based on scene bbox

bbox = self._plot_area._fpl_graphics_scene.get_world_bounding_box()

if self.intersection is None:

if self._plot_area.camera.fov == 0:

# put the rulers in the bottom left corner, clear of their own labels

padding = 4

intersection = self._plot_area.map_screen_to_world(

(xpos + y_margin + padding, ypos + height - x_margin - padding)

)

else:

# force origin since None is not supported for Persepctive projections

self._intersection = (0, 0, 0)

intersection = self._intersection

else:

# axes intersect at the origin

intersection = self.intersection

self.update(bbox, intersection)

self._last_state = state

def update(self, bbox, intersection):

"""

Update the axes using the given bbox and ruler intersection point

Parameters

----------

bbox: np.ndarray

array of shape [2, 3], [[xmin, ymin, zmin], [xmax, ymax, zmax]]

intersection: float, float, float

intersection point of the x, y, z ruler

"""

world_xmin, world_ymin, world_zmin = bbox[0]

world_xmax, world_ymax, world_zmax = bbox[1]

world_x_10, world_y_10, world_z_10 = intersection

# swap min and max for each dimension if necessary

if self._plot_area.camera.local.scale_y < 0:

world_ymin, world_ymax = world_ymax, world_ymin

self.y.tick_side = "right" # swap tick side

self.x.tick_side = "right"

else:

self.y.tick_side = "left"

self.x.tick_side = "right"

if self._plot_area.camera.local.scale_x < 0:

world_xmin, world_xmax = world_xmax, world_xmin

self.x.tick_side = "left"

self.x.start_pos = world_xmin, world_y_10, world_z_10

self.x.end_pos = world_xmax, world_y_10, world_z_10

self.x.start_value = self.x.start_pos[0] - self.offset[0]

statsx = self.x.update(

self._plot_area.camera, self._plot_area.viewport.logical_size

)

self.y.start_pos = world_x_10, world_ymin, world_z_10

self.y.end_pos = world_x_10, world_ymax, world_z_10

self.y.start_value = self.y.start_pos[1] - self.offset[1]

statsy = self.y.update(

self._plot_area.camera, self._plot_area.viewport.logical_size

)

if self._plot_area.camera.fov != 0:

self.z.start_pos = world_x_10, world_y_10, world_zmin

self.z.end_pos = world_x_10, world_y_10, world_zmax

self.z.start_value = self.z.start_pos[2] - self.offset[2]

statsz = self.z.update(

self._plot_area.camera, self._plot_area.viewport.logical_size

)

major_step_z = statsz["tick_step"]

if self.grids:

if self.auto_grid:

major_step_x, major_step_y = statsx["tick_step"], statsy["tick_step"]

self.grids.xy.major_step = major_step_x, major_step_y

self.grids.xy.minor_step = 0.2 * major_step_x, 0.2 * major_step_y

if self._plot_area.camera.fov != 0:

self.grids.xz.major_step = major_step_x, major_step_z

self.grids.xz.minor_step = 0.2 * major_step_x, 0.2 * major_step_z

self.grids.yz.material.major_step = major_step_y, major_step_z

self.grids.yz.minor_step = 0.2 * major_step_y, 0.2 * major_step_z