#!/usr/bin/python3

import networkx as nx
import sys

sys.path.append("/usr/share/botch")
from util import write_graph, read_graph, find_node


def shortest_path(g, source, target, findall=False):
    source = find_node(g, source)
    target = find_node(g, target)
    # networkx has the method shortest_path which only returns one of the
    # shortest paths. Instead, we always retrieve all shortest paths because
    # the method networkx uses produces different output across individual
    # invocations
    paths = nx.all_shortest_paths(g, source, target)
    if findall:
        nodes = set([v for path in paths for v in path])
    else:
        nodes = sorted(paths)[0]
    h = g.subgraph(nodes)
    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="find the shortest path(s) between two vertices of a "
        "graph in GraphML or dot format"
    )
    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(
        "--source",
        required=True,
        type=vertex,
        action="append",
        help="key:value pairs to match the source vertex. "
        "The special key __ID__ allows one to select "
        "the unique vertex identifier.",
    )
    parser.add_argument(
        "--target",
        required=True,
        type=vertex,
        action="append",
        help="key:value pairs to match the target vertex. "
        "The special key __ID__ allows one to select "
        "the unique vertex identifier.",
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="instead of finding an arbitrary shortest path, f"
        "ind all shortest paths",
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="be verbose")
    args = parser.parse_args()
    h = shortest_path(args.g, args.source, args.target, args.all)
    args.h(h)
