#!/usr/bin/env python

"""
Local Processing Unit (LPU) with plugin support for various neuron/synapse models.
"""
import time
import collections
import pickle
import numbers
import copy
import itertools
import functools
import inspect, glob, os, importlib, sys

from future.utils import iteritems
from past.builtins import long
from builtins import zip


import pycuda.gpuarray as garray
from pycuda.tools import dtype_to_ctype
import pycuda.driver as cuda
from pycuda.compiler import SourceModule
import pycuda.elementwise as elementwise

import numpy as np
import networkx as nx

#import time

# Work around bug in networkx < 1.9 that causes networkx to choke on GEXF
# files with boolean attributes that contain the strings 'True' or 'False'
# (bug already observed in https://github.com/networkx/networkx/pull/971)
nx.readwrite.gexf.GEXF.convert_bool['false'] = False
nx.readwrite.gexf.GEXF.convert_bool['False'] = False
nx.readwrite.gexf.GEXF.convert_bool['true'] = True
nx.readwrite.gexf.GEXF.convert_bool['True'] = True

from neurokernel.mixins import LoggerMixin
from neurokernel.core_gpu import Module, CTRL_TAG, GPOT_TAG, SPIKE_TAG
from neurokernel.tools.gpu import get_by_inds
from neurokernel.plsel import Selector

from types import *
from collections import Counter

from .utils.simpleio import *
from .utils import parray

from . import NDComponents as nd
from .NDComponents.NDComponent import NDComponent
from .MemoryManager import MemoryManager

import pdb

PORT_IN_GPOT = 'port_in_gpot'
PORT_IN_SPK = 'port_in_spk'
PORT_OUT_GPOT = 'port_out_gpot'
PORT_OUT_SPK = 'port_out_spk'

class LPU(Module):
    """
    Create an LPU circuit to simulate

    Parameters
    ----------
    dt: float
        simulation time step
    graph_type: - str
                  specify the method to generate/load the graph
                  options:
                  "obj": networkx.Graph object.
                  "gpickle": a networkx.Graph object pickled into a file by networkx.write_gpickle.
                  "pickle": a networkx.Graph object picked into a byte stream.
                  "gexf": a graph specified by a GEXF file, e.g. by networkx.write_gexf.
                  "function": a functools.partial object that have all appropriate arguments supplied.
                - dict (only for backward compatibility)
                  dictionary of components (nodes in a graph) and their properties
                  the dictionary is generated by `graph_to_dicts` method
    graph_dict: - dict
                  A dictionary specifying how to get/create the graph.
                  Depending on the value of graph_type, the following is required:
                  if graph_type is "obj" - "graph": networkx.Graph
                  if graph_type is "gpickle" - "filename": str
                  if graph_type is "pickle" - "stream": pickled byte object
                  if graph_type is "gexf" - "filename": str
                  if graph_type is "function" - "function": functools.partial object
                  The following is optional:
                  "kwargs": dict - keyword arguments for calling graph_to_dicts
                - str or networkx.Graph or byte or functools.partial
                  Depdending on the graph_type value, the corresponding
                  filename or object can be used
                - list (only for backward compatibility)
                  a list of tuples (pre, post, conn) specifying the edges of
                  a graph and their properties
    device: int
            specifies the device number to use (default: 0)
    input_processors: list
                      a list of InputProcessor objects (default: empty list)
    output_processors: list
                      a list of InputProcessor objects (default: empty list)
    ctrl_tag, gpot_tag, spike_tag: int
        MPI tags that respectively identify messages containing control data,
        graded potential port values, and spiking port values transmitted to
        worker nodes. (see neurokernel.core_gpu.Module)
    manager: bool
        Managerless running mode flag. It False, run Module without a
        manager. (default: True).
    rank_to_id: bidict.bidict
        Mapping between MPI ranks and module object IDs (default: None).
    routing_table: neurokernel.routing_table.RoutingTable
        Routing table describing data connections between modules. If no routing
        table is specified, the module will be executed in isolation. (default: None)
    uid_key: str
             key/property name in the graph that specifies the uid of nodes
             (default: 'id')
    debug: boolean
           run LPU in debug mode (default: False)
    columns: list of str
             Interface port attributes.
             Network port for controlling the module instance.
             (default: ['io', 'type', 'interface'])
    cuda_verbose: boolean
                  print out cuda kernel compilation message (defautl: False)
    time_sync: boolean
               count the time spent on communication (default: False)
    default_dtype: np.dtype
                   default numpy dtype to use for simulation (default: np.double)
    control_inteface: not implemented
    id: str
        Module identifier. If no identifier is specified, a unique
        identifier is automatically generated.
    extra_comps: NDComponent class or module or str or list
                 Provide the definition of NDComponent models to use in simulation.
                 NDComponent class: use the class directly
                 module: search for all NDComponent subclasses within module.
                         if module does not contain any NDComponent subclass
                         and is __init__.py, then walk through the folder and
                         its subfolders to retrieve all NDComponent classes
                         from all modules.
                 str: can be a path to a module or path to a directory.
                      If the former, import all NDComponent subclasses within module.
                      If the latter, import all NDComponent subclasses in modules
                      in the folder and any subfolders.
                 list: any combination of the above types.
                 (default: [])
    print_timing: boolean
                  to print out a detailed timing break out for quick profiling
                  (default: False)
    """
    @staticmethod
    def conv_legacy_graph(g):
        """
        Converts a gexf from legacy neurodriver format to one currently
        supported
        """


        # Find maximum ID in given graph so that we can use it to create new nodes
        # with IDs that don't overlap with those that already exist:
        max_id = 0
        for id in g.nodes():
            if isinstance(id, basestring):
                if id.isdigit():
                    max_id = max(max_id, int(id))
                else:
                    raise ValueError('node id must be an integer')
            elif isinstance(id, numbers.Integral):
                max_id = max(max_id, id)
            else:
                raise ValueError('node id must be an integer')
            gen_new_id = itertools.count(max_id+1).next

        # Create LPU and interface nodes and connect the latter to the former via an
        # Owns edge:
        g_new = nx.MultiDiGraph()

        # Transformation:
        # 1. nonpublic neuron node -> neuron node
        # 2. public neuron node -> neuron node with
        #    output edge to output port
        # 3. input port -> input port
        # 4. synapse edge -> synapse node + 2 edges connecting
        #    transformed original input/output nodes
        edges_to_out_ports = [] # edges to new output port nodes:
        for id, data in g.nodes(data=True):

            # Don't clobber the original graph's data:
            data = copy.deepcopy(data)

            if 'public' in data and data['public']:
                new_id = gen_new_id()
                port_data = {'selector': data['selector'],
                             'port_type': 'spike' if data['spiking'] else 'gpot',
                             'port_io': 'out',
                             'class': 'Port'}
                g_new.add_node(new_id, **port_data)
                edges_to_out_ports.append((id, new_id))
                del data['selector']

            if 'model' in data:
                if data['model'] == 'port_in_gpot':
                    for a in data:
                        if a!='selector': del data[a]
                    data['class'] = 'Port'
                    data['port_type'] = 'gpot'
                    data['port_io'] = 'in'
                elif data['model'] == 'port_in_spk':
                    for a in data:
                        if a!='selector': del data[a]
                    data['class'] = 'Port'
                    data['port_type'] = 'spike'
                    data['port_io'] = 'in'
                else:
                    data['class'] = data['model']

                # Don't need to several attributes that are implicit:
                for a in ['model', 'public', 'spiking','extern']:
                    if a in data: del data[a]

                g_new.add_node(id, **data)

        # Create synapse nodes for each edge in original graph and connect them to
        # the source/dest neuron/port nodes:
        for from_id, to_id, data in g.edges(data=True):
            data = copy.deepcopy(data)
            if data['model'] == 'power_gpot_gpot':
                data['class'] = 'PowerGPotGPot'
            else:
                data['class'] = data['model']
            del data['model']

            if 'id' in data: del data['id']

            new_id = gen_new_id()
            g_new.add_node(new_id, **data)
            g_new.add_edge(from_id, new_id)
            g_new.add_edge(new_id, to_id)

        # Connect output ports to the neurons that emit data through them:
        for from_id, to_id in edges_to_out_ports:
            g_new.add_edge(from_id, to_id)

        return g_new

    @staticmethod
    def graph_to_dicts(graph, uid_key=None, class_key='class',
                       remove_edge_id = True):
        """
        Convert graph of LPU neuron/synapse data to Python data structures.

        Parameters
        ----------
        graph : networkx.MultiDiGraph
            NetworkX graph containing LPU data.

        Returns
        -------
        comp_dict : dict
            A dictionary of components of which
            keys are model names, and
            values are dictionaries of parameters/attributes associated
            with the model.
            Keys of a dictionary of parameters are the names of them,
            and values of corresponding keys are lists of value of parameters.
            One of the parameters is called 'id' and by default it
            uses the id of the node in the graph.
            If uid_key is specified, id will use the specified parameter.
            Therefore, comp_dict has the following structure:

            comp_dict = {}
                comp_dict[model_name_1] = {}
                    comp_dict[model_name_1][parameter_1] = []
                    ...
                    comp_dict[model_name_1][parameter_N] = []
                    comp_dict[model_name_1][id] = []

                ...

                comp_dict[model_name_M] = {}
                    comp_dict[model_name_M][parameter_1] = []
                    ...
                    comp_dict[model_name_M][parameter_N] = []
                    comp_dict[model_name_M][id] = []

        conns : list
            A list of edges contained in graph describing the relation
            between components

        Example
        -------
        TODO: Update

        Notes
        -----
        TODO: Update
        """

        comp_dict = {}
        comps = graph.nodes.items()

        all_component_types = list(set([comp[class_key] for uid, comp in comps]))

        for model in all_component_types:
            sub_comps = [comp for comp in comps \
                                   if comp[1][class_key] == model]

            all_keys = [set(comp[1]) for comp in sub_comps]
            key_intersection = set.intersection(*all_keys)
            key_union = set.union(*all_keys)

            # For visually checking if any essential parameter is dropped
            ignored_keys = list(key_union-key_intersection)
            if ignored_keys:
                print('parameters of model {} ignored: {}'.format(model, ignored_keys))

            del all_keys

            sub_comp_keys = list(key_intersection)

            if model == 'Port':
                assert('selector' in sub_comp_keys)

            comp_dict[model] = {
                k: [comp[k] for uid, comp in sub_comps] \
                for k in sub_comp_keys if not k in [uid_key, class_key]}

            comp_dict[model]['id'] = [comp[uid_key] if uid_key else uid \
                                      for uid, comp in sub_comps]

#        for id, comp in comps:
#            model = comp[class_key]
#
#            # For port, make sure selector is specified
#            if model == 'Port':
#                assert('selector' in comp.keys())
#
#            # if the neuron model does not appear before, add it into n_dict
#            if model not in comp_dict:
#                comp_dict[model] = {k:[] for k in comp.keys() + ['id']}
#
#            # Same model should have the same attributes
#            if not set(comp_dict[model].keys()) == set(comp.keys() + ['id']):
#                raise KeyError("keys of component does not match with that of "+\
#                               model+": "+ str(set(comp_dict[model].keys())) +
#                               str(set(comp.keys() + ['id'])))
#
#            # add data to the subdictionary of comp_dict
#            for key in comp.iterkeys():
#                if not key==uid_key:
#                    comp_dict[model][key].append( comp[key] )
#            if uid_key:
#                comp_dict[model]['id'].append(comp[uid_key])
#            else:
#                comp_dict[model]['id'].append( id )
#
#        # Remove duplicate model information:
#        for val in comp_dict.itervalues(): val.pop(class_key)

        # Extract connections
        conns = graph.edges(data=True)
        if remove_edge_id:
            for pre, post, conn in conns:
                conn.pop('id', None)
                conn.pop('tag', None)
        return comp_dict, conns

    @staticmethod
    def lpu_parser(filename):
        """
        GEXF LPU specification parser.

        Extract LPU specification data from a GEXF file and store it in
        Python data structures.
        TODO: Update

        Parameters
        ----------
        filename : str
            GEXF filename.

        Returns
        -------
        TODO: Update
        """

        graph = nx.read_gexf(filename)
        return LPU.graph_to_dicts(graph, remove_edge_id = True)

    @staticmethod
    def lpu_parser_legacy(filename):
        """
        TODO: Update
        """

        graph = nx.read_gexf(filename)
        return LPU.graph_to_dicts(LPU.conv_legacy_graph(graph))

    @classmethod
    def extract_in_gpot(cls, comp_dict, uid_key):
        """
        Return selectors of non-spiking input ports.
        """
        if not 'Port' in comp_dict: return ('',[])
        a = list(zip(*[(sel,uid) for sel,ptype,io,uid in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_type'],
                             comp_dict['Port']['port_io'],
                             comp_dict['Port'][uid_key]) if ptype=='gpot' \
                  and io=='in']))
        if not a: a = ('',[])
        return a

    @classmethod
    def extract_in_spk(cls, comp_dict, uid_key):
        """
        Return selectors of spiking input ports.
        """
        if not 'Port' in comp_dict: return ('',[])
        a = list(zip(*[(sel,uid) for sel,ptype,io,uid in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_type'],
                             comp_dict['Port']['port_io'],
                             comp_dict['Port'][uid_key]) if ptype=='spike' \
                         and io=='in']))
        if not a: a = ('',[])
        return a

    @classmethod
    def extract_out_gpot(cls, comp_dict, uid_key):
        """
        Return selectors of non-spiking output neurons.
        """
        if not 'Port' in comp_dict: return ('',[])
        a = list(zip(*[(sel,uid) for sel,ptype,io,uid in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_type'],
                             comp_dict['Port']['port_io'],
                             comp_dict['Port'][uid_key]) if ptype=='gpot' \
                  and io=='out']))
        if not a: a = ('',[])
        return a

    @classmethod
    def extract_out_spk(cls, comp_dict, uid_key):
        """
        Return selectors of spiking output neurons.
        """
        if not 'Port' in comp_dict: return ('',[])
        a = list(zip(*[(sel,uid) for sel,ptype,io,uid in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_type'],
                             comp_dict['Port']['port_io'],
                             comp_dict['Port'][uid_key]) if ptype=='spike' \
                  and io=='out']))
        if not a: a = ('',[])
        return a

    @classmethod
    def extract_sel_in_gpot(cls, comp_dict):
        """
        Return selectors of non-spiking input ports.
        """
        if not 'Port' in comp_dict: return ''
        return ','.join([sel  for sel,ptype,io in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_type'],
                             comp_dict['Port']['port_io']) \
                             if ptype=='gpot' and io=='in'])

    @classmethod
    def extract_sel_in_spk(cls, comp_dict):
        """
        Return selectors of spiking input ports.
        """
        if not 'Port' in comp_dict: return ''
        return ','.join([sel  for sel,ptype,io in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_type'],
                             comp_dict['Port']['port_io']) \
                             if ptype=='spike' and io=='in'])

    @classmethod
    def extract_sel_out_gpot(cls, comp_dict):
        """
        Return selectors of non-spiking output neurons.
        """
        if not 'Port' in comp_dict: return ''
        return ','.join([sel  for sel,ptype,io in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_type'],
                             comp_dict['Port']['port_io']) \
                             if ptype=='gpot' and io=='out'])

    @classmethod
    def extract_sel_out_spk(cls, comp_dict):
        """
        Return selectors of spiking output neurons.
        """
        if not 'Port' in comp_dict: return ''
        return ','.join([sel for sel,ptype,io,uid in \
                  zip(comp_dict['Port']['selector'],
                      comp_dict['Port']['port_type'],
                      comp_dict['Port']['port_io']) \
                  if ptype=='spike' and io=='out'])

    @classmethod
    def extract_sel_in(cls, comp_dict):
        """
        Return selectors of all input ports.
        """
        if not 'Port' in comp_dict: return ''
        return ','.join([sel for sel, io in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_io']) if io=='in'])

    @classmethod
    def extract_sel_out(cls, comp_dict):
        """
        Return selectors of all output neurons.
        """

        if not 'Port' in comp_dict: return ''
        return ','.join([sel for sel, io in \
                         zip(comp_dict['Port']['selector'],
                             comp_dict['Port']['port_io']) if io=='out'])

    @classmethod
    def extract_sel_all(cls, comp_dict):
        """
        Return selectors for all input ports and output neurons.
        """

        return ','.join(filter(None, \
                    [cls.extract_in(comp_dict), cls.extract_out(comp_dict)]))

    @classmethod
    def import_NDcomponent_from_path(cls, path):
        cls = []
        if not isinstance(path, str):
            raise TypeError('path must be a str')

        if os.path.isdir(path):
            for root, dirs, files in os.walk(path):
                if os.path.split(root)[-1][0] in ['_','.']:
                    continue
                sys.path.append(root)
                for file in files:
                    m, ext = os.path.splitext(file)
                    if ext == '.py' and not file.startswith('_'):
                        try:
                            mod = importlib.import_module(m, root)
                            cls.extend(inspect.getmembers(
                                        mod,
                                        lambda x: inspect.isclass(x) and \
                                        issubclass(x, NDComponent) and \
                                        x.__module__ == mod.__name__))
                        except ImportError:
                            print('failed to import from {} {}'.format(root, m))
                            pass
                sys.path.pop(-1)
        elif os.path.isfile(path):
            p, file = os.path.split(path)
            m, ext = os.path.splitext(file)
            if ext == '.py' and not file.startswith('_'):
                sys.path.append(p)
                try:
                    mod = importlib.import_module(m, p)
                    cls.extend(inspect.getmembers(
                                mod,
                                lambda x: inspect.isclass(x) and \
                                          issubclass(x, NDComponent) and \
                                          x.__module__ == mod.__name__))
                except ImportError:
                    print('failed to import from {} {}'.format(p, m))
                    pass
                sys.path.pop(-1)
        else:
            raise FileNotFoundError('File or directory does not exist: {}'.format(path))
        return cls

    @classmethod
    def import_NDcomponent_from_module(cls, mod):
        if not inspect.ismodule(mod):
            raise TypeError('mod is not a module')
        cls = inspect.getmembers(
                    mod,
                    lambda x: inspect.isclass(x) and issubclass(x, NDComponent))
        if not len(cls):
            # If module does not contain any NDComponent subclasses,
            # check if the module is an __init__.py
            # if it is, then recursively import all importable NDComponent classes
            # down the folder and subfolders
            path, file = os.path.split(mod.__file__)
            if file == '__init__.py':
                base_package = mod.__package__
                to_remove = len(path.split('/'))
                for root, dirs, files in os.walk(path):
                    if os.path.split(root)[-1][0] in ['_', '.']:
                        continue
                    subpackage = '.'.join(['']+root.split('/')[to_remove:])
                    for file in files:
                        m, ext = os.path.splitext(file)
                        if ext == '.py' and not file.startswith('_'):
                            try:
                                mod = importlib.import_module(
                                        '{}.{}'.format(subpackage, m),
                                        base_package)
                                cls.extend(inspect.getmembers(
                                            mod,
                                            lambda x: inspect.isclass(x) and \
                                            issubclass(x, NDComponent) and \
                                            x.__module__ == mod.__name__))
                            except ImportError:
                                print('failed to import from {} {}'.format(m, subpackage))
                                pass
        return cls

    def import_models(self, models):
        if inspect.isclass(models) and issubclass(models, NDComponent):
            # When an NDComponent class is imported directly
            self._comps[models.__name__] = {'accesses': models.accesses ,
                                            'updates':models.updates,
                                            'cls':models}
        elif inspect.ismodule(models):
            # when a module is provided
            cls = self.import_NDcomponent_from_module(models)
            for name, obj in cls:
                self._comps[name] = {'accesses': obj.accesses,
                                     'updates':obj.updates,
                                     'cls':obj}
        elif isinstance(models, str):
            cls = self.import_NDcomponent_from_path(models)
            for name, obj in cls:
                self._comps[name] = {'accesses': obj.accesses,
                                     'updates':obj.updates,
                                     'cls':obj}
        elif isinstance(models, (list, tuple)):
            for model in models:
                self.import_models(model)


    def __init__(self, dt, graph_type, graph_dict, device=0, input_processors=[],
                 output_processors=[], ctrl_tag=CTRL_TAG, gpot_tag=GPOT_TAG,
                 spike_tag=SPIKE_TAG, manager = True,
                 rank_to_id=None, routing_table=None,
                 uid_key='id', debug=False, columns=['io', 'type', 'interface'],
                 cuda_verbose=False, time_sync=False, default_dtype=np.double,
                 control_inteface=None, id=None, extra_comps=[],
                 print_timing=False):
        compilation_start_time = time.time()

        LoggerMixin.__init__(self, 'LPU {}'.format(id))

        assert('io' in columns)
        assert('type' in columns)
        assert('interface' in columns)
        self.LPU_id = id
        self.dt = dt
        self.step = 0
        self._debug = debug
        self.device = device
        self._default_dtype = default_dtype
        self.control_inteface = control_inteface
        self._print_timing = print_timing
        if cuda_verbose:
            self._compile_options = ['--ptxas-options=-v']
        else:
            self._compile_options = []

        if self._print_timing:
            start = time.time()

        if not isinstance(input_processors, list):
            input_processors = [input_processors]
        if not isinstance(output_processors, list):
            output_processors = [output_processors]

        self.output_processors = output_processors
        self.input_processors = input_processors

        G = nx.MultiDiGraph()
        for input_processor in self.input_processors:
            G = nx.compose(G, input_processor.get_new_graph())

        if isinstance(graph_type, dict):
        # for backward compatibility where the two arguments are comp_dict and conn_list
            comp_dict = graph_type
            conn_list = graph_dict
        elif isinstance(graph_type, str):
            if graph_type == 'obj':
                # use networkx.Graph objects,
                # as slow as using comp_dict, conn_list directly
                if isinstance(graph_dict, nx.Graph):
                    kwargs = {}
                    graph = graph_dict
                elif isinstance(graph_dict, dict):
                    kwargs = graph_dict.get('kwargs', {})
                    graph = graph_dict['graph']
                else:
                    raise TypeError("For graph_type = 'obj', graph must be networkx Graphs")
            elif graph_type == 'gpickle':
                # use networkx built-in pickle function,
                # relatively fast but require disk space
                if isinstance(graph_dict, str):
                    kwargs = {}
                    graph = nx.readwrite.gpickle.read_gpickle(graph_dict)
                elif isinstance(graph_dict, dict):
                    kwargs = graph_dict.get('kwargs', {})
                    graph = nx.read_gpickle(graph_dict['filename'])
                else:
                    raise TypeError('graph_dict must be either str (filename) or dict')
            elif graph_type == 'pickle':
                # use pickled byte stream of networkx.Graph
                # fast
                if isinstance(graph_dict, bytes):
                    kwargs = {}
                    graph = pickle.loads(graph_dict)
                elif isinstance(graph_dict, dict):
                    kwargs = graph_dict.get('kwargs', {})
                    graph = pickle.loads(graph_dict['stream'])
                else:
                    raise TypeError('graph_dict must be either pickled byte stream or dict')
            elif graph_type == 'gexf':
                # use networkx build-in read_gexf
                # slow and require disk space
                if isinstance(graph_dict, str):
                    kwargs = {'remove_edge_id': True}
                    graph = nx.read_gexf(graph_dict)
                elif isinstance(graph_dict, dict):
                    kwargs = graph_dict.get('kwargs', {'remove_edge_id': True})
                    graph = nx.read_gexf(graph_dict['filename'])
                else:
                    raise TypeError('graph_dict must be either str (filename) or dict')
            elif graph_type == 'function':
                # call a function created by functools.partial method,
                # where the function and arguments
                # fast and generate the graph during initialization
                if isinstance(graph_dict, (functools.partial, functools.partialmethod)):
                    kwargs = {}
                    graph = graph_dict()
                elif isinstance(graph_dict, dict):
                    kwargs = graph_dict.get('kwargs', {})
                    graph = graph_dict['function']()
            else:
                raise TypeError('Unkown method to generate Graph')
            comp_dict, conn_list = self.graph_to_dicts(nx.compose(graph, G), **kwargs)
        else:
            raise TypeError('Unkown method to generate Graph')

        if self._print_timing:
            self.log_info("Elapsed time for converting graph: {:.3f} seconds".format(time.time()-start))
            start = time.time()

        self._uid_key = uid_key

        self._uid_generator = uid_generator()

        # Load all NDComponents:
        self._comps = {}
        self.import_models(nd)
        if extra_comps:
            self.import_models(extra_comps)

        if self._print_timing:
            self.log_info("Elapsed time for importing component definitions: {:.3f} seconds".format(time.time()-start))
            start = time.time()

        # self._comps = {cls.__name__:{'accesses': cls.accesses ,
        #                              'updates':cls.updates,
        #                              'cls':cls} \
        #                for cls in comp_classes if not cls.__name__[:4]=='Base'}
        # self._load_components(extra_comps=extra_comps)

        # Ignore models without implementation
        models_to_be_deleted = []
        for model in comp_dict:
            if not model in self._comps and not model in ['Port','Input']:
                self.log_info("Ignoring Model %s: Can not find implementation"
                              % model)
                raise NotImplementedError("Model %s: Can not find implementation" % model)
                models_to_be_deleted.append(model)
        for model in models_to_be_deleted:
            del comp_dict[model]

        # Assume zero delay by default
        self._variable_delay_map = {}

        # Generate a uid to model map of components
        self._uid_model_map = {}
        for model,attribs in iteritems(comp_dict):
            for i,uid in enumerate(attribs[uid_key]):
                self._uid_model_map[uid] = model


        # Map from post synaptic component to aggregator uid
        agg_map = {}
        agg = {}

        #start = time.time()

        conns = []
        self._in_port_vars = {}
        self._out_port_conns = []
        comp_uid_order = {}
        for model, attribs in comp_dict.items():
            comp_uid_order[model] = {uid: i for i, uid in enumerate(attribs[uid_key])}

        if self._print_timing:
            self.log_info("Elapsed time for processing comp_dict: {:.3f} seconds".format(time.time()-start))
            start = time.time()
        # Process connections between components, remove inconsitent connections
        # calculate required delays, infer variable if required
        in_port_vars_set = {}
        reverse_set = ['reverse','Vr','VR','reverse_potential']

        comp_updates = {name: v['updates'] if not name=='Port' else [] \
                        for name, v in self._comps.items()}
        comp_updates.update({'Port': []})
        comp_accesses = {name: v['accesses'] if not name=='Port' else [] \
                         for name, v in self._comps.items()}
        comp_accesses.update({'Port': []})

        for pre, post, data in conn_list:
            try:
                pre_model = self._uid_model_map[pre]
            except KeyError:
                continue
            try:
                post_model = self._uid_model_map[post]
            except KeyError:
                continue

            pre_updates = comp_updates[pre_model]
            pre_updates_set = set(pre_updates)
            post_accesses = comp_accesses[post_model]
            post_accesses_set = set(post_accesses)
            pre_post_intersection = pre_updates_set&post_accesses_set
            g_in_pre_update = 'g' in pre_updates_set # not so useful
            V_in_pre_update = 'V' in pre_updates_set

            # Update delay
            delay = max(int(round((data['delay']/dt))) \
                    if 'delay' in data else 0, 1) - 1
            data['delay'] = delay

            if post_model == 'Aggregator':
                agg_map[post] = post
                reverse_key = None
                for k in reverse_set:
                    if k in data:
                        reverse_key = k
                        reverse = data[reverse_key]
                        break
                if reverse_key is None:
                    # else look in the attibutes of the synapse
                    s = (set(['reverse','Vr','VR','reverse_potential'])&
                         set(comp_dict[pre_model]))
                    if s: reverse_key = s.pop()
                    if reverse_key:
                        reverse = comp_dict[pre_model][reverse_key][\
                                comp_uid_order[pre_model][pre]]
                        if g_in_pre_update:
                            data['reverse'] = reverse
                    else:
                        if g_in_pre_update:
                            self.log_info('Assuming reverse potential ' +
                                          'to be zero for connection from' +
                                          '%s to %s'%(pre,post))
                            data['reverse'] = 0
                        reverse = 0

                if post in agg:
                    if g_in_pre_update:
                        agg[post].append({'pre':pre,'reverse':reverse,
                                          'variable':'g'})
                    elif V_in_pre_update:
                        agg[post].append({'pre':pre, 'variable':'V'})
                else:
                    # Make sure aggregator has access to postsynaptic voltage
                    if g_in_pre_update:
                        agg[post] = [{'pre':pre,'reverse':reverse,'variable':'g'}]
                    elif V_in_pre_update:
                        agg[post] = [{'pre':pre,'variable':'V'}]

                if g_in_pre_update:
                    agg[post][-1].update(data)
                    self._variable_delay_map['g'] = max(data['delay'],
                                    self._variable_delay_map['g'] if 'g' in \
                                    self._variable_delay_map else 0)


            # Ensure consistency
            # Insert Aggregator between g->V if required. Assume 'reverse' or
            # 'Vr' or 'VR' or 'reverse_potential' id present as a param in the
            # synapse in that case
            if not pre_post_intersection:
                if g_in_pre_update and 'I' in post_accesses_set:
                    # start2 = time.time()
                    # First look for reverse in the attributes of the edge
                    reverse_key = None
                    for k in reverse_set:
                        if k in data:
                            reverse_key = k
                            reverse = data[reverse_key]
                            break
                    if reverse_key is None:
                        # else look in the attibutes of the synapse
                        s = (set(['reverse','Vr','VR','reverse_potential'])&
                             set(comp_dict[pre_model]))
                        if s: reverse_key = s.pop()
                        if reverse_key:
                            reverse = comp_dict[pre_model][reverse_key][\
                                    comp_uid_order[pre_model][pre]]
                        else:
                            self.log_info('Assuming reverse potential ' +
                                          'to be zero for connection from' +
                                          '%s to %s'%(pre,post))
                            reverse = 0

                    data.update({'pre':pre,'reverse':reverse,
                                 'variable':'g'})
                    if post in agg:
                        agg[post].append(data)
                    else:
                        # Make sure aggregator has access to postsynaptic voltage
                        agg[post] = [{'pre':post,'variable':'V'},
                                     data]
                    if post not in agg_map:
                        uid = self._uid_generator.generate_uid()
                        agg_map[post] = uid
                    self._variable_delay_map['g'] = max(data['delay'],
                                    self._variable_delay_map['g'] if 'g' in \
                                    self._variable_delay_map else 0)

                elif pre_model == 'Port':
                    if not 'variable' in data:
                        data['variable'] = post_accesses[0]
                    if not data['variable'] in self._in_port_vars:
                        self._in_port_vars[data['variable']] = []
                        in_port_vars_set[data['variable']] = set()
                    if pre not in in_port_vars_set[data['variable']]:
                        self._in_port_vars[data['variable']].append(pre)
                        in_port_vars_set[data['variable']].add(pre)
                    conns.append((pre, post, data))
                    self._variable_delay_map[data['variable']] = max(data['delay'],
                            self._variable_delay_map[data['variable']] if \
                            data['variable'] in self._variable_delay_map else 0)
                elif post_model == 'Port':
                    if not 'variable' in data:
                        data['variable'] = pre_updates[0]
                    self._out_port_conns.append((pre, post, data['variable']))
                else:
                    self.log_info("Ignoring connection %s -> %s"%(pre,post))
                continue

            var = data['variable'] if 'variable' in data else None
            if not var:
                var = pre_post_intersection.pop()
            elif not (var in pre_updates_set and var in post_accesses_set):
                continue
            data['variable'] = var
            self._variable_delay_map[data['variable']] = max(data['delay'],
                            self._variable_delay_map[data['variable']] if \
                            data['variable'] in self._variable_delay_map else 0)
            # connection to Aggregator will be added later
            if 'Aggregator' not in post_model:
                conns.append((pre,post,data))

        if self._print_timing:
            self.log_info("Elapsed time for processing conn_list: {:.3f} seconds".format(time.time()-start))
            start = time.time()

        if agg and not 'Aggregator' in comp_dict:
            comp_dict['Aggregator'] = {uid_key: []}
        if agg:
            agg_uid_key_set = set(comp_dict['Aggregator'][uid_key])

        # Add updated aggregator components to component dictionary
        # and create connections for aggregator
        for post, conn_list in agg.items():
            uid = agg_map[post]
            if uid not in agg_uid_key_set:
                keys = [k for k in comp_dict['Aggregator'] if k != uid_key]
                comp_dict['Aggregator'][uid_key].append(uid)
                agg_uid_key_set.add(uid)
                self._uid_model_map[uid] = 'Aggregator'
                for k in keys:
                    comp_dict['Aggregator'][k].append(str(uid))
            for conn in conn_list:
                conns.append((conn['pre'],uid,{k:v for k,v in conn.items()
                                               if k!='pre'}))
            # Add a 'I' connection between Aggregator and neuron if they are
            # automatically generated.
            # This can be checking if the 'pre' attribute in the item
            # in conn_list with 'variable' 'V' is the same neuron as post
            if post == [tmp['pre'] for tmp in conn_list if tmp['variable']=='V'][0]:
                conns.append((uid,post,{'variable':'I', 'delay': 0}))
        if self._print_timing:
            self.log_info("Elapsed time for processing aggregator: {:.3f} seconds".format(time.time()-start))
            start = time.time()

        self.conn_dict = {}

        # RePackage connections
        for (pre, post, data) in conns:
            if not post in self.conn_dict:
                self.conn_dict[post] = {}
            var = data['variable']
            data.pop('variable')
            if not var in self.conn_dict[post]:
                self.conn_dict[post][var] = {k:[] for k in ['pre'] + list(data)}
            self.conn_dict[post][var]['pre'].append(pre)
            for k in data: self.conn_dict[post][var][k].append(data[k])

        if self._print_timing:
            self.log_info("Elapsed time for repackaging connections: {:.3f} seconds".format(time.time()-start))
            start = time.time()

        # Add connections for component with no incoming connections
        for uid, model in iteritems(self._uid_model_map):
            if model == 'Port':
                continue
            for var in self._comps[model]['accesses']:
                if ((not uid in self.conn_dict or not var in self.conn_dict[uid])):
                    pre = self._uid_generator.generate_uid(input=True)
                    if not var in self._variable_delay_map:
                        self._variable_delay_map[var]=0
                    if not uid in self.conn_dict: self.conn_dict[uid] = {}
                    if model == 'Aggregator' and var == 'g':
                        self.conn_dict[uid][var] = {'pre':[pre],'delay':[0],
                                                    'reverse': [0]} #'id': [0],
                    else:
                        self.conn_dict[uid][var] = {'pre':[pre],'delay':[0]}
                    if not 'Input' in comp_dict:
                        comp_dict['Input'] = {}
                    if not var in comp_dict['Input']:
                        comp_dict['Input'][var] = {self._uid_key: []}
                    comp_dict['Input'][var][self._uid_key].append(pre)

        if self._print_timing:
            self.log_info("Elapsed time for adding connections for component with no incoming connections: {:.3f} seconds".format(time.time()-start))
            start = time.time()

        # Optimize ordering (TODO)
        self._uid_ind_map = {m:{uid:i for i,uid in enumerate(n[uid_key])}
                            for m,n in comp_dict.items() if not m=='Input'}

        if 'Input' in comp_dict:
            self._uid_ind_map['Input'] = {var:{uid:i for i, uid in enumerate(d[uid_key])}
                                         for var, d in comp_dict['Input'].items()}

        # Reorder components
        for m, n in comp_dict.items():
            if m=='Input':
                for var, d in n.items():
                    order = np.argsort([self._uid_ind_map[m][var][uid] for uid in d[uid_key]])
                    d[uid_key] = [d[uid_key][i] for i in order]
                continue

            order = np.argsort([self._uid_ind_map[m][uid] for uid in n[uid_key]])
            for k in n:
                n[k] = [n[k][i] for i in order]

        # Reorder input port variables
        for var, uids in self._in_port_vars.items():
            order = np.argsort([self._uid_ind_map['Port'][uid] for uid in uids])
            self._in_port_vars[var] = [uids[i] for i in order]

        if self._print_timing:
            self.log_info("Elapsed time for optimizing ordering: {:.3f} seconds".format(time.time()-start))
            start = time.time()

        # Try to figure out order of stepping through components
        # If a loop of dependencies is present, update order behaviour is undefined
        models = list(comp_dict)
        try:
            models.remove('Port')
        except:
            pass

        try:
            models.remove('Input')
        except:
            pass


        deps = {i:[] for i in range(len(models))}
        for i in range(len(models)):
            for j in range(i+1,len(models)):
                in12 = set(self._comps[models[i]]['updates'])&\
                       set(self._comps[models[j]]['accesses'])
                in21 = set(self._comps[models[i]]['accesses'])&\
                       set(self._comps[models[j]]['updates'])
                if in12 or in21:
                    if len(in12) > len(in21):
                        deps[j].append(i)
                    else:
                        deps[i].append(j)


        self._exec_order = []
        for i, model in enumerate(models):
            if not model in self._exec_order: self._exec_order.append(model)
            for j in deps[i]:
                try:
                    if self._exec_order.index(models[j]) > \
                       self._exec_order.index(model):
                        self._exec_order.remove(models[j])
                        self._exec_order.insert(self._exec_order.index(model),
                                               models[j])
                except ValueError:
                    self._exec_order.insert(self._exec_order.index(model),
                                           models[j])

        var_mod = {}
        for i, model in enumerate(models):
            for var in self._comps[model]['updates']:
                if not var in var_mod: var_mod[var] = []
                var_mod[var].append(model)

        self._model_var_inj = {}
        for var, models in var_mod.items():
            i = 0
            for model in models:
                i = max(self._exec_order.index(model),i)
            if not self._exec_order[i] in self._model_var_inj:
                self._model_var_inj[self._exec_order[i]] = []
            self._model_var_inj[self._exec_order[i]].append(var)

        #Variables not updated by any component (for example those coming from
        #external input or Ports) are slated to be injected at the end of a step
        for var in self._variable_delay_map:
            if not var in var_mod:
                if not self._exec_order[-1] in self._model_var_inj:
                    self._model_var_inj[self._exec_order[-1]] = []
                self._model_var_inj[self._exec_order[-1]].append(var)

        if self._print_timing:
            self.log_info("Elapsed time for ordering execution: {:.3f} seconds".format(time.time()-start))
            start = time.time()
        # Get selectors of input ports:
        sel_in_gpot, self._in_gpot_uids = self.extract_in_gpot(comp_dict,
                                                              self._uid_key)
        self._sel_in_gpot = Selector(','.join(sel_in_gpot))
        sel_in_spk, self._in_spk_uids = self.extract_in_spk(comp_dict,
                                                           self._uid_key)
        self._sel_in_spk = Selector(','.join(sel_in_spk))
        sel_in = Selector.add(self._sel_in_gpot, self._sel_in_spk)

        # Get selectors of output neurons:
        sel_out_gpot, self._out_gpot_uids = self.extract_out_gpot(comp_dict,
                                                                 self._uid_key)
        self._sel_out_gpot = Selector(','.join(sel_out_gpot))
        sel_out_spk, self._out_spk_uids = self.extract_out_spk(comp_dict,
                                                              self._uid_key)
        self._sel_out_spk = Selector(','.join(sel_out_spk))
        sel_out = Selector.add(self._sel_out_gpot, self._sel_out_spk)
        sel_gpot = Selector.add(self._sel_in_gpot, self._sel_out_gpot)
        sel_spk = Selector.add(self._sel_in_spk, self._sel_out_spk)
        sel = Selector.add(sel_gpot, sel_spk)
        if self._print_timing:
            self.log_info("Elapsed time for generating selectors: {:.3f} seconds".format( time.time()-start))
            start = time.time()

        # Save component parameters data in the form
        # [('Model0', {'attrib0': [..], 'attrib1': [..]}), ('Model1', ...)]
        self.comp_list = comp_dict.items()
        self.models = {m:i for i,(m,_) in enumerate(self.comp_list)}

        # Number of components of each model:
        self.model_num = [len(n[uid_key]) if not m=='Input' else
                          len(sum([d[uid_key] for d in n.values()],[]))
                          for m, n in self.comp_list]

        data_gpot = np.zeros(len(self._in_gpot_uids)+len(self._out_gpot_uids),
                             self._default_dtype)
        data_spike = np.zeros(len(self._in_spk_uids)+len(self._out_spk_uids),
                              self._default_dtype)

        if self._print_timing:
            self.log_info("Elapsed time for misc: {:.3f} seconds".format(time.time()-start))
            start = time.time()
        super(LPU, self).__init__(sel=sel, sel_in=sel_in, sel_out=sel_out,
                                  sel_gpot=sel_gpot, sel_spike=sel_spk,
                                  data_gpot=data_gpot, data_spike=data_spike,
                                  columns=columns, ctrl_tag=ctrl_tag, gpot_tag=gpot_tag,
                                  spike_tag=spike_tag, id=self.LPU_id,
                                  manager = manager,
                                  rank_to_id=rank_to_id, routing_table=routing_table,
                                  device=device, debug=debug, time_sync=time_sync,
                                  print_timing=print_timing)

        for model in comp_dict:
            if model == 'Input':
                print('{}: Number of {}: {}'.format(self.id, model, {k: len(v[self._uid_key]) for k, v in comp_dict[model].items()}))
                self.log_info('Number of {}: {}'.format(model, {k: len(v[self._uid_key]) for k, v in comp_dict[model].items()}.__repr__().replace('{','').replace('}','')))
            else:
                print('{}: Number of {}: {}'.format(self.id, model, len(comp_dict[model][self._uid_key])))
                self.log_info('Number of {}: {}'.format(model, len(comp_dict[model][self._uid_key])))

        if self._print_timing:
            cuda.Context.synchronize()
            self.log_info("Elapsed time for initializing parent class: {:.3f} seconds".format(time.time()-start))


        # Integer indices in port map data arrays corresponding to input/output
        # gpot/spiking ports:
        self._in_gpot_inds = np.array(self.pm['gpot'].ports_to_inds(
                                     self._sel_in_gpot), dtype=np.int32)
        self._out_gpot_inds = np.array(self.pm['gpot'].ports_to_inds(
                                      self._sel_out_gpot), dtype=np.int32)
        self._in_spk_inds = np.array(self.pm['spike'].ports_to_inds(
                                    self._sel_in_spk), dtype=np.int32)
        self._out_spk_inds = np.array(self.pm['spike'].ports_to_inds(
                                     self._sel_out_spk), dtype=np.int32)

        self._compile_time = time.time()-compilation_start_time

    # def generate_uid(self, input=False):
    #     if input:
    #         uid = 'input_' + str(np.random.randint(100000))
    #     else:
    #         uid = 'auto_' + str(np.random.randint(100000))
    #     while uid in self.gen_uids:
    #         if input:
    #             uid = 'input_' + str(np.random.randint(100000))
    #         else:
    #             uid = 'auto_' + str(np.random.randint(100000))
    #     return uid

    def pre_run(self):
        pre_run_start_time = time.time()
        if self._print_timing:
            start = time.time()
        super(LPU, self).pre_run()
        if self._print_timing:
            self.log_info("LPU pre_run parent took {} seconds".format(time.time()-start))
            start = time.time()

        self.memory_manager = MemoryManager()
        self.init_variable_memory()
        if self._print_timing:
            cuda.Context.synchronize()
            self.log_info('Elapsed time for initialing variable memory: {:.3f} seconds'.format( time.time()-start))
            start = time.time()

        self.process_connections()
        if self._print_timing:
            self.log_info('Elapsed time for process_connections: {:.3f} seconds'.format(time.time()-start))
            start = time.time()

        self.init_parameters()
        if self._print_timing:
            cuda.Context.synchronize()
            self.log_info('Elapsed time for init_paramseters: {:.3f} seconds'.format( time.time()-start))
            start = time.time()

        self.components = {}
        # Instantiate components
        for model in self.models:
            if model in ['Port','Input']: continue
            self.components[model] = self._instantiate_component(model)
            update_pointers = {}
            for var in self._comps[model]['updates']:
                buff = self.memory_manager.get_buffer(var)
                mind = self.memory_manager.variables[var]['models'].index(model)
                shift = self.memory_manager.variables[var]['cumlen'][mind]
                update_pointers[var] = int(buff.gpudata)+(buff.current*buff.ld+\
                                            shift)*buff.dtype.itemsize
            self.components[model].pre_run(update_pointers)
            for var in self._comps[model]['updates']:
                buff = self.memory_manager.get_buffer(var)
                mind = self.memory_manager.variables[var]['models'].index(model)
                shift = self.memory_manager.variables[var]['cumlen'][mind]
                for j in range(buff.buffer_length):
                    if j is not buff.current:
                        cuda.memcpy_dtod(
                            int(int(buff.gpudata)+(j*buff.ld+\
                                                shift)*buff.dtype.itemsize),
                            int(int(buff.gpudata)+(buff.current*buff.ld+\
                                                shift)*buff.dtype.itemsize),
                            int(buff.dtype.itemsize*self.model_num[self.models[model]]))
        if self._print_timing:
            cuda.Context.synchronize()
            self.log_info('Elapsed time for instantiating components: {:.3f} seconds'.format(time.time()-start))
            start = time.time()
        # Setup ports
        self._setup_input_ports()
        self._setup_output_ports()

        if self._print_timing:
            cuda.Context.synchronize()
            self.log_info('Elapsed time for setting up ports: {:.3f} seconds'.format( time.time()-start))
            start = time.time()

        for p in self.input_processors:
            p.LPU_obj = self
            p._pre_run()

        for p in self.output_processors:
            p.LPU_obj = self
            p._pre_run()

        if self._print_timing:
            cuda.Context.synchronize()
            self.log_info('Elapsed time for prerun input and output processors: {:.3f} seconds'.format( time.time()-start))
            start = time.time()

        self.memory_manager.precompile_fill_zeros()

        if self.control_inteface: self.control_inteface.register(self)
        if self._print_timing:
            cuda.Context.synchronize()
            self.log_info("Elapsed time for LPU pre_run: {:.3f} seconds".format(time.time()-pre_run_start_time))

        self._compile_time += time.time()-pre_run_start_time

        if self._print_timing:
            self.timing = {'read_input': 0, 'input_processors': 0, 'inject_input': 0,
                           'model_run': 0, 'output_processors': 0,
                           'extract_output': 0, 'total': 0}
        cuda.Context.synchronize()
        self.log_info("Total compiling time: {:.3f} seconds".format(self._compile_time))
        print('Compilation of executable circuit completed in {} seconds'.format(
                self._compile_time))

    # TODO: optimize the order of self._out_port_conns beforehand
    def _setup_output_ports(self):
        self._out_port_inds_gpot = {}
        self._out_var_inds_gpot = {}
        self._out_port_inds_spk = {}
        self._out_var_inds_spk = {}
        # assuming that the UIDs are unique
        out_gpot_index = {uid: i for i, uid in enumerate(self._out_gpot_uids)}
        out_spk_index = {uid: i for i, uid in enumerate(self._out_spk_uids)}
        for pre_uid, post_uid, var in self._out_port_conns:
            if not var in self._out_port_inds_gpot:
                self._out_port_inds_gpot[var] = []
                self._out_var_inds_gpot[var] = []
                self._out_port_inds_spk[var] = []
                self._out_var_inds_spk[var] = []
            ind = self.memory_manager.variables[var]['uids'][pre_uid]
            if post_uid in out_gpot_index:
                # self._out_port_inds_gpot[var].append(self._out_gpot_inds[\
                #                             self._out_gpot_uids.index(post_uid)])
                self._out_port_inds_gpot[var].append(self._out_gpot_inds[\
                                            out_gpot_index[post_uid]])
                self._out_var_inds_gpot[var].append(ind)
            else:
                self._out_port_inds_spk[var].append(self._out_spk_inds[\
                                            out_spk_index[post_uid]])
                self._out_var_inds_spk[var].append(ind)

        tmp = self._out_port_inds_gpot.copy()
        for var in tmp:
            if not self._out_port_inds_gpot[var]:
                del self._out_port_inds_gpot[var]
                del self._out_var_inds_gpot[var]
            else:
                self._out_port_inds_gpot[var] = garray.to_gpu(\
                        np.array(self._out_port_inds_gpot[var],np.int32))
                self._out_var_inds_gpot[var] = garray.to_gpu(\
                        np.array(self._out_var_inds_gpot[var],np.int32))

        tmp = self._out_port_inds_spk.copy()
        for var in tmp:
            if not self._out_port_inds_spk[var]:
                del self._out_port_inds_spk[var]
                del self._out_var_inds_spk[var]
            else:
                self._out_port_inds_spk[var] = garray.to_gpu(\
                        np.array(self._out_port_inds_spk[var],np.int32))
                self._out_var_inds_spk[var] = garray.to_gpu(\
                        np.array(self._out_var_inds_spk[var],np.int32))

    def _setup_input_ports(self):
        self.port_inds_gpot = {}
        self.var_inds_gpot = {}
        self.port_inds_spk = {}
        self.var_inds_spk = {}
        # assuming that the UIDs are unique
        in_gpot_index = {uid: i for i, uid in enumerate(self._in_gpot_uids)}
        in_spk_index = {uid: i for i, uid in enumerate(self._in_spk_uids)}
        for var, uids in self._in_port_vars.items():
            self.port_inds_gpot[var] = []
            self.var_inds_gpot[var] = []
            self.port_inds_spk[var] = []
            self.var_inds_spk[var] = []
            mind = self.memory_manager.variables[var]['models'].index('Port')
            shift = self.memory_manager.variables[var]['cumlen'][mind]
            # The following assumes the intersection of set of variables
            # accessed via spiking with those accessed via gpot ports is null
            for i,uid in enumerate(uids):
                if uid in in_gpot_index:
                    self.port_inds_gpot[var].append(self._in_gpot_inds[\
                                            in_gpot_index[uid]])
                    self.var_inds_gpot[var].append(i + shift)
                else:
                    self.port_inds_spk[var].append(self._in_spk_inds[\
                                            in_spk_index[uid]])
                    self.var_inds_spk[var].append(i + shift)
        tmp = self.port_inds_gpot.copy()
        for var in tmp:
            if not self.port_inds_gpot[var]:
                del self.port_inds_gpot[var]
                del self.var_inds_gpot[var]
            else:
                self.port_inds_gpot[var] = garray.to_gpu(\
                        np.array(self.port_inds_gpot[var],np.int32))
                self.var_inds_gpot[var] = garray.to_gpu(\
                        np.array(self.var_inds_gpot[var],np.int32))
        tmp = self.port_inds_spk.copy()
        for var in tmp:
            if not self.port_inds_spk[var]:
                del self.port_inds_spk[var]
                del self.var_inds_spk[var]
            else:
                self.port_inds_spk[var] = garray.to_gpu(\
                        np.array(self.port_inds_spk[var],np.int32))
                self.var_inds_spk[var] = garray.to_gpu(\
                        np.array(self.var_inds_spk[var],np.int32))


    def init_parameters(self):
        for m, n in self.comp_list:
            if not m in ['Port','Input']:
                nn = n.copy()
                nn.pop(self._uid_key)
                # copy integer and boolean parameters into separate dictionary
                nn_int = {k:v for k, v in iteritems(nn) if (isinstance(v, list)
                            and len(v) and type(v[0]) in [int, bool])}
                nn_rest = {k:v for k, v in iteritems(nn) if (
                           (not isinstance(v, list)) or (len(v) and
                           type(v[0]) not in [int, long, bool]))}
                if nn_int:
                    self.memory_manager.params_htod(m, nn_int, np.int32)
                if nn_rest:
                    self.memory_manager.params_htod(m, nn_rest,
                                                    self._default_dtype)

    def init_variable_memory(self):
        var_info = {}
        for (model, attribs) in self.comp_list:
            if model in ['Port']: continue
            # Add memory for external inputs if required
            if model == 'Input':
                for var, d in iteritems(attribs):
                    if not var in var_info:
                        var_info[var] = {'models':[],'len':[],'delay':0,'uids':[]}
                    var_info[var]['models'].append('Input')
                    var_info[var]['len'].append(len(d[self._uid_key]))
                    var_info[var]['uids'].extend(d[self._uid_key])
                continue
            for var in self._comps[model]['updates']:
                if not var in var_info:
                    var_info[var] = {'models':[],'len':[],'delay':0,'uids':[]}
                var_info[var]['models'].append(model)
                var_info[var]['len'].append(len(attribs[self._uid_key]))
                var_info[var]['uids'].extend(attribs[self._uid_key])

        # Add memory for input ports
        for var in self._in_port_vars:
            if not var in var_info:
                var_info[var] = {'models':[],'len':[],'delay':0,'uids':[]}
            var_info[var]['models'].append('Port')
            var_info[var]['len'].append(len(self._in_port_vars[var]))
            var_info[var]['uids'].extend(self._in_port_vars[var])



        for var in self._variable_delay_map:
            var_info[var]['delay'] = self._variable_delay_map[var]

        for var, d in var_info.items():
            d['cumlen'] = np.cumsum([0]+d['len'])
            d['uids'] = {uid:i for i, uid in enumerate(d['uids'])}

            if isinstance(var, str):
                size = d['cumlen'][-1]
            elif isinstance(var, tuple):
                size = d['cumlen'][-1] * len(var)

            self.memory_manager.memory_alloc(var, size, d['delay']+2,\
                dtype=self._default_dtype,
                info=d)

    def process_connections(self):
        for (model, attribs) in self.comp_list:
            if model in ['Port','Input']: continue
            pre = {var:[] for var in self._comps[model]['accesses']}
            npre = {var:[] for var in self._comps[model]['accesses']}
            data = {var:{} for var in self._comps[model]['accesses']}
            for uid in attribs[self._uid_key]:
                cnt = {var:0 for var in self._comps[model]['accesses']}
                if uid in self.conn_dict:
                    for var in self.conn_dict[uid]:
                        for i in range(len(self.conn_dict[uid][var]['pre'])):
                            # Figure out index of the precomponent in the
                            # particular variable memory
                            p = self.conn_dict[uid][var]['pre'][i]
                            ind = self.memory_manager.variables[var]['uids'][p]
                            pre[var].append(ind)

                            cnt[var] += 1
                            for k in self.conn_dict[uid][var]:
                                if k in ['pre','variable']: continue
                                if k not in data[var]: data[var][k] = []
                                data[var][k].append(self.conn_dict[uid][var][k][i])
                            l = len(pre[var])
                            assert(all([len(data[var][k])==l for k in data[var]]))
                    for var,c in cnt.items():
                        npre[var].append(cnt[var])
                else:
                    for n in npre.values(): n.append(0)
            cumpre = {var: np.cumsum([0]+n) for var, n in npre.items()}

            attribs['pre'] = pre
            attribs['cumpre'] = cumpre
            attribs['npre'] = npre
            attribs['conn_data'] = data

    def post_run(self):
        for comp in self.components.values():
            comp.post_run()
        # Cycle through IO processors as well
        for p in self.input_processors: p.post_run()
        for p in self.output_processors: p.post_run()
        super(LPU, self).post_run()
        if self._print_timing:
            print('time spent on:', self.timing)

    def run_step(self):
        super(LPU, self).run_step()


        # Update input ports
        if self._print_timing:
            start_all = time.time()
            start = time.time()
        self._read_LPU_input()
        if self._print_timing:
            cuda.Context.synchronize()
            self.timing['read_input'] += time.time()-start


        # Fetch updated input if available from all input processors
        if self._print_timing:
            start = time.time()
        for p in self.input_processors: p.run_step()
        if self._print_timing:
            cuda.Context.synchronize()
            self.timing['input_processors'] += time.time()-start

        if self._print_timing:
            start = time.time()
        for model in self._exec_order:
            if model in self._model_var_inj:
                for var in self._model_var_inj[model]:
                    # Reset memory for external input to zero if present
                    self.memory_manager.fill_zeros(model='Input', variable=var)
                    for p in self.input_processors:
                        p.inject_input(var)
        if self._print_timing:
            cuda.Context.synchronize()
            self.timing['inject_input'] += time.time()-start

        # Call run_step of components
        if self._print_timing:
            start = time.time()
        for model in self._exec_order:
            # Get correct position in buffer for update
            update_pointers = {}
            for var in self._comps[model]['updates']:
                buff = self.memory_manager.get_buffer(var)
                mind = self.memory_manager.variables[var]['models'].index(model)
                shift = self.memory_manager.variables[var]['cumlen'][mind]
                buffer_current_plus_one = buff.current + 1
                if buffer_current_plus_one >= buff.buffer_length:
                    buffer_current_plus_one = 0
                update_pointers[var] = int(buff.gpudata)+\
                                       (buffer_current_plus_one*buff.ld+\
                                        shift)*buff.dtype.itemsize
            self.components[model].run_step(update_pointers)
        if self._print_timing:
            cuda.Context.synchronize()
            self.timing['model_run'] += time.time()-start

        # Process output processors
        if self._print_timing:
            start = time.time()
        for p in self.output_processors: p.run_step()
        if self._print_timing:
            cuda.Context.synchronize()
            self.timing['output_processors'] += time.time()-start

        # Check for transforms

        # Update output ports
        if self._print_timing:
            start = time.time()
        self._extract_output()
        if self._print_timing:
            cuda.Context.synchronize()
            self.timing['extract_output'] += time.time()-start

        # Step through buffers
        self.memory_manager.step()

        self.step += 1

        # Instruct Control inteface to process any pending commands
        if self.control_inteface: self.control_inteface.process_commands()
        if self._print_timing:
            cuda.Context.synchronize()
            self.timing['total'] += time.time()-start_all

    def _read_LPU_input(self):
        """
        Extract membrane voltages/spike states from LPU's port map data arrays and
        store them in buffers.
        """
        for var in self.port_inds_gpot:
            # Get correct position in buffer for update
            buff = self.memory_manager.get_buffer(var)
            dest_mem = garray.GPUArray((1,buff.size),buff.dtype,
                                       gpudata=int(buff.gpudata)+\
                                       buff.current*buff.ld*\
                                       buff.dtype.itemsize)
            self.set_inds_both(self.pm['gpot'].data, dest_mem,
                               self.port_inds_gpot[var],self.var_inds_gpot[var])
        for var in self.port_inds_spk:
            # Get correct position in buffer for update
            buff = self.memory_manager.get_buffer(var)
            dest_mem = garray.GPUArray((1,buff.size),buff.dtype,
                                       gpudata=int(buff.gpudata)+\
                                       buff.current*buff.ld*\
                                       buff.dtype.itemsize)
            self.set_inds_both(self.pm['spike'].data, dest_mem, \
                          self.port_inds_spk[var],self.var_inds_spk[var])

    def _extract_output(self):
        """
        Extract membrane voltages/spike states from LPU's port map data arrays and
        store them in buffers.
        """

        for var in self._out_port_inds_gpot:
            # Get correct position in buffer for update
            buff = self.memory_manager.get_buffer(var)
            src_mem = garray.GPUArray((1,buff.size),buff.dtype,
                                       gpudata=int(buff.gpudata)+\
                                       buff.current*buff.ld*\
                                      buff.dtype.itemsize)
            self.set_inds_both(src_mem, self.pm['gpot'].data, \
                    self._out_var_inds_gpot[var], self._out_port_inds_gpot[var])
        for var in self._out_port_inds_spk:
            # Get correct position in buffer for update
            buff = self.memory_manager.get_buffer(var)
            src_mem = garray.GPUArray((1,buff.size),buff.dtype,
                                       gpudata=int(buff.gpudata)+\
                                       buff.current*buff.ld*\
                                       buff.dtype.itemsize)
            self.set_inds_both(src_mem, self.pm['spike'].data, \
                    self._out_var_inds_spk[var], self._out_port_inds_spk[var])

    def set_inds_both(self, src, dest, src_inds, dest_inds):
        """
        Set `dest[dest_inds[i]] = src[src_inds[i]] for i in range(len(src_inds))`
        """

        try:
            func = self.set_inds_both.cache[(src_inds.dtype, src.dtype)]
        except KeyError:
            inds_ctype = dtype_to_ctype(src_inds.dtype)
            data_ctype = dtype_to_ctype(src.dtype)
            v = ("{data_ctype} *dest, {inds_ctype} *dest_inds, " +\
                 "{inds_ctype} *src_inds, {data_ctype} *src").format(\
                        data_ctype=data_ctype,inds_ctype=inds_ctype)
            func = elementwise.ElementwiseKernel(v,\
                            "dest[dest_inds[i]] = src[src_inds[i]]")
            self.set_inds_both.cache[(src_inds.dtype, src.dtype)] = func
        func(dest, dest_inds, src_inds, src, range=slice(0, len(src_inds), 1) )

    set_inds_both.cache = {}

    def _instantiate_component(self, comp_name):
        try:
            cls = self._comps[comp_name]['cls']
        except:
            self.log_info("Error instantiating component of model '%s'" \
                          % comp_name)
            return None

        params_dict = self.memory_manager.parameters[comp_name]
        access_buffers = {var:self.memory_manager.get_buffer(var) \
                          for var in self._comps[comp_name]['accesses'] \
                          if var in self.memory_manager.variables}
        return cls(params_dict, access_buffers, self.dt,
                   LPU_id=self.LPU_id, debug=self._debug,
                   cuda_verbose=bool(self._compile_options))


    def _load_components(self, extra_comps=[]):
        """
        Load all available NDcomponents
        """
        child_classes = NDComponent.NDComponent.__subclasses__()
        comp_classes = child_classes[:]
        for cls in child_classes:
            comp_classes.extend(cls.__subclasses__())
        comp_classes.extend(extra_comps)
        self._comps = {cls.__name__:{'accesses': cls.accesses ,
                                     'updates':cls.updates,
                                     'cls':cls} \
                       for cls in comp_classes if not cls.__name__[:4]=='Base'}

    @property
    def time(self):
        return self.step*self.dt

class uid_generator(object):
    def __init__(self):
        self.input_count = 0
        self.auto_count = 0

    def generate_uid(self, input=False):
        if input:
            uid = 'input_' + str(self.input_count)
            self.input_count += 1
        else:
            uid = 'auto_' + str(self.auto_count)
            self.auto_count += 1
        return uid
