import os

from pycallgraph import Config, GlobbingFilter, PyCallGraph
from pycallgraph.color import Color
from pycallgraph.output import GraphvizOutput


class CustomOutput(GraphvizOutput):
    def node_color(self, node):
        value = float(node.time.fraction)
        return Color.hsv(value / 2 + 0.5, value, 0.9)

    def edge_color(self, edge):
        value = float(edge.time.fraction)
        return Color.hsv(value / 2 + 0.5, value, 0.7)


from pypet import Trajectory, load_trajectory

filename = None


def to_test(traj, length):
    for irun in range(length):
        traj._add_run_info(irun)


def test_load():
    load_trajectory(index=-1, filename=filename, load_data=1)


if __name__ == "__main__":
    if not os.path.isdir("./tmp"):
        os.mkdir("tmp")
    graphviz = CustomOutput()
    graphviz.output_file = "./tmp/traj_add_run_info.png"
    service_filter = GlobbingFilter(
        include=[
            "*storageservice.*",
            "*ptcompat.*",
            "*naturalnaming.*",
            "*parameter.*",
            "*trajectory.*",
        ]
    )
    # service_filter = GlobbingFilter(include=['*naturalnaming.*', '*trajectory.*'])

    config = Config(groups=True, verbose=True)
    config.trace_filter = service_filter

    print("RUN PROFILE")
    with PyCallGraph(config=config, output=graphviz):
        to_test(Trajectory(), 10000)
    print("DONE RUN PROFILE")
