#! /usr/bin/python3

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


def extract_neighborhood(g, center, depth):
    center = find_node(g, center)

    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)

    recurse(center, depth + 1)

    return g.subgraph(neighbors)


def vertex(string):
    try:
        key, value = string.split(":", 1)
    except ValueError:
        raise argparse.ArgumentTypeError(
            "key must be separated from value by a colon")
    return key, value

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(
        description='Extract the neighborhood (ancestors and descendants) '
        'around a node in a graph in GraphML format up to a certain depth')
    parser.add_argument("g", type=read_graphml, nargs='?', default='-',
                        help="Input graph in GraphML format (default: stdin)")
    parser.add_argument("h", type=write_graphml, nargs='?', default='-',
                        help="Output graph in GraphML format "
                        "(default: stdout)")
    parser.add_argument('--center', nargs='+', required=True, type=vertex,
                        help="key:value pairs to match the center vertex")
    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.g, args.center, args.depth)
    args.h(h)
