#! /usr/bin/python3

import sys

sys.path.append("/usr/share/botch")
from util import read_graph, write_graph, 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)

    h = g.subgraph(neighbors)
    h.input_file_type = g.input_file_type
    return h


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 or dot format up to a certain "
        "depth"
    )
    parser.add_argument(
        "g",
        type=read_graph,
        nargs="?",
        default="-",
        help="Input graph in GraphML or dot format (default: " "stdin)",
    )
    parser.add_argument(
        "h",
        type=write_graph,
        nargs="?",
        default="-",
        help="Output graph in GraphML or dot format " "(default: stdout)",
    )
    parser.add_argument(
        "--center",
        required=True,
        type=vertex,
        action="append",
        help="key:value pairs to match the center vertex. "
        "The special key __ID__ allows one to select "
        "the unique vertex identifier.",
    )
    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)
