import os

import numpy as np
import scipy.sparse as spsp
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 Environment, Parameter, cartesian_product, load_trajectory
from pypet.tests.testutils.data import add_params, create_param_dict, simple_calculations
from pypet.tests.testutils.ioutils import make_temp_dir

filename = None


def explore(traj):
    explored = {
        "Normal.trial": [0],
        "Numpy.double": [np.array([1.0, 2.0, 3.0, 4.0]), np.array([-1.0, 3.0, 5.0, 7.0])],
        "csr_mat": [spsp.csr_matrix((2222, 22)), spsp.csr_matrix((2222, 22))],
    }

    explored["csr_mat"][0][1, 2] = 44.0
    explored["csr_mat"][1][2, 2] = 33

    traj.f_explore(cartesian_product(explored))


def test_run():

    global filename

    np.random.seed()
    trajname = "profiling"
    filename = make_temp_dir(os.path.join("hdf5", f"test{trajname}.hdf5"))

    env = Environment(
        trajectory=trajname,
        filename=filename,
        file_title=trajname,
        log_stdout=False,
        results_per_run=5,
        derived_parameters_per_run=5,
        multiproc=False,
        ncores=1,
        wrap_mode="LOCK",
        use_pool=False,
        overwrite_file=True,
    )

    traj = env.v_trajectory

    traj.v_standard_parameter = Parameter

    ## Create some parameters
    param_dict = {}
    create_param_dict(param_dict)
    ### Add some parameter:
    add_params(traj, param_dict)

    # remember the trajectory and the environment
    traj = traj
    env = env

    traj.f_add_parameter("TEST", "test_run")
    ###Explore
    explore(traj)

    ### Make a test run
    simple_arg = -13
    simple_kwarg = 13.0
    env.f_run(simple_calculations, simple_arg, simple_kwarg=simple_kwarg)

    size = os.path.getsize(filename)
    size_in_mb = size / 1000000.0
    print(f"Size is {size_in_mb}MB")


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/run_profile_traj_slots.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):
        test_run()
    print("DONE RUN PROFILE")

    graphviz = CustomOutput()
    graphviz.output_file = "./tmp/load_mode_1_profile_traj_slots.png"

    print("LOAD PROFILE")
    with PyCallGraph(config=config, output=graphviz):
        test_load()
    print("DONE LOAD PROFILE")
