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_map = {
            '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):
            io_map = [(inputs, expected_inputs), (outputs, expected_outputs)]

            for io_puts, expected_io_puts in io_map:
                # remove inappropriate sockets
                for socket in io_puts:
                    if not socket in expected_io_puts:
                        io_puts.remove(socket)

                # add desired sockets
                for socket in expected_io_puts:
                    if not socket in io_puts:
                        socket_type = socket_map.get(socket, 'StringsSocket')
                        io_puts.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']

        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']

        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']
        else:
            return

        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()