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',
            'grouped dist': '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)

        standard_inputs = ['Verts', 'Check Verts']

        if mode == 'FIND':
            expected_inputs = standard_inputs
            expected_outputs = ['proximity Verts', 'proximity Indices']
        elif mode == 'FIND_N':
            expected_inputs = standard_inputs + ['n nearest']
            expected_outputs = ['n proxima .co', 'n proxima .idx']
        elif mode == 'FIND_RANGE':
            expected_inputs = standard_inputs + ['radius']
            expected_outputs = ['grouped .co', 'grouped .idx', 'grouped dist']
        else:
            return

        prepare_node(expected_inputs, expected_outputs)

    modes = [
        ("FIND", " 1 ", "Find nearest", 1),
        ("FIND_N", " n ", "Find n nearest", 2),
        ("FIND_RANGE", "radius", "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')

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

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

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

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

        if self.mode == 'FIND':
            ''' [Verts.co,..] =>
                [Verts.idx,.] =>
            => [Main Verts]
            => [cVert,..]
            '''
            pass

        elif self.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
            '''
            pass

        elif self.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
            '''
            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()