Viewport Interactivity #1121

import bpy
import bl_ext.blender_org.molecularnodes as mn  # type: ignore
import MDAnalysis as mda
from bpy_extras.view3d_utils import region_2d_to_origin_3d, region_2d_to_vector_3d
from mathutils import kdtree
from MDAnalysis.tests.datafiles import PSF, DCD


def setup_universe():
    u = mda.Universe(PSF, DCD)
    # add trajectory
    t = mn.Trajectory(u)
    t.add_style("cartoon", color=(0.2, 0.6, 0.2, 0.2))
    t.add_style("ball_and_stick", selection="resid 1")
    t.styles[0].quality = 3
    t.styles[1].quality = 3
    # add annotation to show hovered atom info
    a1 = t.annotations.add_atom_info(
        name="Hovered Atom Info", selection="index 0", show_resid=True, show_segid=True
    )
    a1.text_color = (1, 1, 0, 1)
    a1.text_depth = False
    # add annotation to show com distance from protein center
    a2 = t.annotations.add_com_distance(
        name="Hovered COM Distance",
        selection1="protein",
        selection2="resid 1",
        text1="protein|COM",
    )
    a2.line_arrow_size = 0.05
    # add universe info annotation
    t.annotations.add_universe_info(name="Universe Info")


class MN_OT_Viewport_Interaction(bpy.types.Operator):
    bl_idname = "mn.viewport_interaction"
    bl_label = "MN Viewport Interaction"

    _kdtree_cache = {}
    _kdtree_cache_frame = None

    def _get_kdtree(self, obj):
        mesh = obj.data
        size = len(mesh.vertices)
        if self._kdtree_cache_frame != bpy.context.scene.frame_current:
            # clear the kdtree cache for frame changes
            self._kdtree_cache.clear()
            self._kdtree_cache_frame = bpy.context.scene.frame_current
        key = obj.uuid
        if key not in self._kdtree_cache:
            kd = kdtree.KDTree(size)
            for i, v in enumerate(mesh.vertices):
                kd.insert(v.co, i)
            kd.balance()
            self._kdtree_cache[key] = kd
        return self._kdtree_cache[key]

    def modal(self, context, event):
        area = next((a for a in context.screen.areas if a.type == "VIEW_3D"), None)
        if area and event.type == "MOUSEMOVE":
            region = next((r for r in area.regions if r.type == "WINDOW"), None)
            rv3d = area.spaces.active.region_3d
            if region and rv3d:
                coord = (event.mouse_x - region.x, event.mouse_y - region.y)
                origin = region_2d_to_origin_3d(region, rv3d, coord)
                direction = region_2d_to_vector_3d(region, rv3d, coord)

                # depsgraph = context.evaluated_depsgraph_get()
                depsgraph = context.view_layer.depsgraph
                hit, location, _, _, obj, _ = context.scene.ray_cast(
                    depsgraph, origin, direction
                )

                if hit and obj.uuid:
                    # hit for an MN entity
                    local_location = obj.matrix_world.inverted() @ location
                    kdtree = self._get_kdtree(obj)
                    _, index, _ = kdtree.find(local_location)
                    # update entity values (WARNING: currently hard-coded stuff)
                    entity = bpy.context.scene.MNSession.entities[obj.uuid]
                    resid = entity.universe.atoms[index].resid
                    entity.annotations[0].selection = f"index {index}"
                    entity.annotations[1].selection2 = f"resid {resid}"
                    obj.mn_trajectory_selections[0].string = f"resid {resid}"

        if event.type in {"RIGHTMOUSE", "ESC"}:
            self._kdtree_cache.clear()
            return {"CANCELLED"}

        return {"PASS_THROUGH"}

    def invoke(self, context, event):
        context.window_manager.modal_handler_add(self)
        return {"RUNNING_MODAL"}


def register_op():
    bpy.utils.register_class(MN_OT_Viewport_Interaction)


def unregister_op():
    bpy.utils.unregister_class(MN_OT_Viewport_Interaction)


def run_op():
    bpy.ops.mn.viewport_interaction("INVOKE_DEFAULT")


setup_universe()
register_op()
run_op()