
from collections import OrderedDict
from datetime import datetime

import h5py
import numpy as np

import pycuda.driver as cuda
import pycuda.gpuarray as garray
from neurokernel.LPU.InputProcessors.BaseInputProcessor import BaseInputProcessor


class ArrayInputProcessor(BaseInputProcessor):
    """
    Parameters
    ----------
    inputs: dict
            A dictionary with the following structure:
            {var1: {'uids': ['comp1', 'comp2' ... ], 'data': np.ndarray},
             var2: {'uids': ['comp3', 'comp4' ... ], 'data': np.ndarray},
             ...}
            where in each key/value pair, the key specifies the varialbe to be
            inject input to, value['uids'] is a list that specifies the uids of
            components to be injected to, and value['data'] is an np.ndarray
            of shape (T, N) that specifies the input to each of the N components
            at every time step for a total of T steps, and the ordering the data in
            each column should be the same as the order in value['uids'].
            For 'spike_state', in addition to provide 0 (no spike) or 1 (spike)
            for each time step, spike time and the index of neuron that spiked
            can be provided by value['data']['time'] and value['data']['index'],
            respectively. See examples for more details.

            The input can also be programmatically loaded by other methods.
    mode:     int
              mode in BaseInputProcessor
              0: default, when the file does not cover the entire simulation,
                 inputs default to 0.
              1: when the file does not cover the entire simulation,
                 inputs defaults to the last state.
    cache_length: int
                  length (number of steps) of cache on GPU memory to preload
                  inputs.

    Examples
    --------
    Input h5d must have the following format:
    1. For 'spike_state', there are two options. First,
        {'spike_state':
            {'uids': numpy array of dtype 'S',
                             contains the uids of the nodes recorded,
                             the order of this array determines the index below,
            'data': {
                'time': numpy array of np.float32/np.double,
                      contains a **monotonically** increasing time stamps
                      of spikes from all nodes,
                'index': numpy array of np.int32
                       contains the index of every entry of spike time
                       in 'spike_state/data/time'}.
            }
        }
        loading a output file generated by FileOutputProcessor into a Dictionary
        of the above structure complies with this requirement.

       A second option is to have it as the same format as in 2 below, but have
       only float/double value of 0 (no spike at the time step) or 1 (spike)

    2. For other variables, e.g., 'I':
        {'I':
            {'uids': numpy array of dtype 'S',
                  contains the uids of the nodes recorded,
                  the order of this array determines which column each node
                  is receives as input in the 'I/data' array.,
             'data': numpy array of dtype np.float32/np.double,
                   contains the value of inputs injected to the nodes.
            }
        }
    """
    def __init__(self, inputs = {}, mode = 0, cache_length = 1000,
                 input_file = None, input_interval = 1,
                 sensory_file = None, sensory_interval = 1):
        super(ArrayInputProcessor, self).__init__([], mode, memory_mode = 'gpu',
                                                  input_file = input_file,
                                                  input_interval = input_interval,
                                                  sensory_file = sensory_file,
                                                  sensory_interval = sensory_interval)
        self.cache_length = cache_length
        self.data = {}
        self.add_inputs(inputs)

    def add_inputs(self, inputs):
        for var, g in inputs.items():
            self.add_input(var, g['uids'], g['data'])

    def add_input(self, var, uids, data):
        """
        Add inputs to the input processor

        Parameters
        ----------
        var: str
             the variable to which input will be injected to the components
        uids: list
              the uids of components to inject to
        data: ndarray or dict
              If `variable` is 'spike_state', then data can be in spike event
              format, as a dictionary {'time': np.ndarray, 'index', np.ndarray},\
              specifying the time and index of spikes.
              For all variables, data can be in stepwise input format,
              as a ndarray of shape (T, N)
              that specifies the input to each of the N components
              at every time step for a total of T steps.
              It is expected that if 'spike_state' input is added more than once,
              they must all be of one form, either in spike event format,
              or in stepwise input format.
        """
        if var in self.variables:
            if var == 'spike_state':
                if isinstance(data, dict):
                    assert issubclass(data['index'].dtype.type, np.integer), \
                           'spike index array must be integer dtype'
                    self.data[var]['time'] = np.concatenate(
                            self.data[var]['time'], data['time'])
                    self.data[var]['index'] = np.concatenate(
                            self.data[var]['index'],
                            data['index']+len(self.variables[var]['uids']))
                else:
                    self.data[var] = np.hstack((self.data[var], data))
            else:
                self.data[var] = np.hstack((self.data[var], data))
        else:
            self.data[var] = data
            if var == 'spike_state':
                if isinstance(data, dict):
                    self.spike_state_format = 'event'
                else:
                    self.spike_state_format = 'stepwise'
        self.add_variables([(var, uids)])

    def pre_run(self):
        self.dsets = {}
        self.cache = {}
        self.counts = {}
        self.end_of_var_in_array = {}
        self.block_total = {}
        self.end_of_var = {}
        self.last_read_index = {}

        if 'spike_state' in self.variables:
            if self.spike_state_format == 'event':
                sort_order = np.argsort(self.data['spike_state']['time'])
                self.data['spike_state']['time'] = self.data['spike_state']['time'][sort_order]
                self.data['spike_state']['index'] = self.data['spike_state']['index'][sort_order]

        for var in self.variables:
            data = self.data[var]
            uids = self.variables[var]['uids']
            self.block_total[var] = 0
            if var == 'spike_state':
                if self.spike_state_format == 'event':
                    self.dsets[var] = {'index': data['index'],
                                       'time': data['time']}
                else:
                    self.dsets[var] = data
                self.cache[var] = garray.empty((self.cache_length, len(self.variables[var]['uids'])),
                                               self.variables[var]['input'].dtype)
            else:
                self.dsets[var] = data
                self.cache[var] = garray.empty((self.cache_length, len(self.variables[var]['uids'])),
                                               self.variables[var]['input'].dtype)
            self.last_read_index[var] = 0
            self.counts[var] = 0
            self.end_of_var[var] = False
            self.end_of_var_in_array[var] = False
        # self.end_of_input = False

    def update_input(self):
        for var, dset in self.dsets.items():
            if self.counts[var] == self.block_total[var]:
                if not self.end_of_var_in_array[var]:
                    self.read_to_cache(var)

            if not self.end_of_var[var]:
                cuda.memcpy_dtod(self.variables[var]['input'].gpudata,
                                 int(self.cache[var].gpudata) + \
                                 self.cache[var].shape[1]*self.counts[var]*self.cache[var].dtype.itemsize,
                                 self.variables[var]['input'].nbytes)
                self.counts[var] += 1
            else:
                if not self.end_of_input:
                    self.variables[var]['input'].fill(0)

            if self.counts[var] == self.block_total[var]:
                if self.end_of_var_in_array[var]:
                    self.end_of_var[var] = True

    def read_to_cache(self, var):
        if self.end_of_var_in_array[var]:
            # self.cache[var].fill(0)
            self.counts[var] = 0
            self.block_total[var] = self.cache_length
            return

        if var == 'spike_state':
            if self.spike_state_format == 'event':
                self.read_spike_times(var)
            else:
                self.read_block_input(var)
        else:
            self.read_block_input(var)
        self.counts[var] = 0

    def read_block_input(self, var):
        tmp = self.dsets[var][self.last_read_index[var]:self.last_read_index[var]+self.cache_length,:]
        self.last_read_index[var] += tmp.shape[0]
        if self.last_read_index[var] == self.dsets[var].shape[0]:
            self.end_of_var_in_array[var] = True
        self.block_total[var] = tmp.shape[0]
        if tmp.shape[0] < self.cache_length:
            tmp1 = np.zeros(self.cache[var].shape, self.cache[var].dtype)
            tmp1[:tmp.shape[0]] = tmp
            tmp = tmp1
        self.cache[var].set(tmp)

    def read_spike_times(self, var):
        current_time = self._LPU_obj.time
        next_time = current_time + self.cache_length*self.sim_dt
        spike_times = []
        spike_index = []
        i = 0
        while True:
            tmp = self.dsets[var]['time'][self.last_read_index[var]:self.last_read_index[var]+10000]
            if tmp[0] + self.sim_dt/100 >= next_time:
                break
            last_spike = tmp.shape[0] - np.argmax(tmp[::-1] + self.sim_dt/100 < next_time)
            index = self.dsets[var]['index'][self.last_read_index[var]:self.last_read_index[var]+last_spike]
            spike_times.append(tmp[:last_spike])
            spike_index.append(index[:last_spike])
            self.last_read_index[var] += last_spike
            if self.last_read_index[var] == self.dsets[var]['time'].shape[0]:
                self.end_of_var_in_array[var] = True
                break
            if last_spike < tmp.shape[0]:
                break
        if len(spike_times):
            spike_times = np.concatenate(spike_times) - current_time
            spike_time_ids = np.floor(spike_times/self.sim_dt+np.finfo(spike_times.dtype).eps*(1/self.sim_dt)).astype(np.int32)
            spike_index = np.concatenate(spike_index)
            tmp = np.zeros(self.cache[var].shape, self.cache[var].dtype)
            self.block_total[var] = self.cache[var].shape[0]
            np.add.at(tmp, (spike_time_ids, spike_index), 1)
            self.cache[var].set(tmp)
        else:
            self.cache[var].fill(0)
            self.block_total[var] = self.cache[var].shape[0]

    @property
    def end_of_input(self):
        return all(self.end_of_var.values())

    def is_input_available(self):
        return not self.end_of_input

    def get_input(self, var = None, uids = None, dt = 1, start_time = 0):
        """
        retrieve inputs by variable name `var` and uids

        Parameters
        ----------
        var: str
             Name of the variable to retrieve.
             If None, all variables associated with a uid will be retrieved.
        uids: str or list of str
              uids of the component to retrieve
              If not specified (None), all uids associated with the `var`
              will be retrieved.

        Returns
        -------
        input: dict or OrderedDict
                If uids is a list or tuple, returns a OrderedDict, Otherwise a dict,
                with inputs keyed by uid, and values are either the data/spike_times
                of the component, or a dict keyed by variable name and data/spike_times
                in the value.
        """
        if var is None and uids is None:
            return self.data
        elif var is not None and uids is None:
            return self._get_input_by_var(var, dt = dt, start_time = start_time)
        elif var is None and uids is not None:
            return self._get_input_by_uids(uids, dt = dt, start_time = start_time)
        else:
            return self._get_input_by_var_and_uids(var, uids, dt = dt, start_time = start_time)

    def _get_input_by_var(self, var, dt = 1, start_time = 0):
        uids = self.variables[var]['uids']
        data = self.data[var]
        if var == 'spike_state' and self.spike_state_format == 'event':
            input = {uid: {'data': data['time'][data['index']==i]} for i, uid in enumerate(uids)}
        else:
            t = np.arange(0, data.shape[0])*dt + start_time
            input = {uid: {'time': t, 'data': data[:,i].copy()} for i, uid in enumerate(uids)}
        return input

    def _get_input_by_uid(self, uid, dt = 1, start_time = 0):
        input = {}
        for var in self.data:
            try:
                index = self.variables[var]['uids'].index(uid)
            except ValueError:
                pass
            else:
                if var == 'spike_state' and self.spike_state_format == 'event':
                    input[var] = self.data[var]['time'][self.data[var]['index'] == index]
                else:
                    t = np.arange(0, self.data[var].shape[0])*dt + start_time
                    input[var] = {'time': t, 'data': self.data[var][:,index].copy()}
        return input

    def _get_input_by_uids(self, uids, dt = 1, start_time = 0):
        if isinstance(uids, str):
            return {uids: self._get_input_by_uid(uids, dt = dt, start_time = start_time)}
        elif isinstance(uids, (list, tuple)):
            return OrderedDict([(uid, self._get_input_by_uid(uid, dt = dt, start_time = start_time)) for uid in uids])
        elif isinstance(uids, set):
            return {uid: self._get_input_by_uid(uid, dt = dt, start_time = start_time) for uid in uids}

    def _get_input_by_var_and_uid(self, var, uid, dt = 1, start_time = 0):
        input = None
        try:
            index = self.variables[var]['uids'].index(uid)
        except ValueError:
            pass
        else:
            if var == 'spike_state' and self.spike_state_format == 'event':
                input = {'data': self.data[var]['time'][self.data[var]['index'] == index]}
            else:
                t = np.arange(0, self.data[var].shape[0])*dt + start_time
                input = {'time': t, 'data': self.data[var][:,index].copy()}
        return input

    def _get_input_by_var_and_uids(self, var, uids):
        if isinstance(uids, str):
            return {uids: self._get_input_by_var_and_uid(var, uids)}
        elif isinstance(uids, (list, tuple)):
            return OrderedDict([(uid, self._get_input_by_var_and_uid(var, uid)) for uid in uids])
        elif isinstance(uids, set):
            return {uid: self._get_input_by_var_and_uid(var, uid) for uid in uids}

    def to_file(self, filename):
        """
        Store inputs to file with name `filename`,
        compatible to be used by `FileInputProcessor`

        parameters
        ----------
        filename: str
                  name of the file to store the results
        """
        with h5py.File(filename, 'w') as h5file:
            h5file.create_dataset('metadata', (), 'i')
            h5file['metadata'].attrs['start_time'] = 0.0
            h5file['metadata'].attrs['sample_interval'] = self.input_interval
            #h5file['metadata'].attrs['dt'] = self.dt
            h5file['metadata'].attrs[
                'DateCreated'] = datetime.now().isoformat()
            for var, d in self.variables.items():
                if var == 'spike_state' and self.spike_state_format == 'event':
                    h5file.create_dataset(var + '/data/index',
                                          dtype = np.int32,
                                          maxshape = (None,),
                                          data = self.data[var]['index'])
                    h5file.create_dataset(var + '/data/time',
                                          dtype = np.double,
                                          maxshape = (None,),
                                          data = self.data[var]['time'])
                else:
                    h5file.create_dataset(var + '/data',
                                          dtype = self.data[var]['data'].dtype,
                                          maxshape=(None, len(d['uids'])),
                                          data = self.data[var]['data'])
                h5file.create_dataset(var + '/uids',
                                      data = np.array(d['uids'], dtype = 'S'))

if __name__ == '__main__':
    import networkx as nx
    from neurokernel.LPU.LPU import LPU
    from neurokernel.LPU.OutputProcessors.OutputRecorder import OutputRecorder

    dur = 0.2
    dt = 1e-4
    steps = int(dur/dt)
    G = nx.MultiDiGraph()
    uid = 'synapse1'
    scale = 0.1
    ar = 0.110
    ad = 0.190
    gmax = 0.01
    G.add_node(uid, name = 'alpha_synapse',
               model = 'AlphaSynapse2',
               ar = ar,
               ad = ad,
               scale = scale,
               gmax = gmax)

    spike_times = np.array([0.038, 0.048, 0.078], np.float64)
    input_processor = ArrayInputProcessor(
        {'spike_state':
            {'uids': [uid],
             'data': {'time': spike_times,
                      'index': np.array([0]*spike_times.shape[0], np.int32)}
            }
        })
    output_processor = OutputRecorder([('g', [uid])], dur, dt)
    lpu = LPU(dt, 'obj', {'graph': G, 'kwargs': {'class_key': 'model'}},
              device = 0,
              id = 'test', input_processors = [input_processor],
              output_processors = [output_processor],
              debug = True,
              manager = False,
              print_timing = False,
              extra_comps = [ArrayInputProcessor])
    lpu.run(steps = steps)

    s = np.zeros(steps, np.double)
    t = np.arange(0, dur, dt)
    for tk in spike_times:
        s += (scale/(ar-ad) * (np.exp(-ad*(t-tk)*1000) - np.exp(-ar*(t-tk)*1000) ))*(t>tk)

    if np.abs(output_processor.output['g']['data'].reshape(-1)-s).max() < 1e-12:
        print('Test Passed')
    else:
        print('Test Failed')
