import numpy as np
from scipy import sparse, ndimage as ndi
import networkx as nx
import numba


@numba.jit
def get_nz_neighbors(arr, idx, steps, out):
    pass


def skeleton_to_csgraph(skel):
    skelint = np.zeros(skel.shape, int)
    skelint[skel] = np.arange(1, np.sum(skel) + 1)


def _add_skeleton_edges(values, graph, distances):
    values = values.astype(int)
    center_idx = len(values) // 2
    center = values[center_idx]
    if center == 0:
        return 0
    count = 0.
    for value, distance in zip(values, distances):
        if value != center and value != 0:
            graph.add_edge(center, value, weight=distance)
            count += 1
    node = graph.node[center]
    if count <= 1:
        node['type'] = 'tip'
    elif count == 2:
        node['type'] = 'path'
    else:
        node['type'] = 'junction'
    return count


def _neighbor_distances(ndim):
    center = np.ones((3,) * ndim, dtype=bool)
    center.ravel()[center.size//2] = False
    out = ndi.distance_transform_edt(center).ravel()
    return out


def skeleton_to_nx(skel):  # to do: add pixel spacing
    g = nx.Graph()
    distances = _neighbor_distances(skel.ndim)
    skelint = np.zeros(skel.shape, int)
    num_nodes = np.sum(skel)
    g.add_nodes_from(range(1, num_nodes + 1))
    skelint[skel] = np.arange(1, num_nodes + 1)
    ndi.generic_filter(skelint, function=_add_skeleton_edges,
                       size=3, mode='constant', cval=0,
                       extra_arguments=(g, distances))
    return g

@@ -1508,8 +1508,1496 @@
 eturn g%0A
+%0A%0Adef branch_statistics(g):%0A    %22%22%22Compute the length and type of each branch in a skeleton graph.%0A%0A    Parameters%0A    ----------%0A    g : nx.Graph%0A        A skeleton graph. Its nodes must be marked as 'tip', 'path', or%0A        'junction'.%0A%0A    Returns%0A    -------%0A    branches : array of float, shape (N, 4, 2)%0A        An array containing branch endpoint IDs, length, and branch type.%0A        The types are:%0A        - tip-tip (0)%0A        - tip-junction (1)%0A        - junction-junction (2)%0A    %22%22%22%0A    visited = np.zeros(max(g) + 1, dtype=bool)%0A    type_dict = %7B'tiptip': 0, 'tipjunction': 1, 'junctiontip': 1,%0A                 'junctionjunction': 2%7D%0A    result = %5B%5D%0A    for node, data in g.nodes_iter(data=True):%0A        if data%5B'type'%5D == 'path' and not visited%5Bnode%5D:%0A            # we expand the path in either direction%0A            visited%5Bnode%5D = True%0A            left, right = g.neighbors(node)%0A            id0, d0, kind0 = _expand_path(g, node, left, visited)%0A            id1, d1, kind1 = _expand_path(g, node, right, visited)%0A            result.append(%5Bid0, id1, d0 + d1, type_dict%5Bkind0 + kind1%5D%5D)%0A    return np.array(result)%0A%0A%0Adef _expand_path(g, source, step, visited):%0A    d = g%5Bsource%5D%5Bstep%5D%5B'weight'%5D%0A    while g.node%5Bstep%5D%5B'type'%5D == 'path':%0A        n1, n2 = g.neighbors(step)%0A        nextstep = n1 if n1 != source else n2%0A        source, step = step, nextstep%0A        d += g%5Bsource%5D%5Bstep%5D%5B'weight'%5D%0A        visited%5Bsource%5D = True%0A    return source, d, g.node%5Bstep%5D%5B'type'%5D%0A

Add branch stats function