# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_sampling.ipynb (unless otherwise specified).

__all__ = ['sample_node', 'sample_nodes', 'sample_edges', 'khop_subgraph', 'khop_neighbourhood']

# Cell
from nbdev.showdoc import *
import numpy as np
import networkx as nx
from .graphtools import non_pendant_edges

# Cell
def sample_node(G):
    """Uniformly samples a single node from G."""
    return int(np.random.choice(list(G.nodes())))

def sample_nodes(G, num_nodes=1):
    """Uniformly samples `num_nodes` nodes.

    Nodes are sampled uniformly at random without replacement and returned in a list.
    """
    nodes = G.nodes()
    return list(np.random.choice(list(nodes), num_nodes, replace=False))

# Cell
def sample_edges(G, r, non_pendant=False):
    """Samples `r` edges from `G`.

    Edges are sampled uniformly at random without replacement and returned as a list.

    If `non_pendant` is `True` then pendant edges will not be sampled.
    """
    if non_pendant:
        edges = non_pendant_edges(G)
    else:
        edges = list(G.edges)
    return [edges[i] for i in np.random.choice(range(len(edges)), r, replace=False)]

# Cell
def khop_subgraph(G, k, node=None):
    """Returns a subgraph of `G` which is the k-hop neighbourhood of some `node`.

    If `node=None` then a node is randomly sampled.
    """
    if node is None:
        node = sample_node(G)
    khop = khop_neighbourhood(G, k, node)
    return G.subgraph(khop), node

def khop_neighbourhood(G, k, node):
    """Returns the k-hop neighbourhood of `node`."""
    return list(nx.single_source_shortest_path_length(G, node, cutoff=k).keys())