import networkx

from gamelib.lab import Lab


def construct_research_tree(lab):
    rt = networkx.DiGraph()
    for research in lab.new_research:
        rt.add_node(research, label=research.NAME, color='green')
    for research in lab.new_research:
        for prereq in research.PREREQUISITES:
            prefix = prereq[0].NAME[0:3]
            cur = prereq[0]
            for i in range(1, prereq[1] + 1):
                rt.add_node(prefix + str(i), color='green')
                rt.add_edge(cur, prefix + str(i), color='green')
                cur = prefix + str(i)
            rt.add_edge(cur, research, color='green')
    for schematic in lab.new_schematics:
        rt.add_node(schematic, label=schematic.NAME, color='pink')
    for schematic in lab.new_schematics:
        for prereq in schematic.PREREQUISITES:
            prefix = prereq[0].NAME[0:3]
            cur = prereq[0]
            for i in range(1, prereq[1] + 1):
                rt.add_node(prefix + str(i), color='pink')
                rt.add_edge(cur, prefix + str(i), color='pink')
                cur = prefix + str(i)
            rt.add_edge(cur, schematic, color='pink')
    return rt


def draw(rt):
    agraph = networkx.to_agraph(rt)
    agraph.graph_attr['label'] = 'Research Tree'
    agraph.write('research_tree.dot')


def main():
    lab = Lab({'science': {}})
    research_tree = construct_research_tree(lab)
    draw(research_tree)