import networkx as nx

"""
This script processes files produced by pyreverse.
Get pyreverse by installing pylint: ``pip install pylint``.

The particular files we process are generated by a command like
    ``pyreverse  -k -d <path> -o plantuml pygsti``
where <path> is the relative path to the directory that should contain
the output files.

Notes
-----
 1. The command above will create two files: "classes.plantuml" and "packages.plantuml".
    This script is only concerned with the first such file.

 2. The classes.plantuml file contains data for different types of class relations.
    Most of the data concerns class inheritance. But there's also data on which classes
    are used as members within other classes. At present, this script only tries to
    extract inheritance structure.
"""


def plantuml_classes_to_networkx(plantuml_filename, prefix='pygsti'):
    if prefix[-1] != '.':
        prefix = prefix + '.'
    fi = open(plantuml_filename, 'r')

    def strip_readline():
        return fi.readline().removesuffix('\n').removesuffix('\r')

    def node_test(tup):
        return len(tup) == 5 and tup[0] == 'class' and tup[2] == 'as' and tup[4] == '{'

    def edge_test(tup):
        return len(tup) == 3 and prefix in tup[0] and prefix in tup[2]

    # initial file format check
    tup = strip_readline().split(' ')
    if len(tup) != 2 or tup[0] != '@startuml' or tup[1] != 'classes':
        raise ValueError('Unexpected file format.')
    fi.readline()  # toss

    # read nodes
    nodes = set()
    while True:
        tup = strip_readline().split(' ')
        if node_test(tup):
            nodes.add(tup[3].removeprefix(prefix))
            ell = strip_readline()
            if ell != '}':
                raise ValueError('Unexpected file format.')
        else:
            if edge_test(tup):
                break
            else:
                raise ValueError('Unexpected file format.')

    # read inheritance edges
    edges = set()
    while tup[1] == '--|>':
        edges.add((tup[0].removeprefix(prefix), tup[2].removeprefix(prefix)))
        tup = strip_readline().split(' ')

    fi.close()

    G = nx.DiGraph()
    G.add_nodes_from(nodes)
    G.add_edges_from(edges)
    return G


if __name__ == '__main__':
    # We assume this script is run from a directory that contains "classes.plantml".
    # We save the class inheritance graph in GML format. This format is widely
    # used, and can be read by the free software called "yEd".
    G = plantuml_classes_to_networkx('classes.plantuml')
    nx.write_gml(G, 'pygsti.gml')
    # Note: once the graph is loaded into yEd it can be saved as a richer format
    # like GraphML (that records visual positioning data). You can also delete
    # nodes if they seem conceptually unimportant.
