import networkx

from gamelib.lab import Lab


COLORS = {
    'schematic': 'green',
    'research': 'pink',
    }


def color(science):
    return COLORS[science.SCIENCE_TYPE]


def add_node(rt, science):
    rt.add_node(science, label=science.NAME, color=color(science))


def add_pre_node(rt, prior, science, points):
    name = "%s%s" % (science.NAME[:3], points)
    if name not in rt:
        rt.add_node(name, color=color(science))
        rt.add_edge(prior, name, color=color(science))
    return name


def add_prereqs(rt, science):
    for pre, points in science.PREREQUISITES:
        prior = pre
        for i in range(points):
            prior = add_pre_node(rt, prior, pre, i + 1)
        rt.add_edge(prior, science, color=color(science))


def construct_research_tree(lab):
    rt = networkx.DiGraph()
    sciences = lab.new_research + lab.new_schematics
    for science in sciences:
        add_node(rt, science)
    for science in sciences:
        add_prereqs(rt, science)
    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)