#!/usr/bin/env python3

import sys
sys.path.append('/usr/share/botch')
from util import read_graphml, write_graphml


def extract_neighborhood(g, nodename, depth):
    matching = [n for n, attrdict in list(
        g.node.items()) if attrdict['name'] == nodename]

    neighbors = set()

    def recurse(node, depth):
        if depth == 0:
            return
        if node in neighbors:
            return
        neighbors.add(node)
        for n in g.predecessors(node):
            recurse(n, depth - 1)
        for n in g.successors(node):
            recurse(n, depth - 1)

    for n in matching:
        recurse(n, depth + 1)

    return g.subgraph(neighbors)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(
        description='Extract the neighborhood around a node')
    parser.add_argument('ingraph', type=read_graphml, help='input graph')
    parser.add_argument('outgraph', type=write_graphml, help='output graph')
    parser.add_argument('nodename', help="name of node around which to " +
                                         "extract the neighborhood")
    parser.add_argument('--depth', type=int, default=1,
                        help='range of neighborhood to extract (Default: 1)')
    parser.add_argument(
        '-v', '--verbose', action='store_true', help='be verbose')
    args = parser.parse_args()
    h = extract_neighborhood(args.ingraph, args.nodename, args.depth)
    args.outgraph(h)
