from modelspec.utils import evaluate
from modelspec.utils import save_to_json_file
from modelspec.utils import locate_file
from modelspec.utils import parse_list_like

from neuromllite.utils import load_network
from neuromllite.utils import print_v
from neuromllite.utils import get_pops_vs_cell_indices_seg_ids


import numpy as np
import os
import random

import lems.api as lems  # from pylems

DEFAULT_NET_GENERATION_SEED = 1234
DEFAULT_SIMULATION_SEED = 5678



def __TEMP_CHECK_IF_SET(value):
    """
    Should just be a check for None! Update in modelspec...
    """
    if value is None: return False
    if value == '': return False
    return True

def _get_rng_for_network(nl_model):
    """
    Get a random number generator and the seed used to generate it
    """
    seed = (
        nl_model.seed
        if nl_model and nl_model.seed is not None
        else DEFAULT_NET_GENERATION_SEED
    )
    rng = random.Random(seed)

    return rng, seed


def generate_network(
    nl_model,
    handler,
    always_include_props=False,
    include_connections=True,
    include_inputs=True,
    base_dir=None,
):
    """
    Generate the network model as described in NeuroMLlite in a specific handler,
    e.g. NeuroMLHandler, PyNNHandler, etc.
    """

    pop_locations = {}
    cell_objects = {}
    synapse_objects = {}

    print_v(
        "Starting net generation for %s%s..."
        % (nl_model.id, " (base dir: %s)" % base_dir if base_dir else "")
    )

    rng, seed = _get_rng_for_network(nl_model)

    if nl_model.network_reader:

        exec(
            "from neuromllite.%s import %s"
            % (nl_model.network_reader.type, nl_model.network_reader.type)
        )

        network_reader = eval("%s()" % (nl_model.network_reader.type))
        network_reader.parameters = nl_model.network_reader.parameters

        network_reader.parse(handler)
        pop_locations = network_reader.get_locations()

    else:
        from neuromllite import __version__ as nmlv

        notes = "Generated by NeuroMLlite v%s" % nmlv
        notes += "\n    Generated network: %s" % nl_model.id
        notes += "\n    Generation seed: %i" % (seed)
        if nl_model.parameters:
            notes += "\n    NeuroMLlite parameters: "
            for p in sorted(nl_model.parameters.keys()):
                if not p == "__builtins__":
                    notes += "\n        %s = %s" % (p, nl_model.parameters[p])
        handler.handle_document_start(nl_model.id, notes)
        temperature = "%sdegC" % nl_model.temperature if nl_model.temperature is not None else None
        handler.handle_network(nl_model.id, nl_model.notes, temperature=temperature)

    nml2_doc_temp = _extract_pynn_components_to_neuroml(nl_model)

    for c in nl_model.cells:

        if c.neuroml2_source_file:
            from pyneuroml import pynml

            nml2_doc = pynml.read_neuroml2_file(
                locate_file(c.neuroml2_source_file, base_dir), include_includes=True
            )
            cell_objects[c.id] = nml2_doc.get_by_id(c.id)

        if c.pynn_cell:
            cell_objects[c.id] = nml2_doc_temp.get_by_id(c.id)

    for s in nl_model.synapses:
        if s.neuroml2_source_file:
            from pyneuroml import pynml

            nml2_doc = pynml.read_neuroml2_file(
                locate_file(s.neuroml2_source_file, base_dir), include_includes=True
            )
            synapse_objects[s.id] = nml2_doc.get_by_id(s.id)

        if hasattr(s, 'pynn_synapse') and s.pynn_synapse is not None:
            synapse_objects[s.id] = nml2_doc_temp.get_by_id(s.id)

    for p in nl_model.populations:

        size = evaluate(p.size, nl_model.parameters)
        properties = p.properties if p.properties else {}

        if p.random_layout:
            properties["region"] = p.random_layout.region
        if p.relative_layout:
            properties["region"] = p.relative_layout.region

        if (
            not p.random_layout
            and not p.single_location
            and not p.relative_layout
            and not always_include_props
        ):

            # If there are no positions (abstract network), and <property>
            # is added to <population>, jLems doesn't like it... (it has difficulty
            # interpreting pop0[0]/v, etc.)
            # So better not to give properties...
            properties = {}

        if p.notes:
            handler.handle_population(
                p.id,
                p.component,
                size,
                cell_objects[p.component] if p.component in cell_objects else None,
                properties=properties,
                notes=p.notes,
            )
        else:
            handler.handle_population(
                p.id,
                p.component,
                size,
                cell_objects[p.component] if p.component in cell_objects else None,
                properties=properties,
            )

        pop_locations[p.id] = np.zeros((size, 3))

        for i in range(size):
            if p.random_layout:
                region = nl_model.get_child(p.random_layout.region, "regions")

                x = region.x + rng.random() * region.width
                y = region.y + rng.random() * region.height
                z = region.z + rng.random() * region.depth
                pop_locations[p.id][i] = (x, y, z)

                handler.handle_location(i, p.id, p.component, x, y, z)

            if p.single_location:
                loc = p.single_location.location
                x = loc.x
                y = loc.y
                z = loc.z
                pop_locations[p.id][i] = (x, y, z)

                handler.handle_location(i, p.id, p.component, x, y, z)

            if p.relative_layout:
                print_v("Generating population with layout: %s" % p.relative_layout)
                region = nl_model.get_child(p.relative_layout.region, "regions")
                x = p.relative_layout.x + region.x
                y = p.relative_layout.y + region.y
                z = p.relative_layout.z + region.z
                pop_locations[p.id][i] = (x, y, z)
                handler.handle_location(i, p.id, p.component, x, y, z)

        if hasattr(handler, "finalise_population"):
            handler.finalise_population(p.id)

    if include_connections:
        for p in nl_model.projections:

            type = p.type if p.type else "projection"

            synapse_obj = (
                synapse_objects[p.synapse] if p.synapse in synapse_objects else None
            )
            pre_synapse_obj = (
                synapse_objects[p.pre_synapse]
                if p.pre_synapse in synapse_objects
                else None
            )

            delay = evaluate(p.delay, nl_model.parameters, rng) if p.delay else 0
            weight = evaluate(p.weight, nl_model.parameters, rng) if p.weight else 1

            if weight != 0:
                handler.handle_projection(
                    p.id,
                    p.presynaptic,
                    p.postsynaptic,
                    p.synapse,
                    synapse_obj=synapse_obj,
                    pre_synapse_obj=pre_synapse_obj,
                    type=type,
                )

                if pre_synapse_obj is None and p.pre_synapse == "silentSyn_%s" % p.id:
                    if hasattr(handler, "nml_doc"):
                        for ss in handler.nml_doc.silent_synapses:
                            if ss.id == "silentSyn_%s" % p.id:
                                handler.nml_doc.silent_synapses.remove(
                                    ss
                                )  # A default silent synapse was generated by libNeuroML

                conn_count = 0
                if p.random_connectivity:
                    for pre_i in range(len(pop_locations[p.presynaptic])):
                        for post_i in range(len(pop_locations[p.postsynaptic])):
                            flip = rng.random()
                            # print("Is cell %i conn to %i, prob %s - %s"%(pre_i, post_i, flip, p.random_connectivity.probability))
                            if flip < evaluate(
                                p.random_connectivity.probability,
                                nl_model.parameters,
                                rng,
                            ):

                                delay = (
                                    evaluate(p.delay, nl_model.parameters, rng)
                                    if p.delay
                                    else 0
                                )
                                weight = (
                                    evaluate(p.weight, nl_model.parameters, rng)
                                    if p.weight
                                    else 1
                                )
                                # print_v("Adding connection %i with weight: %s, delay: %s"%(conn_count, weight, delay))
                                handler.handle_connection(
                                    p.id,
                                    conn_count,
                                    p.presynaptic,
                                    p.postsynaptic,
                                    p.synapse,
                                    pre_i,
                                    post_i,
                                    preSegId=0,
                                    preFract=0.5,
                                    postSegId=0,
                                    postFract=0.5,
                                    delay=delay,
                                    weight=weight,
                                )
                                conn_count += 1

                if p.convergent_connectivity:

                    for post_i in range(len(pop_locations[p.postsynaptic])):

                        for count in range(int(p.convergent_connectivity.num_per_post)):
                            found = False
                            while not found:
                                pre_i = int(
                                    rng.random() * len(pop_locations[p.presynaptic])
                                )
                                if p.presynaptic == p.postsynaptic and pre_i == post_i:
                                    found = False
                                else:
                                    found = True

                            print_v(
                                "Adding connection %i (%i->%i; %i to %s of post) with weight: %s, delay: %s"
                                % (
                                    conn_count,
                                    pre_i,
                                    post_i,
                                    count,
                                    p.convergent_connectivity.num_per_post,
                                    weight,
                                    delay,
                                )
                            )
                            handler.handle_connection(
                                p.id,
                                conn_count,
                                p.presynaptic,
                                p.postsynaptic,
                                p.synapse,
                                pre_i,
                                post_i,
                                preSegId=0,
                                preFract=0.5,
                                postSegId=0,
                                postFract=0.5,
                                delay=delay,
                                weight=weight,
                            )
                            conn_count += 1

                elif p.one_to_one_connector:
                    for i in range(
                        min(
                            len(pop_locations[p.presynaptic]),
                            len(pop_locations[p.postsynaptic]),
                        )
                    ):

                        # print_v("Adding connection %i with weight: %s, delay: %s"%(conn_count, weight, delay))
                        handler.handle_connection(
                            p.id,
                            conn_count,
                            p.presynaptic,
                            p.postsynaptic,
                            p.synapse,
                            i,
                            i,
                            preSegId=0,
                            preFract=0.5,
                            postSegId=0,
                            postFract=0.5,
                            delay=delay,
                            weight=weight,
                        )
                        conn_count += 1

                handler.finalise_projection(
                    p.id, p.presynaptic, p.postsynaptic, p.synapse
                )

    if include_inputs:
        for input in nl_model.inputs:

            handler.handle_input_list(
                input.id,
                input.population,
                input.input_source,
                size=0,
                input_comp_obj=None,
            )

            input_count = 0

            if __TEMP_CHECK_IF_SET(input.cell_ids):
                if __TEMP_CHECK_IF_SET(input.percentage):
                    raise Exception(
                        "On the input: %s, only one of percentage or cell_ids is allowed"
                        % input
                    )
                for i in input.cell_ids:
                    weight = input.weight if input.weight else 1

                    if input.number_per_cell and input.segment_ids:
                        raise Exception(
                            "On input: %s, only one of number_per_cell or segment_ids is allowed"
                            % input
                        )

                    if input.number_per_cell:
                        number_per_cell = evaluate(
                            input.number_per_cell, nl_model.parameters
                        )
                        seg_ids = [0] * number_per_cell
                    elif input.segment_ids:
                        seg_ids = parse_list_like(input.segment_ids)
                    else:
                        seg_ids = [0]

                    for seg_id in seg_ids:
                        handler.handle_single_input(
                            input.id,
                            input_count,
                            i,
                            segId=seg_id,
                            weight=evaluate(weight, nl_model.parameters),
                        )
                        input_count += 1

            if __TEMP_CHECK_IF_SET(input.percentage):
                if __TEMP_CHECK_IF_SET(input.cell_ids):
                    raise Exception(
                        "On input: %s, only one of percentage or cell_ids is allowed"
                        % input
                    )
                for i in range(len(pop_locations[input.population])):
                    flip = rng.random()
                    weight = input.weight if input.weight else 1
                    if flip * 100.0 < input.percentage:

                        if input.number_per_cell and input.segment_ids:
                            raise Exception(
                                "On input: %s, only one of number_per_cell or segment_ids is allowed"
                                % input
                            )

                        if input.number_per_cell:
                            number_per_cell = evaluate(
                                input.number_per_cell, nl_model.parameters
                            )
                            seg_ids = [0] * number_per_cell
                        elif input.segment_ids:
                            seg_ids = parse_list_like(input.segment_ids)
                        else:
                            seg_ids = [0]

                        for seg_id in seg_ids:
                            handler.handle_single_input(
                                input.id,
                                input_count,
                                i,
                                segId=seg_id,
                                weight=evaluate(weight, nl_model.parameters),
                            )
                            input_count += 1

            handler.finalise_input_source(input.id)

    if hasattr(handler, "finalise_document"):
        handler.finalise_document()
    else:
        pass


def check_to_generate_or_run(argv, sim):
    """
    Useful method for calling in main method after network and simulation are
    generated, to handle some standard export options like -jnml, -graph etc.
    """

    print_v(
        "Checking arguments: %s to see whether anything should be run in simulation %s (net: %s)..."
        % (argv, sim.id, sim.network)
    )

    if len(argv) == 1:
        print_v("No arguments found. Currently supported export formats:")
        print_v(
            "   -nml |  -nmlh5 | -jnml | -jnmlnrn | -jnmlnetpyne | -netpyne | -pynnnrn "
            + "| -pynnnest | -pynnbrian | -pynnneuroml | -sonata | -matrix[1-2] | -graph[1-6 n/d/f/c]"
        )

    if "-pynnnest" in argv:
        generate_and_run(sim, simulator="PyNN_NEST")

    elif "-pynnnrn" in argv:
        generate_and_run(sim, simulator="PyNN_NEURON")

    elif "-pynnbrian" in argv:
        generate_and_run(sim, simulator="PyNN_Brian")

    elif "-arbor" in argv:
        generate_and_run(sim, simulator="Arbor")

    elif "-bindsnet" in argv:
        generate_and_run(sim, simulator="BindsNET")

    elif "-jnml" in argv:
        generate_and_run(sim, simulator="jNeuroML")

    elif "-jnmlnrn" in argv:
        generate_and_run(sim, simulator="jNeuroML_NEURON")

        """
        work in progress...
    elif '-jnmlbrian2' in argv:
        generate_and_run(sim, simulator='jNeuroML_Brian2')
"""

    elif "-netpyne" in argv:
        generate_and_run(sim, simulator="NetPyNE")

    elif "-pynnneuroml" in argv:
        generate_and_run(sim, simulator="PyNN_NeuroML")

    elif "-sonata" in argv:
        generate_and_run(sim, simulator="sonata")

    elif "-pnl" in argv or "-psyneulink" in argv:
        generate_and_run(sim, simulator="psyneulink")

    elif "-mdf" in argv:
        generate_and_run(sim, simulator="mdf")

    elif "-nml" in argv or "-neuroml" in argv:

        network = load_network(sim.network)
        generate_neuroml2_from_network(network, simulation=sim, validate=True)

    elif "-nmlh5" in argv or "-neuromlh5" in argv:

        network = load_network(sim.network)
        generate_neuroml2_from_network(
            network, simulation=sim, validate=True, format="hdf5"
        )

    else:
        for a in argv:

            if "-jnmlnetpyne" in a:
                num_processors = 1
                if len(a) > len("-jnmlnetpyne"):
                    num_processors = int(a[12:])
                generate_and_run(
                    sim, simulator="jNeuroML_NetPyNE", num_processors=num_processors
                )
            elif "graph" in a:  # e.g. -graph3c
                show_graph = True
                if "-nogui" in argv:
                    show_graph = False
                generate_and_run(sim, simulator=a[1:], nogui= not show_graph)  # Will not "run" obviously...
            elif "matrix" in a:  # e.g. -matrix2
                generate_and_run(sim, simulator=a[1:])  # Will not "run" obviously...


def _extract_pynn_components_to_neuroml(nl_model, nml_doc=None):
    """
    Parse the NeuroMLlite description for cell, synapses and inputs described as
    PyNN elements (e.g. IF_cond_alpha, DCSource) and parameters, and convert
    these to the equivalent elements in a NeuroMLDocument
    """

    import copy

    if nml_doc == None:
        from neuroml import NeuroMLDocument

        nml_doc = NeuroMLDocument(id="temp")

    for c in nl_model.cells:
        if c.pynn_cell:

            if nml_doc.get_by_id(c.id) == None:
                import pyNN.neuroml

                cell_params = copy.deepcopy(c.parameters) if c.parameters else {}
                for p in cell_params:
                    cell_params[p] = evaluate(cell_params[p], nl_model.parameters)
                for proj in nl_model.projections:

                    synapse = nl_model.get_child(proj.synapse, "synapses")
                    post_pop = nl_model.get_child(proj.postsynaptic, "populations")
                    if post_pop.component == c.id:
                        # print("--------- Cell %s in post pop %s of %s uses %s"%(c.id,post_pop.id, proj.id, synapse))

                        if synapse.pynn_receptor_type == "excitatory":
                            post = "_E"
                        elif synapse.pynn_receptor_type == "inhibitory":
                            post = "_I"
                        for p in synapse.parameters:
                            cell_params["%s%s" % (p, post)] = evaluate(
                                synapse.parameters[p], nl_model.parameters
                            )

                temp_cell = eval("pyNN.neuroml.%s(**cell_params)" % c.pynn_cell)
                if c.pynn_cell != "SpikeSourcePoisson":
                    temp_cell.default_initial_values["v"] = temp_cell.parameter_space[
                        "v_rest"
                    ].base_value

                cell_id = temp_cell.add_to_nml_doc(nml_doc, None)
                cell = nml_doc.get_by_id(cell_id)
                cell.id = c.id

    for s in nl_model.synapses:
        if nml_doc.get_by_id(s.id) == None:

            if s.pynn_synapse_type and s.pynn_receptor_type:
                import neuroml

                tau_syn = evaluate(s.parameters["tau_syn"], nl_model.parameters)
                if "e_rev" in s.parameters:
                    e_rev = evaluate(s.parameters["e_rev"], nl_model.parameters)
                if s.pynn_synapse_type == "cond_exp":
                    syn = neuroml.ExpCondSynapse(id=s.id, tau_syn=tau_syn, e_rev=e_rev)
                    nml_doc.exp_cond_synapses.append(syn)
                elif s.pynn_synapse_type == "cond_alpha":
                    syn = neuroml.AlphaCondSynapse(
                        id=s.id, tau_syn=tau_syn, e_rev=e_rev
                    )
                    nml_doc.alpha_cond_synapses.append(syn)
                elif s.pynn_synapse_type == "curr_exp":
                    syn = neuroml.ExpCurrSynapse(id=s.id, tau_syn=tau_syn)
                    nml_doc.exp_curr_synapses.append(syn)
                elif s.pynn_synapse_type == "curr_alpha":
                    syn = neuroml.AlphaCurrSynapse(id=s.id, tau_syn=tau_syn)
                    nml_doc.alpha_curr_synapses.append(syn)

    for i in nl_model.input_sources:

        if i.pynn_input:
            import pyNN.neuroml

            input_params = copy.deepcopy(i.parameters) if i.parameters else {}
            for ip in input_params:
                input_params[ip] = evaluate(input_params[ip], nl_model.parameters)
            temp_input = eval("pyNN.neuroml.%s(**input_params)" % (i.pynn_input))
            # print('%s'%(['%s->%s'%(p.id,p.amplitude) for p in nml_doc.pulse_generators]))
            pg_id = temp_input.add_to_nml_doc(nml_doc, None)
            # print('>> %s ±   %s; %s; %s'%(i,temp_input,pg_id,['%s->%s'%(p.id,p.amplitude) for p in nml_doc.pulse_generators]))
            pg = nml_doc.get_by_id(pg_id)
            pg.id = i.id
            # print('%s'%(['%s->%s'%(p.id,p.amplitude) for p in nml_doc.pulse_generators]))

    return nml_doc


def generate_neuroml2_from_network(
    nl_model,
    nml_file_name=None,
    print_summary=True,
    format="xml",
    base_dir=None,
    copy_included_elements=False,
    target_dir=None,
    validate=False,
    simulation=None,
):
    """
    Generate and save NeuroML2 file (in either XML or HDF5 format) from the
    NeuroMLlite description
    """

    print_v(
        "Generating NeuroML2 for %s%s..."
        % (
            nl_model.id,
            " (base dir: %s; target dir: %s)" % (base_dir, target_dir)
            if base_dir or target_dir
            else "",
        )
    )

    import neuroml
    from neuroml.hdf5.NetworkBuilder import NetworkBuilder

    neuroml_handler = NetworkBuilder()

    generate_network(nl_model, neuroml_handler, base_dir=base_dir)

    nml_doc = neuroml_handler.get_nml_doc()

    extra_lems_components = lems.Model()
    extra_lems_file = "%s__lems.xml" % nl_model.id

    if simulation is not None:
        if simulation.dt is not None:
            nml_doc.networks[0].properties.append(
                neuroml.Property("recommended_dt_ms", simulation.dt)
            )
        if simulation.duration is not None:
            nml_doc.networks[0].properties.append(
                neuroml.Property("recommended_duration_ms", simulation.duration)
            )

    for i in nl_model.input_sources:

        if i.lems_source_file:
            fname = locate_file(i.lems_source_file, base_dir)

            print_v(
                "Need to use for this input parameters: %s in %s from %s (%s); will place in %s"
                % (i.parameters, i.id, i.lems_source_file, fname, extra_lems_file)
            )
            model = lems.Model()
            model.import_from_file(fname)
            for comp in model.components:
                if i.id == comp.id:
                    print_v("Found a component: %s in %s" % (comp, fname))
                    if hasattr(i, 'parameters') and len(i.parameters) > 0:
                        for p in i.parameters:
                            comp.set_parameter(
                                p, evaluate(i.parameters[p], nl_model.parameters)
                            )
                    extra_lems_components.add(comp)

            for d in model.dimensions:
                print_v("Found a dimension: %s in %s" % (d, fname))
                extra_lems_components.add(d)
            for u in model.units:
                print_v("Found a unit: %s in %s" % (u, fname))
                extra_lems_components.add(u)
            for ct in model.component_types:
                print_v("Found a component type: %s in %s" % (ct, fname))
                extra_lems_components.add(ct)

            for inc in model.includes:
                print_v("Found an include: %s in %s" % (inc, fname))
                extra_lems_components.add(inc)

            incl = neuroml.IncludeType(extra_lems_file)
            if not incl in nml_doc.includes:
                nml_doc.includes.append(incl)

        if nml_doc.get_by_id(i.id) == None:
            if i.neuroml2_source_file:

                incl = neuroml.IncludeType(
                    locate_file(i.neuroml2_source_file, base_dir)
                )
                if not incl in nml_doc.includes:
                    nml_doc.includes.append(incl)

            if hasattr(i, 'neuroml2_input') and i.neuroml2_input is not None:
                input_params = i.parameters if i.parameters else {}

                # TODO make more generic...

                if i.neuroml2_input.lower() == "pulsegenerator":
                    input = neuroml.PulseGenerator(id=i.id)
                    nml_doc.pulse_generators.append(input)

                elif i.neuroml2_input.lower() == "pulsegeneratordl":
                    input = neuroml.PulseGeneratorDL(id=i.id)
                    nml_doc.pulse_generator_dls.append(input)

                elif i.neuroml2_input.lower() == "poissonfiringsynapse":
                    input = neuroml.PoissonFiringSynapse(id=i.id)
                    nml_doc.poisson_firing_synapses.append(input)

                for p in input_params:
                    exec(
                        'input.%s = "%s"'
                        % (p, evaluate(input_params[p], nl_model.parameters))
                    )

    for c in nl_model.cells:
        if c.neuroml2_source_file:

            incl = neuroml.IncludeType(locate_file(c.neuroml2_source_file, base_dir))
            found_cell = False
            for cell in nml_doc.cells:
                if cell.id == c.id:
                    nml_doc.cells.remove(
                        cell
                    )  # Better to use imported cell file; will have channels, etc.
                    nml_doc.includes.append(incl)
                    found_cell = True

            if not found_cell:
                for p in nl_model.populations:
                    if p.component == c.id:
                        pass

            if nml_doc.get_by_id(c.id) == None:  # cell (of any type) is not in nml_doc
                if not incl in nml_doc.includes:
                    nml_doc.includes.append(incl)

        if c.lems_source_file:
            fname = locate_file(c.lems_source_file, base_dir)

            print_v(
                "Need to use for this cell parameters: %s in %s from %s (%s); will place in %s"
                % (c.parameters, c.id, c.lems_source_file, fname, extra_lems_file)
            )
            model = lems.Model()
            model.import_from_file(fname)
            for comp in model.components:
                if c.id == comp.id:
                    print_v("Found component: %s in %s" % (comp, fname))
                    if c.parameters is not None and len(c.parameters) > 0:
                        for p in c.parameters:
                            comp.set_parameter(
                                p, evaluate(c.parameters[p], nl_model.parameters)
                            )
                    extra_lems_components.add(comp)

            for d in model.dimensions:
                print_v("Found a dimension: %s in %s" % (d, fname))
                extra_lems_components.add(d)
            for u in model.units:
                print_v("Found a unit: %s in %s" % (u, fname))
                extra_lems_components.add(u)
            for ct in model.component_types:
                print_v("Found component type: %s in %s" % (ct, fname))
                extra_lems_components.add(ct)
            for inc in model.includes:
                print_v("Found include: %s in %s" % (inc, fname))
                extra_lems_components.add(inc)

            incl = neuroml.IncludeType(extra_lems_file)
            if not incl in nml_doc.includes:
                nml_doc.includes.append(incl)

        if c.neuroml2_cell:

            cell_params = c.parameters if c.parameters else {}
            # TODO make more generic...
            if c.neuroml2_cell.lower() == "spikegenerator":
                cell = neuroml.SpikeGenerator(id=c.id)
                nml_doc.spike_generators.append(cell)
            elif c.neuroml2_cell.lower() == "spikegeneratorpoisson":
                cell = neuroml.SpikeGeneratorPoisson(id=c.id)
                nml_doc.spike_generator_poissons.append(cell)
            elif c.neuroml2_cell.lower() == "spikegeneratorrefpoisson":
                cell = neuroml.SpikeGeneratorRefPoisson(id=c.id)
                nml_doc.spike_generator_ref_poissons.append(cell)
            elif c.neuroml2_cell.lower() == "izhikevich2007cell":
                cell = neuroml.Izhikevich2007Cell(id=c.id)
                nml_doc.izhikevich2007_cells.append(cell)
            elif c.neuroml2_cell.lower() == "iafcell":
                cell = neuroml.IafCell(id=c.id)
                nml_doc.iaf_cells.append(cell)
            elif c.neuroml2_cell.lower() == "fitzhughnagumo1969cell":
                cell = neuroml.FitzHughNagumo1969Cell(id=c.id)
                nml_doc.fitz_hugh_nagumo1969_cells.append(cell)
            else:
                raise Exception(
                    "The neuroml2_cell: %s is not yet supported here..."
                    % c.neuroml2_cell
                )

            for p in cell_params:
                exec(
                    'cell.%s = "%s"'
                    % (p, evaluate(cell_params[p], nl_model.parameters))
                )

    for s in nl_model.synapses:
        if nml_doc.get_by_id(s.id) == None:
            if s.neuroml2_source_file:
                incl = neuroml.IncludeType(
                    locate_file(s.neuroml2_source_file, base_dir)
                )
                if not incl in nml_doc.includes:
                    nml_doc.includes.append(incl)

            if s.lems_source_file:
                fname = locate_file(s.lems_source_file, base_dir)

                print_v(
                    "Need to use for this synapse parameters: %s in %s from %s (%s); will place in %s"
                    % (s.parameters, s.id, s.lems_source_file, fname, extra_lems_file)
                )
                model = lems.Model()
                model.import_from_file(fname)
                for comp in model.components:
                    if s.id == comp.id:
                        print_v("Found component: %s in %s" % (comp, fname))
                        if s.parameters is not None and len(s.parameters) > 0:
                            for p in s.parameters:
                                comp.set_parameter(
                                    p, evaluate(s.parameters[p], nl_model.parameters)
                                )
                        extra_lems_components.add(comp)

                for d in model.dimensions:
                    print_v("Found a dimension: %s in %s" % (d, fname))
                    extra_lems_components.add(d)
                for u in model.units:
                    print_v("Found a unit: %s in %s" % (u, fname))
                    extra_lems_components.add(u)
                for ct in model.component_types:
                    print_v("Found component type: %s in %s" % (ct, fname))
                    extra_lems_components.add(ct)
                for inc in model.includes:
                    print_v("Found include: %s in %s" % (inc, fname))
                    extra_lems_components.add(inc)

                incl = neuroml.IncludeType(extra_lems_file)
                if not incl in nml_doc.includes:
                    nml_doc.includes.append(incl)

    # Look for and add the PyNN based elements to the NeuroMLDocument
    _extract_pynn_components_to_neuroml(nl_model, nml_doc)

    if print_summary:
        # Print info
        print_v(nml_doc.summary())

    # Save to file
    if target_dir == None:
        target_dir = base_dir
    if format == "xml":
        if not nml_file_name:
            nml_file_name = locate_file("%s.net.nml" % nml_doc.id, target_dir)
        from neuroml.writers import NeuroMLWriter

        NeuroMLWriter.write(nml_doc, nml_file_name)

    if format == "hdf5":
        if not nml_file_name:
            nml_file_name = locate_file("%s.net.nml.h5" % nml_doc.id, target_dir)
        from neuroml.writers import NeuroMLHdf5Writer

        NeuroMLHdf5Writer.write(nml_doc, nml_file_name)

    print_v("Written NeuroML to %s" % nml_file_name)

    if len(extra_lems_components.components) > 0:
        extra_lems_components.export_to_file(extra_lems_file)

    if validate and format == "xml":

        from pyneuroml import pynml

        success = pynml.validate_neuroml2(nml_file_name, verbose_validate=False)
        if success:
            print_v("Generated file is valid NeuroML2!")
        else:
            print_v("Generated file is NOT valid NeuroML2!")

    return nml_file_name, nml_doc


locations_mods_loaded_from = []


def _generate_neuron_files_from_neuroml(network, verbose=False, dir_for_mod_files=None):
    """
    Generate NEURON hoc/mod files from the NeuroML files which are marked as
    included in the NeuroMLlite description; also compiles the mod files
    """

    print_v(
        "-------------   Generating NEURON files from NeuroML for %s (default dir: %s)..."
        % (network.id, dir_for_mod_files)
    )
    nml_src_files = []

    from neuroml import NeuroMLDocument
    import neuroml.writers as writers

    temp_nml_doc = NeuroMLDocument(id="temp")

    dirs_for_mod_files = []
    if dir_for_mod_files != None:
        dirs_for_mod_files.append(os.path.abspath(dir_for_mod_files))

    for c in network.cells:
        if c.neuroml2_source_file:
            nml_src_files.append(c.neuroml2_source_file)
            dir_for_mod = os.path.dirname(os.path.abspath(c.neuroml2_source_file))
            if not dir_for_mod in dirs_for_mod_files:
                dirs_for_mod_files.append(dir_for_mod)

    for s in network.synapses:
        if s.neuroml2_source_file:
            nml_src_files.append(s.neuroml2_source_file)
            dir_for_mod = os.path.dirname(os.path.abspath(s.neuroml2_source_file))
            if not dir_for_mod in dirs_for_mod_files:
                dirs_for_mod_files.append(dir_for_mod)

    for i in network.input_sources:
        if i.neuroml2_source_file:
            nml_src_files.append(i.neuroml2_source_file)
            dir_for_mod = os.path.dirname(os.path.abspath(i.neuroml2_source_file))
            if not dir_for_mod in dirs_for_mod_files:
                dirs_for_mod_files.append(dir_for_mod)

    temp_nml_doc = _extract_pynn_components_to_neuroml(network)

    summary = temp_nml_doc.summary()

    if "IF_" in summary:
        import tempfile

        temp_nml_file = tempfile.NamedTemporaryFile(
            delete=False, suffix=".nml", dir=dir_for_mod_files
        )
        print_v("Writing temporary NML file to: %s, summary: " % temp_nml_file.name)
        print_v(summary)

        writers.NeuroMLWriter.write(temp_nml_doc, temp_nml_file.name)
        nml_src_files.append(temp_nml_file.name)

    for f in nml_src_files:
        from pyneuroml import pynml

        print_v("Generating/compiling hoc/mod files for: %s" % f)
        pynml.run_lems_with_jneuroml_neuron(
            f, nogui=True, only_generate_scripts=True, compile_mods=True, verbose=False
        )

    for dir_for_mod_files in dirs_for_mod_files:
        if not dir_for_mod_files in locations_mods_loaded_from:
            print_v(
                "Generated NEURON code; loading mechanisms from %s (cwd: %s; already loaded: %s)"
                % (dir_for_mod_files, os.getcwd(), locations_mods_loaded_from)
            )
            try:

                from neuron import load_mechanisms

                if os.getcwd() == dir_for_mod_files:
                    print_v(
                        "That's current dir => importing neuron module loads mods here..."
                    )
                else:
                    load_mechanisms(dir_for_mod_files)
                locations_mods_loaded_from.append(dir_for_mod_files)
            except:
                print_v("Failed to load mod file mechanisms...")
        else:
            print_v(
                "Already loaded mechanisms from %s (all loaded: %s)"
                % (dir_for_mod_files, locations_mods_loaded_from)
            )


def generate_and_run(
    simulation,
    simulator,
    network=None,
    return_results=False,
    base_dir=None,
    target_dir=None,
    num_processors=1,
    nogui=False,
):
    """
    Generates the network in the specified simulator and runs, if appropriate
    """

    if network == None:
        network = load_network(simulation.network)

    print_v(
        "Generating network %s and running in simulator: %s..."
        % (network.id, simulator)
    )

    try:
        if simulator == "NEURON":  # NOT YET WORKING...

            _generate_neuron_files_from_neuroml(network, dir_for_mod_files=target_dir)

            from neuromllite.NeuronHandler import NeuronHandler

            nrn_handler = NeuronHandler()

            for c in network.cells:
                if c.neuroml2_source_file:
                    src_dir = os.path.dirname(os.path.abspath(c.neuroml2_source_file))
                    nrn_handler.executeHoc('load_file("%s/%s.hoc")' % (src_dir, c.id))

            generate_network(network, nrn_handler, generate_network, base_dir)
            if return_results:
                raise NotImplementedError(
                    "Reloading results not supported in Neuron yet..."
                )

        elif simulator.lower() == "mdf":

            from neuromllite.MDFHandler import MDFHandler

            mdf_handler = MDFHandler(nl_network=network)

            generate_network(network, mdf_handler)

        elif simulator.lower() == "psyneulink":

            from neuromllite.PsyNeuLinkHandler import PsyNeuLinkHandler

            pnl_handler = PsyNeuLinkHandler(nl_network=network)

            generate_network(network, pnl_handler)
            from neuromllite import __version__ as nmlliteversion

            run_pnl_script = """# Generated by NeuroMLlite v{1}

import psyneulink as pnl

conv_json_filename = '{0}.bids-mdf.json'
with open(conv_json_filename, 'r') as infile:
    json = infile.read()

load_model = 'load_{0}_bids_mdf'
py_filename = load_model+'.py'
with open(py_filename, 'w') as outfi:
    outfi.write(pnl.generate_script_from_json(json))

exec('from %s import {0}'%load_model)

print('Ready to run network: {0} imported using %s'%py_filename)

dt = {2}
simtime = {3}

time_step_size=dt
num_trials=int(simtime/dt)

{0}.run(inputs={{}}, log=True, num_trials=num_trials)
print('Finished simulation run of %s ms'%simtime)

def generate_time_array(node, context='{0}', param='value'):
    return [entry.time.trial for entry in getattr(node.parameters, param).log[context]]

def generate_value_array(node, index, context='{0}', param='value'):
    return [float(entry.value[index]) for entry in getattr(node.parameters, param).log[context]]

import matplotlib.pyplot as plt
fig, axes = plt.subplots()

all_data_files = []

################################################################################
#          Note: this currently just plots/saved (almost) everything...
################################################################################

for node in {0}.nodes:
    var_num = 2 if node.function.__class__.__name__=='FitzHughNagumoIntegrator' else 1
    for i in range(var_num):
        x_values = generate_time_array(node)
        y_values = generate_value_array(node, i)

        dat_file = '{0}_%s_%i.dat'%(node.name,i)
        all_data_files.append(dat_file)
        fout = open(dat_file,'w')

        for index in range(len(x_values)):
            #                                                       1000 to convert ms to s
            fout.write('%s\\t%s\\n'%(x_values[index]*time_step_size/1000.0, y_values[index]))
        fout.close()

        times = [t*time_step_size/1000.0 for t in x_values]

        axes.plot(
            times,
            y_values,
            label=f'Value of {{i}} {{node.name}}, {{node.function.__class__.__name__}}'
        )

axes.set_xlabel('Time (s)')
axes.legend()

print('Data from simulation saved to: %s'%all_data_files)
plt.show()

#{0}.show_graph()
            """.format(
                network.id, nmlliteversion, simulation.dt, simulation.duration
            )

            run_pnl_file = open("run_pnl_%s.py" % network.id, "w")

            run_pnl_file.write(run_pnl_script)

            print_v("Done with PsyNeuLink...")

        elif simulator.lower() == "sonata":

            target_simulator = "NEST"

            from neuromllite.SonataHandler import SonataHandler

            sonata_handler = SonataHandler()

            generate_network(
                network, sonata_handler, always_include_props=True, base_dir=base_dir
            )

            import pyNN.neuroml

            for c in network.cells:
                if c.pynn_cell:

                    cell_params = {}
                    if c.parameters:
                        for p in c.parameters:
                            cell_params[p] = evaluate(
                                c.parameters[p], network.parameters
                            )

                    temp_cell = eval("pyNN.neuroml.%s(**cell_params)" % c.pynn_cell)

                    mappings = {
                        "i_offset": "I_e",
                        "tau_m": "tau_m",
                        "cm": "C_m",
                        "tau_refrac": "t_ref",
                        "v_rest": "E_L",
                        "v_thresh": "V_th",
                        "v_reset": "V_reset",
                    }
                    scales = {"cm": 1000}

                    print_v("Converting %s (%s) to Sonata format..." % (c, temp_cell))
                    comp_file = (
                        "./components/point_neuron_models_dir/%s.json" % c.id
                    )  # Oi! Hardcoding!
                    comp_file_info = {}

                    for p in temp_cell.default_parameters:
                        print(
                            "  Param %s = %s"
                            % (p, temp_cell.parameter_space[p].base_value)
                        )
                        if p in mappings:
                            sim_name = mappings[p]
                            sim_val = temp_cell.parameter_space[p].base_value * (
                                scales[p] if p in scales else 1
                            )
                            # print('    Translated %s = %s'%(sim_name, sim_val))
                            comp_file_info[sim_name] = sim_val

                    save_to_json_file(comp_file_info, comp_file, indent=2)

            sim_file_info = {}

            sim_file_info["run"] = {}
            sim_file_info["run"]["tstop"] = simulation.duration
            sim_file_info["run"]["dt"] = simulation.dt

            sim_file_info["target_simulator"] = target_simulator
            sim_file_info["manifest"] = {}
            sim_file_info["manifest"]["$OUTPUT_DIR"] = "./output"
            sim_file_info["manifest"]["$INPUT_DIR"] = "./"
            sim_file_info["output"] = {}
            sim_file_info["output"]["output_dir"] = "$OUTPUT_DIR"
            sim_file_info["output"]["log_file"] = "log.txt"
            sim_file_info["output"]["spikes_file"] = "spikes.h5"
            sim_file_info["output"]["spikes_sort_order"] = "time"

            sim_file_info["node_sets"] = {}
            sim_file_info["node_sets"]["point_nodes"] = {}
            sim_file_info["node_sets"]["point_nodes"]["model_type"] = "point_process"

            sim_file_info["reports"] = {}
            sim_file_info["reports"]["membrane_potential"] = {}
            sim_file_info["reports"]["membrane_potential"]["cells"] = "point_nodes"
            sim_file_info["reports"]["membrane_potential"]["variable_name"] = "V_m"
            sim_file_info["reports"]["membrane_potential"][
                "module"
            ] = "multimeter_report"
            sim_file_info["reports"]["membrane_potential"]["sections"] = "soma"
            sim_file_info["reports"]["membrane_potential"]["enabled"] = True

            # print sonata_handler.input_info

            temp_nml_doc = _extract_pynn_components_to_neuroml(network)

            summary = temp_nml_doc.summary()
            from pyneuroml.pynml import convert_to_units

            sim_file_info["inputs"] = {}
            for input in sonata_handler.input_info:
                sim_file_info["inputs"][input] = {}
                input_comp = sonata_handler.input_info[input][1]
                c = temp_nml_doc.get_by_id(input_comp)
                ref = "inputset_%s_%s" % (input, input_comp)
                if c.__class__.__name__ == "PulseGenerator":
                    sim_file_info["inputs"][input]["input_type"] = "current_clamp"
                    sim_file_info["inputs"][input]["module"] = "IClamp"
                    sim_file_info["inputs"][input]["amp"] = convert_to_units(
                        c.amplitude, "pA"
                    )
                    sim_file_info["inputs"][input]["delay"] = convert_to_units(
                        c.delay, "ms"
                    )
                    sim_file_info["inputs"][input]["duration"] = convert_to_units(
                        c.duration, "ms"
                    )
                    sim_file_info["inputs"][input]["node_set"] = ref
                else:
                    raise Exception("Can't yet handle input of type: %s" % c)

                sim_file_info["node_sets"][ref] = {}
                sim_file_info["node_sets"][ref]["model_type"] = "point_process"
                sim_file_info["node_sets"][ref][
                    "population"
                ] = sonata_handler.input_info[input][0]
                # sim_file_info["node_sets"][ref]['node_id'] = sonata_handler.input_info[input][2]

            save_to_json_file(sim_file_info, "simulation_config.json", indent=2)

            run_bmtk_template = """#!/bin/env python

    import sys

    def run(config_file, simulator):

        if simulator=='NEURON':
            from bmtk.simulator import bionet
            conf = bionet.Config.from_json(config_file, validate=True)
            conf.build_env()
            net = bionet.BioNetwork.from_config(conf)
            sim = bionet.BioSimulator.from_config(conf, network=net)


        elif simulator=='NEST':
            from bmtk.simulator import pointnet
            conf = pointnet.Config.from_json(config_file)
            conf.build_env()
            net = pointnet.PointNetwork.from_config(conf)
            sim = pointnet.PointSimulator.from_config(conf, net)

        sim.run()


    if __name__ == '__main__':

            run('config.json', '%s')

            """

            run_bmtk_file = open("run_bmtk.py", "w")
            run_bmtk_file.write(run_bmtk_template % (target_simulator))

            print_v("Done with Sonata...")

        elif simulator.lower().startswith("graph"):  # Will not "run" obviously...

            from neuromllite.GraphVizHandler import GraphVizHandler, engines

            try:
                if simulator[-1].isalpha():

                    engine = engines[simulator[-1]]
                    level = int(simulator[5:-1])
                else:
                    engine = "dot"
                    level = int(simulator[5:])

            except Exception as e:
                print_v("Error parsing: %s: %s" % (simulator, e))
                print_v(
                    "Graphs of the network structure can be generated at many levels of detail (1-6, required) and laid out using GraphViz engines (d - dot (default); c - circo; n - neato; f - fdp), so use: -graph3c, -graph2, -graph4f etc."
                )
                return

            handler = GraphVizHandler(level, engine=engine, nl_network=network, view_on_render = not nogui)

            generate_network(
                network, handler, always_include_props=True, base_dir=base_dir
            )

            print_v("Done with GraphViz...")

        elif simulator.lower().startswith("matrix"):  # Will not "run" obviously...

            from neuromllite.MatrixHandler import MatrixHandler

            try:
                level = int(simulator[6:])
            except:
                print_v("Error parsing: %s" % simulator)
                print_v(
                    "Matrices of the network structure can be generated at many levels of detail (1-n, required), so use: -matrix1, -matrix2, etc."
                )
                return

            handler = MatrixHandler(level, nl_network=network)

            generate_network(
                network, handler, always_include_props=True, base_dir=base_dir
            )

            print_v("Done with MatrixHandler...")

        elif simulator == "BindsNET":

            from neuromllite.BindsNETHandler import BindsNETHandler
            import bindsnet

            print("\n   ********************************************************")
            print("   *** Warning: Support for BindsNET is very preliminary!! ***")
            print("   ********************************************************\n")

            bindsnet_handler = BindsNETHandler(network)

            generate_network(
                network, bindsnet_handler, always_include_props=True, base_dir=base_dir
            )

            monitors_v = {}
            monitors_s = {}
            for pop in bindsnet_handler.pops_vs_components:
                comp = bindsnet_handler.pops_vs_components[pop]
                bn_layer = bindsnet_handler.pops_vs_bn_layers[pop]
                print("Monitoring all spikes in %s (%s) -> %s" % (pop, comp, bn_layer))
                m = bindsnet.network.monitors.Monitor(obj=bn_layer, state_vars=["s"])
                bindsnet_handler.bn_network.add_monitor(
                    monitor=m, name="%s_spikes" % pop
                )
                monitors_s[pop] = m

                if comp != "Input":
                    print("Monitoring all v in %s (%s)" % (pop, comp))
                    mv = bindsnet.network.monitors.Monitor(bn_layer, ["v"])
                    bindsnet_handler.bn_network.add_monitor(
                        monitor=mv, name="%s_voltage" % pop
                    )
                    monitors_v[pop] = mv

            traces = {}
            events = {}

            times = [
                t / 1000.0 for t in range(int(simulation.duration / simulation.dt) + 1)
            ]
            traces["t"] = times

            bindsnet_handler.bn_network.run(inputs={}, time=simulation.duration)

            for pop in monitors_v:
                vs = monitors_v[pop].get("v").numpy()
                print("Shape of voltage results for %s: %s" % (pop, vs.shape))
                v = vs
                print(v)
                traces["%s/0" % pop] = v

            for pop in monitors_s:
                ss = monitors_s[pop].get("s").numpy()
                print("Shape of spike results for %s: %s" % (pop, ss.shape))
                s = ss
                print(s)
                # traces['%s/0'%pop]= ss

            print_v("Finished BindsNET simulation")
            _print_result_info(traces, events)

            if return_results:
                return traces, events

        elif simulator == "Arbor":

            from neuromllite.ArborHandler import ArborHandler
            import arbor

            print("\n   *********************************************************")
            print("   *** Warning: Support for Arbor is very preliminary!!! ***")
            print("   *********************************************************\n")

            arbor_handler = ArborHandler(network)

            generate_network(
                network, arbor_handler, always_include_props=True, base_dir=base_dir
            )

            arbor_recipe = arbor_handler.neuroML_arbor_recipe

            trace_pop_indices_seg_ids = get_pops_vs_cell_indices_seg_ids(
                simulation.record_traces, network
            )
            spike_pop_indices = get_pops_vs_cell_indices_seg_ids(
                simulation.record_spikes, network
            )
            """
            for pop_id in arbor_handler.populations:
                arbor_pop = arbor_handler.populations[pop_id]
                if arbor_pop.label in trace_pop_indices_seg_ids:
                    if arbor_pop.can_record('v'):
                        indices = trace_pop_indices_seg_ids[pop_id].keys()
                        print_v("Recording v in %s in cells %s" % (indices, pop_id))
                        for index in indices:
                            #arbor_handler.sim.record('v',arbor_pop[index],'PP_%s_%s_%s.pkl'%(pop_id,index,'v'))
                if arbor_pop.label in spike_pop_indices:
                    if arbor_pop.can_record('spikes'):
                        indices = spike_pop_indices[pop_id].keys()
                        print_v("Recording spikes in %s in cells %s" % (indices, pop_id))
                        for index in indices:
                            #arbor_handler.sim.record('spikes',arbor_pop[index],'PP_%s_%s.pkl'%(pop_id,'spike'))
            """

            ncells = arbor_recipe.num_cells()
            print_v(
                "Starting Arbor simulation with %i cells of duration %sms (dt: %sms)"
                % (ncells, simulation.duration, simulation.dt)
            )

            # arbor_model.run(tfinal=simulation.duration)

            # (12) Create a default execution context, domain decomposition and simulation
            context = arbor.context()
            decomp = arbor.partition_load_balance(arbor_recipe, context)
            sim = arbor.simulation(arbor_recipe, decomp, context)

            # (13) Set spike generators to record
            sim.record(arbor.spike_recording.all)

            # (14) Attach a sampler to the voltage probe on cell 0.
            handles = [
                sim.sample((gid, 0), arbor.regular_schedule(simulation.dt))
                for gid in range(ncells)
            ]
            print_v("handles: %s, %s" % (handles, range(ncells)))
            # (15) Run simulation
            sim.run(simulation.duration)

            print_v("Finished Arbor simulation")
            # arbor_handler.sim.end()

            traces = {}
            events = {}

            for pop_id in trace_pop_indices_seg_ids:
                indices = trace_pop_indices_seg_ids[pop_id].keys()
                for index in indices:

                    filename = "%s.%s.%s.v.dat" % (simulation.id, pop_id, index)
                    # all_columns = []
                    gid = arbor_recipe.get_gid(pop_id, index)

                    samples, meta = sim.samples(handles[gid])[0]

                    data = samples[:, 1]
                    times = samples[:, 0]

                    print_v(
                        "Writing data for %s[%s] (gid: %s) to %s: (%s) %s"
                        % (pop_id, index, gid, filename, meta, data)
                    )

                    f = open(filename, "w")
                    for ti in range(len(times)):
                        f.write("%s\t%s\n" % (times[ti], data[ti]))

                    ref = "%s[%i]" % (pop_id, index)

                    vm_si = [d / 1000.0 for d in data]
                    traces[ref] = vm_si
                    t_si = [t / 1000.0 for t in times]
                    traces["t"] = t_si

            for pop_id in spike_pop_indices:
                indices = spike_pop_indices[pop_id]

                filename = "%s.%s.spikes" % (simulation.id, pop_id)
                all_columns = []
                print_v(
                    "Writing spike data for %s %s to %s" % (pop_id, indices, filename)
                )
                """
                arbor_pop = arbor_handler.populations[pop_id]
                data =  arbor_pop.get_data('spikes', gather=False)
                spiketrains = data.segments[0].spiketrains

                ff = open(filename, 'w')

                for spiketrain in spiketrains:
                    source_id = spiketrain.annotations['source_id']
                    source_index = spiketrain.annotations['source_index']
                    if source_index in indices:
                        #print_v("Writing spike data for cell %s[%s] (gid: %i): %i spikes "%(arbor_pop.label,source_index, source_id, len(spiketrain)), self.verbose)
                        ref = '%s/%i/???'%(arbor_pop.label,source_index)
                        events[ref] = [t.magnitude/1000. for t in spiketrain]
                        for t in spiketrain:
                            #ff.write('%s\t%i\n'%(t.magnitude/1000.,source_index))
                            ff.write('%i\t%s\n'%(source_index,t.magnitude/1000.))
                ff.close()"""

            if return_results:
                _print_result_info(traces, events)
                return traces, events

        elif simulator.startswith("PyNN"):

            # _generate_neuron_files_from_neuroml(network)
            simulator_name = simulator.split("_")[1].lower()

            from neuromllite.PyNNHandler import PyNNHandler

            pynn_handler = PyNNHandler(simulator_name, simulation.dt, network.id)

            syn_cell_params = {}
            for proj in network.projections:

                synapse = network.get_child(proj.synapse, "synapses")
                post_pop = network.get_child(proj.postsynaptic, "populations")

                if not post_pop.component in syn_cell_params:
                    syn_cell_params[post_pop.component] = {}
                for p in synapse.parameters:
                    post = ""
                    if synapse.pynn_receptor_type == "excitatory":
                        post = "_E"
                    elif synapse.pynn_receptor_type == "inhibitory":
                        post = "_I"
                    syn_cell_params[post_pop.component]["%s%s" % (p, post)] = evaluate(
                        synapse.parameters[p], network.parameters
                    )

            cells = {}
            for c in network.cells:
                if c.pynn_cell:
                    cell_params = {}
                    if c.parameters:
                        for p in c.parameters:
                            cell_params[p] = evaluate(
                                c.parameters[p], network.parameters
                            )

                    dont_set_here = ["tau_syn_E", "e_rev_E", "tau_syn_I", "e_rev_I"]
                    for d in dont_set_here:
                        if d in c.parameters:
                            print("Problem with %s, %s" % (c.id, c.parameters))
                            raise Exception(
                                "Synaptic parameters like %s should be set " % d
                                + "in individual synapses, not in the list of parameters associated with the cell"
                            )
                    if c.id in syn_cell_params:
                        cell_params.update(syn_cell_params[c.id])
                    print_v("Creating cell with params: %s" % cell_params)
                    exec(
                        'cells["%s"] = pynn_handler.sim.%s(**cell_params)'
                        % (c.id, c.pynn_cell)
                    )

                    if c.pynn_cell != "SpikeSourcePoisson":
                        exec(
                            "cells['%s'].default_initial_values['v'] = cells['%s'].parameter_space['v_rest'].base_value"
                            % (c.id, c.id)
                        )

            pynn_handler.set_cells(cells)

            receptor_types = {}
            for s in network.synapses:
                if s.pynn_receptor_type:
                    receptor_types[s.id] = s.pynn_receptor_type

            pynn_handler.set_receptor_types(receptor_types)

            for input_source in network.input_sources:

                # TODO: change this when noisyCurrentSource is integrated into the core of NeuroML
                if input_source.pynn_input or input_source.lems_source_file:
                    pynn_handler.add_input_source(input_source, network)

            generate_network(
                network, pynn_handler, always_include_props=True, base_dir=base_dir
            )

            trace_pop_indices_seg_ids = get_pops_vs_cell_indices_seg_ids(
                simulation.record_traces, network
            )
            spike_pop_indices = get_pops_vs_cell_indices_seg_ids(
                simulation.record_spikes, network
            )

            for pop_id in pynn_handler.populations:
                pynn_pop = pynn_handler.populations[pop_id]
                if pynn_pop.label in trace_pop_indices_seg_ids:
                    if pynn_pop.can_record("v"):
                        indices = trace_pop_indices_seg_ids[pop_id].keys()
                        print_v("Recording v in %s in cells %s" % (indices, pop_id))
                        for index in indices:
                            pynn_handler.sim.record(
                                "v",
                                pynn_pop[index],
                                "PP_%s_%s_%s.pkl" % (pop_id, index, "v"),
                            )
                if pynn_pop.label in spike_pop_indices:
                    if pynn_pop.can_record("spikes"):
                        indices = spike_pop_indices[pop_id].keys()
                        print_v(
                            "Recording spikes in %s in cells %s" % (indices, pop_id)
                        )
                        for index in indices:
                            pynn_handler.sim.record(
                                "spikes",
                                pynn_pop[index],
                                "PP_%s_%s.pkl" % (pop_id, "spike"),
                            )

            print_v(
                "Starting PyNN simulation of duration %sms (dt: %sms)"
                % (simulation.duration, simulation.dt)
            )
            pynn_handler.sim.run(simulation.duration)
            print_v("Finished PyNN simulation")
            pynn_handler.sim.end()

            traces = {}
            events = {}

            if not "NeuroML" in simulator:

                for pop_id in trace_pop_indices_seg_ids:
                    pynn_pop = pynn_handler.populations[pop_id]
                    indices = trace_pop_indices_seg_ids[pop_id].keys()

                    filename = "%s.%s.v.dat" % (simulation.id, pop_id)
                    all_columns = []

                    print_v(
                        "Writing data for %s %s to %s" % (pop_id, indices, filename)
                    )

                    data = pynn_pop.get_data("v", gather=False)
                    analogsignal = data.segments[0].analogsignals[0]

                    source_ids = analogsignal.annotations["source_ids"]
                    analogsignal_index = {}
                    for sii, glob_id in enumerate(source_ids):
                        index_in_pop = pynn_pop.id_to_index(glob_id)
                        # print('Index %i in analogsignal is pynn global id %i, pop index %i'%(sii,glob_id,index_in_pop))
                        analogsignal_index[index_in_pop] = sii

                    for index in indices:
                        ref = "%s[%i]" % (pop_id, index)
                        ref = "%s/%i/???/v" % (pop_id, index)
                        traces[ref] = []
                        vm = analogsignal.transpose()[analogsignal_index[index]]

                        if len(all_columns) == 0:
                            tt = np.array(
                                [t * simulation.dt / 1000.0 for t in range(len(vm))]
                            )
                            all_columns.append(tt)
                        vm_si = np.array(vm / 1000.0)

                        traces[ref] = vm_si
                        if not "t" in traces:
                            traces["t"] = tt

                        all_columns.append(vm_si)

                        times_vm = np.array(all_columns).transpose()

                    np.savetxt(filename, times_vm, delimiter="\t", fmt="%s")

                for pop_id in spike_pop_indices:
                    pynn_pop = pynn_handler.populations[pop_id]
                    indices = spike_pop_indices[pop_id]

                    filename = "%s.%s.spikes" % (simulation.id, pop_id)
                    all_columns = []
                    print_v(
                        "Writing spike data for %s %s to %s"
                        % (pop_id, indices, filename)
                    )

                    data = pynn_pop.get_data("spikes", gather=False)
                    spiketrains = data.segments[0].spiketrains

                    ff = open(filename, "w")

                    for spiketrain in spiketrains:
                        source_id = spiketrain.annotations["source_id"]
                        source_index = spiketrain.annotations["source_index"]
                        if source_index in indices:
                            # print_v("Writing spike data for cell %s[%s] (gid: %i): %i spikes "%(pynn_pop.label,source_index, source_id, len(spiketrain)), self.verbose)
                            ref = "%s/%i/???" % (pynn_pop.label, source_index)
                            events[ref] = [t.magnitude / 1000.0 for t in spiketrain]
                            for t in spiketrain:
                                # ff.write('%s\t%i\n'%(t.magnitude/1000.,source_index))
                                ff.write(
                                    "%i\t%s\n" % (source_index, t.magnitude / 1000.0)
                                )
                    ff.close()

            if return_results:
                _print_result_info(traces, events)
                return traces, events

        elif simulator == "NetPyNE":

            if target_dir == None:
                target_dir = "./"

            _generate_neuron_files_from_neuroml(network, dir_for_mod_files=target_dir)

            from netpyne import specs
            from netpyne import sim

            # Note NetPyNE from this branch is required: https://github.com/Neurosim-lab/netpyne/tree/neuroml_updates
            from netpyne.conversion.neuromlFormat import NetPyNEBuilder

            import pprint

            pp = pprint.PrettyPrinter(depth=6)

            netParams = specs.NetParams()
            simConfig = specs.SimConfig()
            netpyne_handler = NetPyNEBuilder(
                netParams, simConfig=simConfig, verbose=True
            )

            generate_network(network, netpyne_handler, base_dir=base_dir)

            netpyne_handler.finalise()

            simConfig = specs.SimConfig()
            simConfig.tstop = simulation.duration
            simConfig.duration = simulation.duration
            simConfig.dt = simulation.dt
            simConfig.seed = simulation.seed
            simConfig.recordStep = simulation.dt

            simConfig.recordCells = ["all"]
            simConfig.recordTraces = {}

            trace_pop_indices_seg_ids = get_pops_vs_cell_indices_seg_ids(
                simulation.record_traces, network
            )

            for pop_id in trace_pop_indices_seg_ids:
                for index in trace_pop_indices_seg_ids[pop_id]:
                    seg_ids = trace_pop_indices_seg_ids[pop_id][index]
                    if seg_ids == None:
                        simConfig.recordTraces[
                            "%s.%s.%s.v" % (simulation.id, pop_id, index)
                        ] = {
                            "sec": "soma",
                            "loc": 0.5,
                            "var": "v",
                            "conds": {"pop": pop_id, "cellLabel": index},
                        }
                    else:
                        for seg_id in seg_ids:
                            simConfig.recordTraces[
                                "%s.%s.%s.%s.v" % (simulation.id, pop_id, index, seg_id)
                            ] = {
                                "sec": seg_id,
                                "loc": 0.5,
                                "var": "v",
                                "conds": {"pop": pop_id, "cellLabel": index},
                            }

            # simConfig.saveDat = True

            print_v("NetPyNE netParams: ")
            pp.pprint(netParams.todict())
            # print_v("NetPyNE simConfig: ")
            # pp.pprint(simConfig.todict())

            sim.initialize(
                netParams, simConfig
            )  # create network object and set cfg and net params

            sim.net.createPops()
            cells = (
                sim.net.createCells()
            )  # instantiate network cells based on defined populations

            for proj_id in netpyne_handler.projection_infos.keys():
                (
                    projName,
                    prePop,
                    postPop,
                    synapse,
                    ptype,
                ) = netpyne_handler.projection_infos[proj_id]
                print_v(
                    "Creating connections for %s (%s): %s->%s via %s"
                    % (projName, ptype, prePop, postPop, synapse)
                )

                preComp = netpyne_handler.pop_ids_vs_components[prePop]

                for conn in netpyne_handler.connections[projName]:

                    (
                        pre_id,
                        pre_seg,
                        pre_fract,
                        post_id,
                        post_seg,
                        post_fract,
                        delay,
                        weight,
                    ) = conn

                    # connParam = {'delay':delay,'weight':weight,'synsPerConn':1, 'sec':post_seg, 'loc':post_fract, 'threshold':threshold}
                    connParam = {
                        "delay": delay,
                        "weight": weight,
                        "synsPerConn": 1,
                        "sec": post_seg,
                        "loc": post_fract,
                    }

                    if ptype == "electricalProjection":

                        if weight != 1:
                            raise Exception(
                                "Cannot yet support inputs where weight !=1!"
                            )
                        connParam = {
                            "synsPerConn": 1,
                            "sec": post_seg,
                            "loc": post_fract,
                            "gapJunction": True,
                            "weight": weight,
                        }
                    else:
                        connParam = {
                            "delay": delay,
                            "weight": weight,
                            "synsPerConn": 1,
                            "sec": post_seg,
                            "loc": post_fract,
                        }
                        #'threshold': threshold}

                    connParam["synMech"] = synapse

                    if (
                        post_id in sim.net.gid2lid
                    ):  # check if postsyn is in this node's list of gids
                        sim.net._addCellConn(connParam, pre_id, post_id)

            stims = (
                sim.net.addStims()
            )  # add external stimulation to cells (IClamps etc)
            simData = (
                sim.setupRecording()
            )  # setup variables to record for each cell (spikes, V traces, etc)
            sim.runSim()  # run parallel Neuron simulation
            sim.gatherData()  # gather spiking data and cell info from each node
            sim.saveData()  # save params, cell info and sim output to file (pickle,mat,txt,etc)

            print_v("Finished running simulation in NetPyNE")

            print_v("Extracting results from %s" % (sim.allSimData.keys()))

            times = [t / 1000.0 for t in sim.allSimData["t"]]
            traces = {}
            traces["t"] = times
            events = {}

            for pop_id in trace_pop_indices_seg_ids:
                indices_segids = trace_pop_indices_seg_ids[pop_id]
                print("Saving in population %s: %s" % (pop_id, indices_segids))
                save_ref = "%s.%s.v.dat" % (simulation.id, pop_id)
                pop_v_file = open(save_ref, "w")
                all_v = []
                for index in indices_segids:
                    seg_ids = indices_segids[index]
                    if seg_ids == None:
                        ref = "%s.%s.%s.v" % (simulation.id, pop_id, index)
                        # print('Ref: %s, existing: %s'%(ref, sim.allSimData.keys()))
                        v = [
                            v / 1000.0
                            for v in sim.allSimData[ref][list(sim.allSimData[ref])[0]]
                        ]
                        # print _print_info(v)

                        ref = "%s/%i/???/v" % (pop_id, index)
                        traces[ref] = v
                        all_v.append(v)
                    else:
                        for seg_id in seg_ids:
                            ref = "%s.%s.%s.%s.v" % (
                                simulation.id,
                                pop_id,
                                index,
                                seg_id,
                            )
                            print(
                                "Ref: %s, existing: %s" % (ref, sim.allSimData.keys())
                            )
                            v = [
                                v / 1000.0
                                for v in sim.allSimData[ref][
                                    list(sim.allSimData[ref])[0]
                                ]
                            ]
                            # print _print_info(v)

                            ref = "%s/%i/%i/v" % (pop_id, index, seg_id)
                            traces[ref] = v
                            all_v.append(v)

                for ti in range(len(times)):
                    vs = ""
                    for v in all_v:
                        vs += "\t%s" % (v[ti])
                    pop_v_file.write("%s\t" % (times[ti]) + vs + "\n")
                pop_v_file.close()

            spike_pop_indices = sorted(
                get_pops_vs_cell_indices_seg_ids(
                    simulation.record_spikes, network
                ).keys()
            )

            if (
                return_results or True
            ):  ###########################################################
                return traces, events

        elif (
            simulator == "jNeuroML"
            or simulator == "jNeuroML_norun"
            or simulator == "jNeuroML_NEURON"
            or simulator == "jNeuroML_NetPyNE"
            or simulator == "jNeuroML_Brian2"
        ):

            from pyneuroml.lems import generate_lems_file_for_neuroml
            from pyneuroml import pynml

            lems_file_name = "LEMS_%s.xml" % simulation.id

            nml_file_name, nml_doc = generate_neuroml2_from_network(
                network, simulation=simulation, base_dir=base_dir, target_dir=target_dir
            )
            included_files = ["PyNN.xml"]
            """ Needed?
            for c in network.cells:
                if c.lems_source_file:
                    included_files.append(c.lems_source_file)

            if network.cells:
                for c in network.cells:
                    included_files.append(c.neuroml2_source_file)

            if network.synapses:
                for s in network.synapses:
                    if s.lems_source_file:
                        fname = locate_file(s.lems_source_file, base_dir)
                        # more ..?
            """
            print_v("Generating LEMS file prior to running in %s" % simulator)

            pops_plot_save = []
            pops_spike_save = []
            gen_plots_for_quantities = {}
            gen_saves_for_quantities = {}
            gen_spike_saves_for_cells = {}

            trace_pop_indices_seg_ids = get_pops_vs_cell_indices_seg_ids(
                simulation.record_traces, network
            )
            spike_pop_indices = get_pops_vs_cell_indices_seg_ids(
                simulation.record_spikes, network
            )
            rate_pop_indices = get_pops_vs_cell_indices_seg_ids(
                simulation.record_rates, network
            )

            pops = network.populations
            pops.sort(key=id)
            for p in pops:

                if p.id in spike_pop_indices:
                    pops_spike_save.append(p.id)

                    save_ref = "%s.%s.spikes" % (simulation.id, p.id)
                    gen_spike_saves_for_cells[save_ref] = []

                    for i in spike_pop_indices[p.id].keys():
                        quantity = "%s/%i/%s" % (p.id, i, p.component)
                        if not p.has_positions():
                            quantity = "%s[%i]" % (p.id, i)
                        gen_spike_saves_for_cells[save_ref].append(quantity)

                if p.id in trace_pop_indices_seg_ids:
                    plot_ref = "%s_v" % (p.id)
                    save_ref = "%s.%s.v.dat" % (simulation.id, p.id)
                    gen_plots_for_quantities[plot_ref] = []
                    gen_saves_for_quantities[save_ref] = []

                    for cell_index in trace_pop_indices_seg_ids[p.id]:
                        seg_ids = trace_pop_indices_seg_ids[p.id][cell_index]
                        print_v(
                            "Recording in pop %s, cell index %i, seg ids: %s"
                            % (p.id, cell_index, seg_ids)
                        )
                        if seg_ids == None:
                            quantity = "%s/%i/%s/v" % (p.id, cell_index, p.component)
                            if not p.has_positions():
                                quantity = "%s[%i]/v" % (p.id, cell_index)
                            gen_plots_for_quantities[plot_ref].append(quantity)
                            gen_saves_for_quantities[save_ref].append(quantity)
                        else:
                            for seg_id in seg_ids:
                                quantity = "%s/%i/%s/%s/v" % (
                                    p.id,
                                    cell_index,
                                    p.component,
                                    seg_id,
                                )
                                if not p.has_positions():
                                    quantity = "%s[%i]/%s/v" % (
                                        p.id,
                                        cell_index,
                                        seg_id,
                                    )
                                gen_plots_for_quantities[plot_ref].append(quantity)
                                gen_saves_for_quantities[save_ref].append(quantity)

                if p.id in rate_pop_indices:
                    plot_ref = "%s_r" % (p.id)
                    save_ref = "%s.%s.r.dat" % (simulation.id, p.id)
                    gen_plots_for_quantities[plot_ref] = []
                    gen_saves_for_quantities[save_ref] = []

                    for i in rate_pop_indices[p.id].keys():
                        quantity = "%s/%i/%s/r" % (p.id, i, p.component)
                        gen_plots_for_quantities[plot_ref].append(quantity)
                        gen_saves_for_quantities[save_ref].append(quantity)

                if simulation.record_variables:
                    for var in simulation.record_variables:
                        to_rec = simulation.record_variables[var]
                        if "all" in to_rec or p.id in to_rec:
                            size = evaluate(p.size, network.parameters)
                            for i in range(size):
                                quantity = "%s/%i/%s/%s" % (p.id, i, p.component, var)
                                if not p.has_positions():
                                    quantity = "%s[%i]/%s" % (p.id, i, var)
                                gen_plots_for_quantities[
                                    "%s_%i_%s" % (p.id, i, var)
                                ] = [quantity]
                                safe_var = var.replace("/", "_")
                                gen_saves_for_quantities[
                                    "%s_%i.%s.dat" % (p.id, i, safe_var)
                                ] = [quantity]

            generate_lems_file_for_neuroml(
                simulation.id,
                nml_file_name,
                network.id,
                simulation.duration,
                simulation.dt,
                lems_file_name,
                target_dir=target_dir if target_dir else ".",
                nml_doc=nml_doc,  # Use this if the nml doc has already been loaded (to avoid delay in reload)
                include_extra_files=included_files,
                gen_plots_for_all_v=False,
                plot_all_segments=False,
                gen_plots_for_quantities=gen_plots_for_quantities,  # Dict with displays vs lists of quantity paths
                gen_plots_for_only_populations=pops_plot_save,  # List of populations, all pops if = []
                gen_saves_for_all_v=False,
                save_all_segments=False,
                gen_saves_for_only_populations=pops_plot_save,  # List of populations, all pops if = []
                gen_saves_for_quantities=gen_saves_for_quantities,  # Dict with file names vs lists of quantity paths
                gen_spike_saves_for_all_somas=False,
                gen_spike_saves_for_only_populations=[],  # List of populations, all pops if = []
                gen_spike_saves_for_cells=gen_spike_saves_for_cells,  # Dict with file names vs lists of quantity paths
                spike_time_format="ID_TIME",
                copy_neuroml=True,
                lems_file_generate_seed=12345,
                report_file_name="report.%s.txt" % simulation.id,
                simulation_seed=simulation.seed
                if simulation.seed
                else DEFAULT_SIMULATION_SEED,
                verbose=True,
            )

            lems_file_name = locate_file(lems_file_name, target_dir)

            if simulator == "jNeuroML_norun":
                return lems_file_name

            if simulator == "jNeuroML":
                results = pynml.run_lems_with_jneuroml(
                    lems_file_name,
                    nogui=True,
                    load_saved_data=return_results,
                    reload_events=return_results,
                )
            elif simulator == "jNeuroML_NEURON":
                results = pynml.run_lems_with_jneuroml_neuron(
                    lems_file_name,
                    nogui=True,
                    load_saved_data=return_results,
                    reload_events=return_results,
                )
            elif simulator == "jNeuroML_NetPyNE":
                results = pynml.run_lems_with_jneuroml_netpyne(
                    lems_file_name,
                    nogui=True,
                    verbose=True,
                    load_saved_data=return_results,
                    reload_events=return_results,
                    num_processors=num_processors,
                )

            elif simulator == "jNeuroML_Brian2":
                """Work in progress..."""
                results = pynml.run_lems_with_jneuroml_brian2(
                    lems_file_name,
                    nogui=True,
                    load_saved_data=return_results,
                    reload_events=return_results,
                )

            from pyneuroml import __version__, JNEUROML_VERSION

            print_v(
                "Finished running LEMS file %s in %s (pyNeuroML v%s containing jNeuroML v%s; returning results: %s)"
                % (
                    lems_file_name,
                    simulator,
                    __version__,
                    JNEUROML_VERSION,
                    return_results,
                )
            )

            if return_results:
                traces, events = results
                _print_result_info(traces, events)
                return results  # traces, events =
    except Exception as e:
        raise e
        raise Exception(
            "Exception while trying to run in simulator %s:\n%s" % (simulator, e)
        ) from e


def _print_info(l):
    print(
        "List %s of %i elements of type %s: %s -> %s (min: %s, max: %s)"
        % (type(l), len(l), type(l[0]), l[0], l[-1], min(l), max(l))
    )


def _print_result_info(traces, events):
    """
    Print a summary of the returned (voltage) traces and spike times
    """
    count = 0
    max_print = 4
    verbose = False
    print_v("Returning %i traces:" % len(traces))
    for r in sorted(traces.keys()):
        x = traces[r]
        if not verbose and count >= max_print:
            if count == max_print:
                print_v("  ...")
                count += 1
        else:
            print_v(
                "  %s (%s): %s -> %s (min: %s, max: %s, len: %i)"
                % (r, type(x), x[0], x[-1], min(x), max(x), len(x))
            )
            count += 1

    count = 0
    print_v("Returning %i events:" % len(events))
    for r in sorted(events.keys()):
        x = events[r]
        if not verbose and count >= max_print:
            if count == max_print:
                print_v("  ...")
                count += 1
        else:
            print_v(
                "  %s: %s -> %s (len: %i)"
                % (r, x[0] if len(x) > 0 else "-", x[-1] if len(x) > 0 else "-", len(x))
            )
            count += 1
