import sys
from gprof2dot import gprof2dot
import pygraphviz as pgv
import os.path

PRUNE = True
DEFAULT_NODE_TRESHOLD = 0.5
DEFAULT_EDGE_TRESHOLD = 0.1

if __name__ == '__main__':
    # pstats -> dot
    if len(sys.argv) == 1:
        print >> sys.stderr, 'Invalid arguments!'
        sys.exit(1)
    pstatsfile = sys.argv[1]
    dotfilebasename, _ = os.path.splitext(pstatsfile)
    dotfile = dotfilebasename + '.dot'

    parser = gprof2dot.PstatsParser(pstatsfile)
    profile = parser.parse()

    with open(dotfile, 'wt') as f:
        dot = gprof2dot.DotWriter(f)

        if PRUNE:
            node_treshold = DEFAULT_NODE_TRESHOLD
            edge_treshold = DEFAULT_EDGE_TRESHOLD
        else:
            node_treshold = 0
            edge_treshold = 0

        profile.prune(node_treshold / 100, edge_treshold / 100)
        dot.graph(profile, gprof2dot.TEMPERATURE_COLORMAP)

    # dot -> png
    graph = pgv.AGraph(dotfile)
    graph.layout('dot')
    pngfilebasename, _ = os.path.splitext(dotfile)
    pngfilename = pngfilebasename + '.png'
    graph.draw(pngfilename)