graph

@utils.map_neuronlist(desc="Classifying", allow_parallel=True)
@utils.lock_neuron
def classify_nodes(x: "core.NeuronObject", categorical=True, inplace: bool = True):
    """Classify neuron's nodes into end nodes, branches, slabs or root.

    Adds a `'type'` column to `x.nodes` table.

    Parameters
    ----------
    x :         TreeNeuron | NeuronList
                Neuron(s) whose nodes to classify.
    categorical : bool
                If True (default), will use categorical data type which takes
                up much less memory at a small run-time overhead.
    inplace :   bool, optional
                If `False`, nodes will be classified on a copy which is then
                returned leaving the original neuron unchanged.

    Returns
    -------
    TreeNeuron/List

    Examples
    --------
    >>> import navis
    >>> nl = navis.example_neurons(2)
    >>> _ = navis.graph.classify_nodes(nl, inplace=True)

    """
    if not inplace:
        x = x.copy()

    if not isinstance(x, core.TreeNeuron):
        raise TypeError(f'Expected TreeNeuron(s), got "{type(x)}"')

    if x.nodes.empty:
        x.nodes["type"] = None
        return x

    # Make sure there are nodes to classify
    # Note: I have tried to optimized the s**t out of this, i.e. every
    # single line of code here has been tested for speed. Do not
    # change anything unless you know what you're doing!

    # Turns out that numpy.isin() recently started to complain if the
    # node_ids are uint64 and the parent_ids are int64 (but strangely
    # not with 32bit integers). If that's the case we have to convert
    # the node_ids to int64.
    node_ids = x.nodes.node_id.values
    parent_ids = x.nodes.parent_id.values

    if node_ids.dtype == np.uint64:
        node_ids = node_ids.astype(np.int64)

    cl = np.full(len(x.nodes), "slab", dtype="<U6")
    cl[~np.isin(node_ids, parent_ids)] = "end"
    bp = x.nodes.parent_id.value_counts()
    bp = bp.index.values[bp.values > 1]
    cl[np.isin(node_ids, bp)] = "branch"
    cl[parent_ids < 0] = "root"
    if categorical:
        cl = pd.Categorical(
            cl, categories=["end", "branch", "root", "slab"], ordered=False
        )
    x.nodes["type"] = cl

    return x

def connected_subgraph(
    x: Union["core.TreeNeuron", nx.DiGraph], ss: Sequence[Union[str, int]]
) -> Tuple[np.ndarray, Union[int, str]]:
    """Return set of nodes necessary to connect all nodes in subset `ss`.

    Parameters
    ----------
    x :         navis.TreeNeuron | nx.DiGraph
                Neuron (or graph thereof) to get subgraph for.
    ss :        list | array-like
                Node IDs of node to subset to.

    Returns
    -------
    np.ndarray
                Node IDs of connected subgraph.
    root ID
                ID of the node most proximal to the old root in the
                connected subgraph.

    Examples
    --------
    >>> import navis
    >>> n = navis.example_neurons(1)
    >>> ends = n.nodes[n.nodes.type.isin(['end', 'root'])].node_id.values
    >>> sg, root = navis.graph.graph_utils.connected_subgraph(n, ends)
    >>> # Since we asked for a subgraph connecting all terminals + root,
    >>> # we expect to see all nodes in the subgraph
    >>> sg.shape[0] == n.nodes.shape[0]
    True

    """
    if isinstance(x, core.NeuronList):
        if len(x) == 1:
            g = x[0].graph
    elif isinstance(x, core.TreeNeuron):
        g = x.graph
    elif isinstance(x, nx.DiGraph):
        g = x
    else:
        raise TypeError(f'Input must be a single TreeNeuron or graph, got "{type(x)}".')

    ss = set(ss)
    missing = ss - set(g.nodes)
    if np.any(missing):
        missing = np.array(list(missing)).astype(str)  # do NOT remove list() here!
        raise ValueError(f'Nodes not found: {",".join(missing)}')

    # Find nodes that are leafs WITHIN the subset
    g_ss = g.subgraph(ss)
    in_degree = dict(g_ss.in_degree)
    leafs = ss & {n for n, d in in_degree.items() if not d}

    # Run this for each connected component of the neuron
    include = set()
    new_roots = []
    for cc in nx.connected_components(g.to_undirected()):
        # Walk from each node to root and keep track of path
        paths = []
        for n in leafs & cc:
            this_path = []
            while n is not None:
                this_path.append(n)
                n = next(g.successors(n), None)
            paths.append(this_path)

        # If none of these cc in subset there won't be paths
        if not paths:
            continue

        # Find the nodes that all paths have in common
        common = set.intersection(*[set(p) for p in paths])

        # Now find the first (most distal from root) common node
        longest_path = sorted(paths, key=lambda x: len(x))[-1]
        first_common = sorted(common, key=lambda x: longest_path.index(x))[0]

        # Now go back to paths and collect all nodes until this first common node
        for p in paths:
            it = iter(p)
            n = next(it, None)
            while n is not None:
                if n in include:
                    break
                if n == first_common:
                    include.add(n)
                    break
                include.add(n)
                n = next(it, None)

        # In cases where there are even more distal common ancestors
        # (first common will typically be a branch point)
        this_ss = ss & cc
        if this_ss - include:
            # Make sure the new root is set correctly
            nr = sorted(this_ss - include, key=lambda x: longest_path.index(x))[-1]
            new_roots.append(nr)
            # Add those nodes to be included
            include = set.union(include, this_ss)
        else:
            new_roots.append(first_common)

    return np.array(list(include)), new_roots

def generate_list_of_childs(x: "core.NeuronObject") -> Dict[int, List[int]]:
    """Return list of childs.

    Parameters
    ----------
    x :     TreeNeuron | NeuronList
            If List, must contain a SINGLE neuron.

    Returns
    -------
    dict
        `{parent_id: [child_id, child_id, ...]}`

    """
    assert isinstance(x, core.TreeNeuron)
    # Grab graph once to avoid overhead from stale checks
    g = x.graph
    return {n: [e[0] for e in g.in_edges(n)] for n in g.nodes}

def node_label_sorting(
    x: "core.TreeNeuron", weighted: bool = False
) -> List[Union[str, int]]:
    """Return nodes ordered by node label sorting according to Cuntz
    et al., PLoS Computational Biology (2010).

    Parameters
    ----------
    x :         TreeNeuron
    weighted :  bool
                If True will use actual distances instead of just node count.
                Depending on how evenly spaced your points are, this might not
                make much difference.

    Returns
    -------
    list
        `[root, node_id, node_id, ...]`

    """
    if isinstance(x, core.NeuronList) and len(x) == 1:
        x = x[0]

    if not isinstance(x, core.TreeNeuron):
        raise TypeError(f'Expected a singleTreeNeuron, got "{type(x)}"')

    if len(x.root) > 1:
        raise ValueError("Unable to process multi-root neurons!")

    # Get relevant terminal nodes
    term = x.nodes[x.nodes.type == "end"].node_id.values

    # Get directed (!) distance from terminals to all other nodes
    geo = geodesic_matrix(
        x, from_=x.nodes[x.nodes.type.isin(("end", "root", "branch"))].node_id.values, directed=True, weight="weight" if weighted else None
    )
    # Set distance between unreachable points to None
    # Need to reinitialise SparseMatrix to replace float('inf') with NaN
    # dist_mat[geo == float('inf')] = None
    dist_mat = pd.DataFrame(
        np.where(
            geo == float("inf"),  # type: ignore  # no stubs for SparseDataFrame
            np.nan,
            geo,
        ),
        columns=geo.columns,
        index=geo.index,
    )

    # Get starting points (i.e. branches off the root) and sort by longest
    # path to a terminal (note we're operating on the simplified version
    # of the skeleton)
    G = graph.simplify_graph(x.graph)
    curr_points = sorted(
        list(G.predecessors(x.root[0])),
        key=lambda n: dist_mat[n].max() + dist_mat.loc[n, x.root[0]],
        reverse=True,
    )

    # Walk from root towards terminals, prioritising longer branches
    nodes_walked = []
    while curr_points:
        nodes_walked.append(curr_points.pop(0))
        # If the current point is a terminal point, stop here
        if nodes_walked[-1] in term:
            pass
        else:
            new_points = sorted(
                list(G.predecessors(nodes_walked[-1])),
                # Use distance to the farthest terminal + distance to current node as sorting key
                key=lambda n: dist_mat[n].max() + dist_mat.loc[n, nodes_walked[-1]],
                reverse=True,
            )
            curr_points = new_points + curr_points

    # Translate into segments
    node_list = [x.root[0:]]
    # Note that we're inverting here so that the segments are ordered
    # proximal -> distal (i.e. root to tips)
    seg_dict = {s[0]: s[::-1] for s in _break_segments(x)}

    for n in nodes_walked:
        # Note that we're skipping the first (proximal) node to avoid double
        # counting nodes
        node_list.append(seg_dict[n][1:])

    return np.concatenate(node_list, dtype=int)

def simplify_graph(G, inplace=False):
    """Simplify skeleton graph (networkX or igraph).

    This function will simplify the graph by keeping only roots, leafs and
    branch points. Preserves branch lengths (i.e. weights)!

    Parameters
    ----------
    G :         networkx.DiGraph | igraph.Graph
                The skeleton graph to simplify.
    inplace :   bool
                If True, will modify the graph in place.

    Returns
    -------
    G :         networkx.DiGraph | networkx.DiGraph
                Simplified graph.

    Examples
    --------
    >>> import navis
    >>> n = navis.example_neurons(1, kind='skeleton')
    >>> # Simplify skeleton's NetworkX graph representation
    >>> G_simp_nx = navis.graph.simplify_graph(n.graph)
    >>> # Check that we have the expected number of nodes
    >>> assert len(G_simp_nx.nodes) == (n.n_branches + n.n_root + n.n_leafs)
    >>> # Simplify skeleton's iGraph graph representation
    >>> G_simp_ig = navis.graph.simplify_graph(n.igraph)
    >>> # Check that we have the expected number of nodes
    >>> assert len(G_simp_ig.vs) == (n.n_branches + n.n_root + n.n_leafs)

    """
    if not inplace:
        G = G.copy()

    if isinstance(G, nx.Graph):
        # Find all leaf and branch points
        leafs = {n for n in G.nodes if G.in_degree(n) == 0 and G.out_degree(n) != 0}
        branches = {n for n in G.nodes if G.in_degree(n) > 1 and G.out_degree(n) != 0}
        roots = {n for n in G.nodes if G.out_degree(n) == 0}

        stop_nodes = roots | leafs | branches

        # Walk from each leaf/branch point to the next leaf, branch or root
        to_remove = []
        for start_node in leafs | branches:
            dist = 0
            node = start_node
            while True:
                parent = next(G.successors(node))
                dist += G.edges[node, parent]["weight"]

                if parent in stop_nodes:
                    G.add_weighted_edges_from([(start_node, parent, dist)])
                    break

                to_remove.append(parent)
                node = parent

        G.remove_nodes_from(to_remove)
    else:
        # Find all leaf and branch points
        leafs = G.vs.select(_indegree=0, _outdegree_ne=0)
        branches = G.vs.select(_indegree_gt=1, _outdegree_ne=0)
        roots = G.vs.select(_outdegree=0)

        stop_nodes = np.concatenate((roots.indices, leafs.indices, branches.indices))

        # Walk from each leaf/branch point to the next leaf, branch or root
        to_remove = []
        for start_node in np.concatenate((leafs.indices, branches.indices)):
            dist = 0
            node = start_node
            while True:
                parent = G.successors(node)[0]
                dist += G.es[G.get_eid(node, parent)]["weight"]

                if parent in stop_nodes:
                    G.add_edge(start_node, parent, weight=dist)
                    break

                to_remove.append(parent)
                node = parent

        G.delete_vertices(to_remove)

    return G

def skeleton_adjacency_matrix(
    x: "core.NeuronObject", sort: bool = True
) -> pd.DataFrame:
    """Generate adjacency matrix for a skeleton.

    Parameters
    ----------
    x :         TreeNeuron
                Neuron for which to generate adjacency matrix.
    sort :      bool, optional
                If True, will sort the adjacency matrix by topology.

    Returns
    -------
    pd.DataFrame
                Adjacency matrix where rows are nodes and columns are
                their parents.

    See Also
    --------
    [`navis.geodesic_matrix`][]
        For distances between all points.
    [`navis.distal_to`][]
        Check if a node A is distal to node B.
    [`navis.dist_between`][]
        Get point-to-point geodesic ("along-the-arbor") distances.

    """
    if isinstance(x, core.NeuronList):
        if len(x) == 1:
            x = x[0]
        else:
            raise ValueError("Cannot process more than a single neuron.")
    elif not isinstance(x, (core.TreeNeuron,)):
        raise ValueError(f'Unable to process data of type "{type(x)}"')

    # Generate the empty adjacency matrix
    adj = pd.DataFrame(
        np.zeros((len(x.nodes), len(x.nodes)), dtype=bool),
        index=x.nodes.node_id.values,
        columns=x.nodes.node_id.values,
    )

    # Fill in the parent-child relationships
    not_root = x.nodes.parent_id.values >= 0
    node_ix = np.arange(len(x.nodes))[not_root]
    parent_ids = x.nodes.parent_id.values[not_root]
    parent_ix = np.searchsorted(x.nodes.node_id.values, parent_ids)
    adj.values[node_ix, parent_ix] = True

    if sort:
        sort = node_label_sorting(x)
        adj = adj.loc[sort, sort]

    return adj