node_KDTree.py

from node_s import *
from util import *
from mathutils import Vector, Euler, kdtree
import bpy


class SvKDTreeNode(Node, SverchCustomTreeNode):

    bl_idname = 'SvKDTreeNode'
    bl_label = 'Kdtree search'
    bl_icon = 'OUTLINER_OB_EMPTY'

    def mode_change(self, context):
        mode = self.mode
        inputs = self.inputs
        outputs = self.outputs

        socket_mapping = {
            'Verts': 'VerticesSocket',
            'Check Verts': 'VerticesSocket',
            'proximity Verts': 'VerticesSocket',
            'proximity Indices': 'StringsSocket',
            'n nearest': 'StringsSocket',
            'n proxima .co': 'VerticesSocket',
            'n proxima .idx': 'StringsSocket',
            'grouped .co': 'VerticesSocket',
            'grouped .idx': 'StringsSocket'
        }

        def prepare_node(expected_inputs, expected_outputs):
            # refactpr lazybones.

            # (input) remove inappropriate sockets
            for socket in inputs:
                if not socket in expected_inputs:
                    inputs.remove(socket)

            # (input) add desired sockets
            for socket in expected_inputs:
                if not socket in inputs:
                    socket_type = socket_mapping.get(socket, 'StringsSocket')
                    inputs.new(socket_type, socket, socket)

            # (output) remove inappropriate sockets
            for socket in outputs:
                if not socket in expected_outputs:
                    outputs.remove(socket)

             # (output) add desired sockets
            for socket in expected_outputs:
                if not socket in outputs:
                    socket_type = socket_mapping.get(socket, 'StringsSocket')
                    outputs.new(socket_type, socket, socket)

        if mode == 'FIND':
            '''         [Verts.co,..] =>
                        [Verts.idx,.] =>
                => [Main Verts]
                => [cVert,..]
            '''
            expected_inputs = ['Verts', 'Check Verts']
            expected_outputs = ['proximity Verts', 'proximity Indices']
            prepare_node(expected_inputs, expected_outputs)

        elif mode == 'FIND_N':
            '''         [[Verts.co,..n],..c] => from MainVerts closest to v.co
                        [[Verts.idx,..n],.c] => from MainVerts closest to v.co
                => [Main Verts]
                => [cVert,..]
                => [n, max n nearest
            '''
            expected_inputs = ['Verts', 'Check Verts', 'n nearest']
            expected_outputs = ['n proxima .co', 'n proxima .idx']
            prepare_node(expected_inputs, expected_outputs)
            pass

        elif mode == 'FIND_RANGE':
            '''     [grouped [.co for p in MainVerts in r of v in cVert]] =>
                    [grouped [.idx for p in MainVerts in r of v in cVert]] =>
                    [grouped [.dist for p in MainVerts in r of v in cVert]] =>
                => [Main Verts]
                => [cVert,..]
                => n
            '''
            expected_inputs = ['Verts', 'Check Verts', 'radius']
            expected_outputs = ['grouped .co', 'grouped .idx']
            prepare_node(expected_inputs, expected_outputs)
            pass

    modes = [
        ("FIND", " 1 ", "Find nearest", 1),
        ("FIND_N", " n ", "Find n nearest", 2),
        ("FIND_RANGE", "distance", "Find within Distance", 3)
    ]

    mode = bpy.props.EnumProperty(
        items=modes, default='FIND', update=mode_change)

    def draw_buttons(self, context, layout):
        layout.label("Search mode:")
        layout.prop(self, "mode", expand=True)

    def init(self, context):
        self.inputs.new('VerticesSocket', 'Verts', 'Verts')
        self.inputs.new('VerticesSocket', 'Check Verts', 'Check Verts')

        self.outputs.new(
            'VerticesSocket', 'proximity Verts', 'proximity Verts')
        self.outputs.new(
            'StringsSocket', 'proximity Indices', 'proximity Indices')

    def update(self):
        inputs = self.inputs
        outputs = self.outputs

        if not ('Vertices' in inputs and inputs['Vertices'].links):
            return

        verts = SvGetSocketAnyType(self, inputs['Vertices'])

        if self.mode == 'FIND':

            pass

        elif self.mode == 'FIND_N':
            pass

        elif self.mode == 'FIND_RANGE':
            pass

        # SvSetSocketAnyType(self, 'PolyEdge', edges_out)

    def update_socket(self, context):
        self.update()


def register():
    bpy.utils.register_class(SvKDTreeNode)


def unregister():
    bpy.utils.unregister_class(SvKDTreeNode)

if __name__ == "__main__":
    register()