path: root/mesh_snap_utilities_line/snap_context_l/__init__.py

# SPDX-License-Identifier: GPL-2.0-or-later

__all__ = ("SnapContext",)

import gpu
from mathutils import Vector

VERT = 1
EDGE = 2
FACE = 4

class _Internal:
    global_snap_context = None

    @classmethod
    def snap_context_free(cls):
        if cls.global_snap_context != None:
            cls.global_snap_context.free()
            del cls

    from .mesh_drawing import (
        gpu_Indices_enable_state,
        gpu_Indices_restore_state,
        gpu_Indices_use_clip_planes,
        gpu_Indices_mesh_cache_clear,
        )

    from .utils_projection import (
        region_2d_to_orig_and_view_vector,
        intersect_boundbox_threshold,
        intersect_ray_segment_fac,
        project_co_v3,
        )

    from mathutils.geometry import intersect_line_plane


class _SnapObjectData():
    __slots__ = ('data', 'mat')
    def __init__(self, data, omat):
        self.data = data
        self.mat = omat


class _SnapOffscreen():
    bound = None
    def __init__(self, width, height):
        self._fbo = None

        self.width = width
        self.height = height

        self._framebuffer_config()

    def _framebuffer_config(self):
        self._framebuffer_free()

        self._tex_color = gpu.types.GPUTexture((self.width, self.height), format='R32UI')
        self._tex_depth = gpu.types.GPUTexture((self.width, self.height), format='DEPTH_COMPONENT32F')
        self._fbo = gpu.types.GPUFrameBuffer(depth_slot=self._tex_depth, color_slots=self._tex_color)

    def _framebuffer_free(self):
        self._fbo = self._tex_color = self._tex_depth = None

    def bind(self):
        return self._fbo.bind()

    def clear(self):
        #self._fbo.clear(color=(0.0, 0.0, 0.0, 0.0), depth=1.0)
        self._tex_color.clear(format='UINT', value=(0,))
        self._tex_depth.clear(format='FLOAT', value=(1.0,))

    def resize(self, width, height):
        self.width = int(width)
        self.height = int(height)
        self._framebuffer_config()

    def __del__(self):
        self._framebuffer_free()

    def free(self):
        self._framebuffer_free()


class SnapContext():
    """
    Initializes the snap context with the region and space where the snap objects will be added.

    .. note::
        After the context has been created, add the objects with the `add_obj` method.

    :arg region: region of the 3D viewport, typically bpy.context.region.
    :type region: :class:`bpy.types.Region`
    :arg space: 3D region data, typically bpy.context.space_data.
    :type space: :class:`bpy.types.SpaceView3D`
    """

    __slots__ = (\
        '_dist_px',
        '_dist_px_sq',
        '_offscreen',
        '_offset_cur',
        '_snap_buffer',
        '_snap_mode',
        'depsgraph',
        'depth_range',
        'drawn_count',
        'freed',
        'last_ray',
        'mval',
        'proj_mat',
        'region',
        'rv3d',
        'snap_objects',
        'threshold',
        'winsize',)

    def __init__(self, depsgraph, region, space):
        self.freed = False
        self.snap_objects = []
        self.drawn_count = 0
        self._offset_cur = 1 # Starts with index 1

        self.proj_mat = None
        self.mval = Vector((0.0, 0.0))
        self.last_ray = Vector((1.0, 0.0, 0.0))
        self._snap_mode = VERT | EDGE | FACE

        self.set_pixel_dist(12)

        self.region = region
        self.depsgraph = depsgraph

        if not (depsgraph or region or space):
            self.rv3d = None
            self.depth_range = Vector((-1.0, 1.0))
            self._offscreen = _SnapOffscreen(1, 1)
        else:
            self.rv3d = space.region_3d

            if self.rv3d.is_perspective:
                self.depth_range = Vector((space.clip_start, space.clip_end))
            else:
                self.depth_range = Vector((-space.clip_end, space.clip_end))

            self._offscreen = _SnapOffscreen(self.region.width, self.region.height)
            self._offscreen.clear()

        self.winsize = Vector((self._offscreen.width, self._offscreen.height))


    ## PRIVATE ##

    def _get_snap_obj_by_index(self, index):
        if index:
            for snap_obj in self.snap_objects[:self.drawn_count]:
                data = snap_obj.data[1]
                if index < data.first_index + data.get_tot_elems():
                    return snap_obj
        return None

    def _read_buffer(self):
        self._snap_buffer = self._offscreen._tex_color.read()

    def _get_nearest_index(self, mval):
        r_snap_obj = None
        r_value = 0

        loc_curr = [int(mval[1]), int(mval[0])]
        rect = ((max(0, loc_curr[0] - self.threshold), min(self._snap_buffer.dimensions[0], loc_curr[0] + self.threshold)),
                (max(0, loc_curr[1] - self.threshold), min(self._snap_buffer.dimensions[1], loc_curr[1] + self.threshold)))

        if loc_curr[0] < rect[0][0] or loc_curr[0] >= rect[0][1] or loc_curr[1] < rect[1][0] or loc_curr[1] >= rect[1][1]:
            return r_snap_obj, r_value

        find_next_index = self._snap_mode & FACE and self._snap_mode & (VERT | EDGE)
        last_value = -1 if find_next_index else 0
        spiral_direction = 0
        for nr in range(1, 2 * self.threshold):
            for a in range(2):
                for b in range(0, nr):
                    # TODO: Make the buffer flat.
                    value = int(self._snap_buffer[loc_curr[0]][loc_curr[1]])
                    if value != last_value:
                        r_value = value
                        if find_next_index:
                            last_value = value
                            r_snap_obj = self._get_snap_obj_by_index(value)
                            if r_snap_obj is not None:
                                snap_data = r_snap_obj.data[1]
                                if value < (snap_data.first_index + len(snap_data.tri_verts)):
                                    # snap to a triangle
                                    find_next_index = False
                                else:
                                    return r_snap_obj, r_value
                        else:
                            if (r_snap_obj is None) or \
                               (value < r_snap_obj.data[1].first_index) or \
                               (value >= (r_snap_obj.data[1].first_index + r_snap_obj.data[1].get_tot_elems())):
                                r_snap_obj = self._get_snap_obj_by_index(value)

                            return r_snap_obj, r_value

                    # Next spiral step.
                    if (spiral_direction == 0):
                        loc_curr[1] += 1 # right
                    elif (spiral_direction == 1):
                        loc_curr[0] -= 1 # down
                    elif (spiral_direction == 2):
                        loc_curr[1] -= 1 # left
                    else:
                        loc_curr[0] += 1 # up

                    if (loc_curr[not a] < rect[not a][0] or loc_curr[not a] >= rect[not a][1]):
                        return r_snap_obj, r_value

                spiral_direction = (spiral_direction + 1) % 4
        return r_snap_obj, r_value

    def _get_loc(self, snap_obj, index):
        index -= snap_obj.data[1].first_index
        gpu_data = snap_obj.data[1]

        if gpu_data.draw_tris:
            num_tris = len(snap_obj.data[1].tri_verts)
            if index < num_tris:
                tri_verts = gpu_data.get_tri_verts(index)
                tri_co = [snap_obj.mat @ Vector(v) for v in gpu_data.get_tri_co(index)]
                #loc = _Internal.intersect_ray_tri(*tri_co, self.last_ray[0],
                #self.last_ray[1], False)
                nor = (tri_co[1] - tri_co[0]).cross(tri_co[2] - tri_co[0]).normalized()
                loc = _Internal.intersect_line_plane(self.last_ray[1], self.last_ray[1] + self.last_ray[0], tri_co[0], nor)
                return loc, tri_verts, tri_co

            index -= num_tris

        if gpu_data.draw_edges:
            num_edges = len(snap_obj.data[1].edge_verts)
            if index < num_edges:
                edge_verts = gpu_data.get_edge_verts(index)
                edge_co = [snap_obj.mat @ Vector(v) for v in gpu_data.get_edge_co(index)]
                fac = _Internal.intersect_ray_segment_fac(*edge_co, *self.last_ray)

                if (self._snap_mode) & VERT and (fac < 0.25 or fac > 0.75):
                    co = edge_co[0] if fac < 0.5 else edge_co[1]
                    proj_co = _Internal.project_co_v3(self, co)
                    dist = self.mval - proj_co
                    if abs(dist.x) < self._dist_px and abs(dist.y) < self._dist_px:
                        return co, (edge_verts[0] if fac < 0.5 else edge_verts[1],), co

                if fac <= 0.0:
                    co = edge_co[0]
                elif fac >= 1.0:
                    co = edge_co[1]
                else:
                    co = edge_co[0] + fac * (edge_co[1] - edge_co[0])

                return co, edge_verts, edge_co

            index -= num_edges

        if gpu_data.draw_verts:
            if index < len(snap_obj.data[1].looseverts):
                co = snap_obj.mat @ Vector(gpu_data.get_loosevert_co(index))
                return co, (gpu_data.get_loosevert_index(index),), co

        return None, None, None


    def _get_snap_obj_by_obj(self, obj):
        for snap_obj in self.snap_objects:
            if obj == snap_obj.data[0]:
                return snap_obj

    def __del__(self):
        if not self.freed:
            self._offscreen.free()
            self._snap_buffer = None
            # Some objects may still be being referenced
            for snap_obj in self.snap_objects:
                if len(snap_obj.data) == 2:
                    snap_obj.data[1].free()
                    del snap_obj.data[1:]

                del snap_obj.data
                del snap_obj.mat
                del snap_obj
            del self.snap_objects

    ## PUBLIC ##

    def update_viewport_context(self, depsgraph, region, space, resize=False):
        rv3d = space.region_3d

        if not resize and self.rv3d == rv3d and self.region == region:
            return

        self.depsgraph = depsgraph
        self.region = region
        self.rv3d = rv3d

        if self.rv3d.is_perspective:
            self.depth_range = Vector((space.clip_start, space.clip_end))
        else:
            self.depth_range = Vector((-space.clip_end, space.clip_end))

        winsize = Vector((self.region.width, self.region.height))

        if winsize != self.winsize:
            self.winsize = winsize
            self._offscreen.resize(*self.winsize)

    def clear_snap_objects(self, clear_offscreen=False):
        for snap_obj in self.snap_objects:
            if len(snap_obj.data) == 2:
                snap_obj.data[1].free()
                del snap_obj.data[1:]

        self.update_drawing(clear_offscreen)

        self.snap_objects.clear()
        _Internal.gpu_Indices_mesh_cache_clear()

    def update_all(self):
        for snap_obj in self.snap_objects:
            if len(snap_obj.data) == 2:
                snap_obj.data[1].free()
                del snap_obj.data[1:]

        self.update_drawing()

    def update_drawing(self, clear_offscreen=True):
        self.drawn_count = 0
        self._offset_cur = 1
        if clear_offscreen:
            self._offscreen.clear()

        _Internal.gpu_Indices_use_clip_planes(self.rv3d, True)

    def tag_update_drawn_snap_object(self, snap_obj):
        if len(snap_obj.data) > 1:
            snap_obj.data[1].free()
            del snap_obj.data[1:]
            #self.update_drawing()
            # Update on next snap_get call #
            self.proj_mat = None

    def update_drawn_snap_object(self, snap_obj):
        _Internal.gpu_Indices_enable_state(self.rv3d.window_matrix, self.rv3d.view_matrix)

        from .mesh_drawing import GPU_Indices_Mesh
        snap_vert = self._snap_mode & VERT != 0
        snap_edge = self._snap_mode & EDGE != 0
        snap_face = self._snap_mode & FACE != 0

        if len(snap_obj.data) > 1:
            snap_obj.data[1].free()
            del snap_obj.data[1:]

        data = GPU_Indices_Mesh(self.depsgraph, snap_obj.data[0], snap_face, snap_edge, snap_vert)
        snap_obj.data.append(data)

        _Internal.gpu_Indices_restore_state()

    def set_pixel_dist(self, dist_px):
        self._dist_px = int(dist_px)
        self._dist_px_sq = self._dist_px ** 2
        self.threshold = 2 * self._dist_px + 1

    def set_snap_mode(self, snap_to_vert, snap_to_edge, snap_to_face):
        snap_mode = 0
        if snap_to_vert:
            snap_mode |= VERT
        if snap_to_edge:
            snap_mode |= EDGE
        if snap_to_face:
            snap_mode |= FACE

        if snap_mode != self._snap_mode:
            self._snap_mode = snap_mode
            self.update_all()

    def add_obj(self, obj, matrix):
        self.snap_objects.append(_SnapObjectData([obj], matrix.copy()))

        return self.snap_objects[-1]

    def get_ray(self, mval):
        self.last_ray = _Internal.region_2d_to_orig_and_view_vector(self.region, self.rv3d, mval)
        return self.last_ray

    def snap_get(self, mval, main_snap_obj=None):
        ret = None, None, None
        self.mval[:] = mval
        snap_vert = self._snap_mode & VERT != 0
        snap_edge = self._snap_mode & EDGE != 0
        snap_face = self._snap_mode & FACE != 0

        _Internal.gpu_Indices_enable_state(self.rv3d.window_matrix, self.rv3d.view_matrix)
        gpu.state.depth_mask_set(True)
        gpu.state.depth_test_set('LESS_EQUAL')
        gpu.state.program_point_size_set(True)

        with self._offscreen.bind():
            update_buffer = False
            proj_mat = self.rv3d.perspective_matrix.copy()
            if self.proj_mat != proj_mat:
                self.proj_mat = proj_mat
                self.update_drawing()
                update_buffer = True

            ray_dir, ray_orig = self.get_ray(mval)
            for i, snap_obj in enumerate(self.snap_objects[self.drawn_count:], self.drawn_count):
                obj = snap_obj.data[0]
                try:
                    bbmin = Vector(obj.bound_box[0])
                    bbmax = Vector(obj.bound_box[6])
                except ReferenceError:
                    self.snap_objects.remove(snap_obj)
                    continue

                if bbmin != bbmax:
                    MVP = proj_mat @ snap_obj.mat
                    mat_inv = snap_obj.mat.inverted_safe()
                    ray_orig_local = mat_inv @ ray_orig
                    ray_dir_local = mat_inv.to_3x3() @ ray_dir
                    in_threshold = _Internal.intersect_boundbox_threshold(self, MVP, ray_orig_local, ray_dir_local, bbmin, bbmax)
                else:
                    proj_co = _Internal.project_co_v3(self, snap_obj.mat.translation)
                    dist = self.mval - proj_co
                    in_threshold = abs(dist.x) < self._dist_px and abs(dist.y) < self._dist_px
                    #snap_obj.data[1] = primitive_point

                if in_threshold:
                    if len(snap_obj.data) == 1:
                        from .mesh_drawing import GPU_Indices_Mesh
                        is_bound = obj.display_type == 'BOUNDS'
                        draw_face = snap_face and not is_bound and obj.display_type != 'WIRE'
                        draw_edge = snap_edge and not is_bound
                        draw_vert = snap_vert and not is_bound
                        snap_obj.data.append(GPU_Indices_Mesh(self.depsgraph, obj, draw_face, draw_edge, draw_vert))

                    snap_obj.data[1].set_draw_mode(snap_face, snap_edge, snap_vert)

                    if snap_obj == main_snap_obj:
                        snap_obj.data[1].Draw(self._offset_cur, snap_obj.mat, 0.0001)
                    else:
                        snap_obj.data[1].Draw(self._offset_cur, snap_obj.mat)
                    self._offset_cur += snap_obj.data[1].get_tot_elems()

                    tmp = self.snap_objects[self.drawn_count]
                    self.snap_objects[self.drawn_count] = self.snap_objects[i]
                    self.snap_objects[i] = tmp

                    self.drawn_count += 1
                    update_buffer = True

            if update_buffer:
                self._read_buffer()
                #import numpy as np
                #a = np.array(self._snap_buffer)
                #print(a)

            snap_obj, index = self._get_nearest_index(mval)
            #print("index:", index)
            if snap_obj:
                ret = self._get_loc(snap_obj, index)

        gpu.state.program_point_size_set(False)
        gpu.state.depth_mask_set(False)
        gpu.state.depth_test_set('NONE')
        _Internal.gpu_Indices_restore_state()

        return (snap_obj, *ret)

    def free(self):
        self.__del__()
        self.freed = True


def global_snap_context_get(depsgraph, region, space):
    if _Internal.global_snap_context == None:
        import atexit

        _Internal.global_snap_context = SnapContext(depsgraph, region, space)

        # Make sure we only registered the callback once.
        atexit.unregister(_Internal.snap_context_free)
        atexit.register(_Internal.snap_context_free)

    elif (depsgraph and region and space):
        _Internal.global_snap_context.update_viewport_context(depsgraph, region, space, True)

    return _Internal.global_snap_context