tree1.py

from collections import defaultdict
import mathutils
from mathutils import Vector, kdtree


def grow(bv, tv, br, tr):

    basesize = len(bv)
    size = len(tv)
    kd = kdtree.KDTree(size)

    for i, vtx in enumerate(bv):
        kd.insert(Vector(vtx), i)
    kd.balance()

    kd5 = defaultdict(list)

    for i, vtx in enumerate(tv):
        co, index, dist = kd.find(vtx)
        kd5[index].append(i+basesize)

    print(kd5)
    edges = []
    for key, values in kd5.items():
        for idx in values:
            edges.append((idx, key))

    intermediate_vertices = bv + tv

    def get_median(v_list):
        x, y, z = 0, 0, 0
        n = len(v_list)
        for v in v_list:
            x += v[0]
            y += v[1]
            z += v[2]
        return (x/n, y/n, z/n)

    additional = []
    for key, values in kd5.items():
        if len(values) > 1:
            vecs = [intermediate_vertices[i] for i in values + [key]]
            new_vec = get_median(vecs)
            print(new_vec)
            additional.append(new_vec)

    final_verts = intermediate_vertices + additional

    return final_verts, edges


def sv_main(base_verts=[[]], base_r=1.0, tip_verts=[[]], tip_r=0.5):

    verts_out = []
    edges_out = []
    radii_out = []

    in_sockets = [
        ['v', 'base_verts', base_verts],
        ['s', 'base_radius', base_r],
        ['v', 'tip_verts', tip_verts],
        ['s', 'tip_radius', tip_r]
    ]

    out_sockets = [
        ['v', 'verts', [verts_out]],
        ['s', 'edges', [edges_out]],
        ['s', 'radii', radii_out]
    ]

    if base_verts and base_verts[0]:
        bv = base_verts[0]
        if tip_verts and tip_verts[0]:
            tv = tip_verts[0]
            v, e = grow(bv, tv, base_r, tip_r)
            verts_out.extend(v)
            edges_out.extend(e)

    return in_sockets, out_sockets