#! /usr/bin/python3

from __future__ import print_function
import sys
import os

sys.path.append("/usr/share/botch")
from util import read_graph, write_graph, find_nodes
import networkx as nx


def extract_scc(g):
    result = []
    for scc_nodes in nx.components.strongly_connected_components(g):
        scc = g.subgraph(scc_nodes)
        if len(scc) <= 1:
            continue
        if hasattr(g, "input_file_type"):
            scc.input_file_type = g.input_file_type
        result.append(scc)
    return result


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 strongly "
        + "connected components "
        + "(SCC) from input graph"
    )
    parser.add_argument(
        "graph",
        metavar="graph.xml",
        type=read_graph,
        help="graph in GraphML or dot format",
    )
    parser.add_argument(
        "prefix",
        metavar="prefix",
        nargs="?",
        default="scc",
        help="output graph prefix (default: " + "scc)",
    )
    parser.add_argument(
        "--outdir",
        metavar="dir",
        type=str,
        default=".",
        help="output directory. Default is " + "current directory",
    )
    parser.add_argument(
        "--outfnamevert",
        type=vertex,
        action="append",
        help="vertices to pick to generate output filename "
        "as key:value pairs. The special key __ID__ "
        "allows one to select the unique vertex id.",
    )
    parser.add_argument(
        "--outfnameattr",
        type=str,
        help="vertex attribute to pick to generate output" "file name",
    )
    parser.add_argument(
        "--extension", type=str, help="filename extension of the output files"
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="be verbose")
    args = parser.parse_args()
    res = extract_scc(args.graph)
    for scc in res:
        # get the source node with the lowest name
        if args.outfnamevert:
            fnameverts = find_nodes(scc, args.outfnamevert)
        else:
            fnameverts = nx.nodes(scc)
        if args.outfnameattr:
            m = min(scc.nodes[n][args.outfnameattr] for n in fnameverts)
        else:
            m = str(min(fnameverts))
        if args.extension:
            ext = args.extension
        else:
            ext = "xml"
        if args.prefix:
            fname = "%s_%s_%d.%s" % (args.prefix, m, len(scc), ext)
        else:
            fname = "%s_%d.%s" % (m, len(scc), ext)
        fullfname = os.path.join(args.outdir, fname)
        write_graph(fullfname)(scc)
        print(fname)
