import networkx

import matplotlib.pyplot as plt

import pylab

from gamelib import research
from gamelib import schematics
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:
            rt.add_edge(prereq[0], research, weight=1.0 / prereq[1], \
                label=prereq[1], color='green')
    for schematic in lab.new_schematics:
        rt.add_node(schematic, label=schematic.NAME, color='red')
    for schematic in lab.new_schematics:
        for prereq in schematic.PREREQUISITES:
            rt.add_edge(prereq[0], schematic, weight=1.0 / prereq[1], \
                label=prereq[1], color='red', size=schematic.COST/10)
    return rt


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


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