"""
Simulation input, output, and analysis.

See analysis.py for tools for working with groups of simulations.

Robert A McDougal, 2011
"""

import os
import sys
import analysis
import cPickle
import array
from string import ascii_letters, digits
from tokenizer import tokens_to_string
from warnings import warn
import zlib


_snnet_data_header = 'snnet data\x1a\x04'

inf = float('inf')

class SimData(object):
    def save(self, sim, output_folder, run_id):
        """Save a simulation as the current SimData type."""
        file_mode = self._initialize(sim, output_folder, run_id)
        _save(self, output_folder, run_id, file_mode)
        
        
    def _initialize(self, sim, output_folder, run_id):
        """Initialize a SimData object without saving it."""
        self._run_id = run_id
        self._output_folder = output_folder
        self._save_data = sim.save_data()
        self._cell_groups = sim._cell_groups
        self._time_points = sim._time_points
        self._compartments = []
        self.connections = sim._connections
        self._algebraic_history = None
        for n in sim._neurons:
            self._compartments.append(n._compartments)
        return self._extract_data(sim)
    
    
    def _extract_data(self):
        raise Exception('Filetype does not implement method: _extract_data')
    
    
    def expression(self, var, id, t, global_dynamics = False):
        """Return the algebraic expression for a given variable at a given time.
        
        Parameters:
            var -- the name of the variable
            id -- the cell (or global_dynamics) id
            t -- the time
            global_dynamics -- set True to get global dynamics id (instead of neuron)
        """
        if self._algebraic_history is None:
            raise Exception('Filetype does not implement method: algebraic_equation')
        
        # check for state variables
        if var + str(id) in self.state_variables():
            raise Exception('algebraic_equation for algebraic (not differential) variables only.')            
        
        t = self.time(t)
        
        values = None
        for h in self._algebraic_history:
            if h['t'] > t:
                break
            if global_dynamics:
                values = h['global']
            else:
                values = h['neurons']
        try:
            return tokens_to_string(values[id][var])
        except:
            raise Exception('No such algebraic expression.')
            
            
    
    
    def _get_spike_times(self):
        """Return a list of lists of spike times for each cell.
        
        Not supported by this file type. Throws an exception.
        """
        raise Exception('Filetype does not implement method: _get_spike_times')
    
    def _get_names(self):
        """Return a list of the cell types of the neurons.
        
        Not supported by this file type. Throws an exception.
        """
        raise Exception('Filetype does not implement method: _get_names')
    
    def _get_run_time(self):
        raise Exception('Filetype does not implement method: _get_run_time')
    
    def _get_times(self):
        """Return a list of the data point times.
        
        Not supported by this file type. Throws an exception.
        """
        raise Exception('Filetype does not implement method: _get_times')
    
    def get_values(self, var, neuron_id):
        raise Exception('Filetype does not implement method: get_values')
    
    def index(self, var, neuron_id):
        raise Exception('Filetype does not implement method: index')
    
    def __getitem__(self, item):
        """Allow indexing [] access."""
        if type(item) == str:
            return self.get_values(item, None)
        if len(item) != 2:
            raise Exception('Access via [var] or [var, neuron_id]')
        return self.get_values(item[0], item[1])


    def get_indices(self, cells = None):
        """Return the indices of the cells specified.
        
        Parameter:
            cells -- a list of indices or a string naming a cell_type or cell_group
        
        Note:
        
        If cells is None, returns all indices.
        """
        if cells is None:
            return range(len(self.types))
        if isinstance(cells, list):
            return cells
        if cells in self._cell_groups:
            return self._cell_groups[cells]
        return [i for i in range(len(self.types)) if self.types[i] == cells]
        
        
    def field_potential(self, cells = None):
        """Return the average membrane potential of a group of cells.
        
        Parameters:
        
            cells -- can be a named group, a filename, or a list of indices
        
        Note:
        
        If cells == None, returns the average membrane potential of all the cells.
        """
        indices = self.get_indices(cells)
        num = float(len(indices))
        if num == 0:
            raise Exception('No such cells: ' + repr(cells))
        f_p = self.potential(indices[0])
        for i in indices[1 :]:
            new_v = self.potential(i)
            f_p = [old + v for old, v in zip(f_p, new_v)]
        return [v / num for v in f_p]            
        
    def active(self, times, cells = None, cell_type = None, cell_group = None,
               indices = None, activity_window = 100):
        """Return the number of cell_type cells active at times.
        
        Parameters:
            times -- list of times
            cells -- the cells to check (indices, cell type, or cell group)
            activity_window -- how long ago a cell may have fired and still
                               be considered active
        
        Deprecated Parameters:
            cell_type -- the type of cells
            cell_group -- the name of the group of cells to check
            indices -- the indices of the cells to check

        Note:
        
        If a value is passed to cells, then cell_type, cell_group, and indices are
        ignored.
        
        Setting more than one of cell_type, indices, or cell_group is an error.
        If all (including cells) are None, then counts over all cells.
        """
        types = self.types
        result = []
        spikes = self.spike_times
        if cells:
            indices = self.get_indices(cells)
        else:
            if indices == None and cell_type is None and cell_group is None:
                indices = range(len(types))
            else:
                warn(DeprecationWarning('indices, cell_type, and cell_group are deprecated. Use the cells parameter to active.'))
                if indices is None:
                    if cell_group is None and cell_type is None:
                        indices = range(len(types))
                    elif cell_group is not None and cell_type is None:
                        indices = self._cell_groups[cell_group]
                    elif cell_group is not None and cell_type is not None:
                        raise Exception('Cannot define both cell_type and cell_group')
                    else:
                        indices = [i for i in range(len(types)) if types[i] == cell_type]
                else:
                    if cell_type is not None:
                        raise Exception('Cannot define both cell_type and indices')
                    if cell_group is not None:
                        raise Exception('Cannot define both cell_group and indices')
        
        # TODO:  make more efficient version
        for t in times:
            active_count = 0
            for i in indices:
                current_spikes = spikes[i]
                spikes_in_window = [s for s in current_spikes if (s >= t - activity_window and s <= t)]
                if len(spikes_in_window) > 0:
                    active_count += 1
            result.append(active_count)
        return result

    def potential(self, i):
        """Return the membrane potential time course for neuron i.
        
        Not supported by this file type. Throws an exception.
        """
        raise Exception('Filetype does not implement method: potential')

    def _get_field_potentials(self):
        """Return a dict of the field potentials.
        
        Not supported by this file type. Throws an exception.
        """
        raise Exception('Filetype does not implement method: _get_field_potentials')
    
    def time(self, t):
        """Return the time corresponding to t.
        
        Note:
        
        t may either be a numeric value or the name of a saved time point.
        """
        if type(t) in (int, long, float):
            return t
        if t in self._time_points:
            return self._time_points[t]
        raise Exception('Unknown time point: ' + t)

    
    def state_variables(self):
        """Return a list of accessible state variables."""
        return []
    
    # TODO: Modify to only show those functions supported by the type.    
    def __dir__(self):
        return (self.state_variables() +
            ['t', 'run_time', 'types', 'spike_times', 'field_potentials',
             'data', 'get_values', 'index', 'get_indices', 'active', 'expression',
             'potential', 'isi', 'cv', 'state_variables', 'time', 'compartments'])
    
    def __getattr__(self, var):
        """Allow access via dot notation."""
        neuron_id = None
        if var == 't':
            return self._get_times()
        elif var == 'run_time':
            return self._get_run_time()
        elif var == 'types':
            return self._get_names()
        elif var == 'spike_times':
            return self._get_spike_times()
        elif var == 'field_potentials':
            return self._get_field_potentials()
        elif var == 'data':
            return self._save_data
        elif var == 'compartments':
            return self._compartments
        try:
            var, neuron_id = _split_name(var)
        except:
            raise AttributeError(var)
        try:
            return self.get_values(var, neuron_id)
        except:
            raise AttributeError(var)
    

    def isi(self, cells=None, cell_type=None, time_range=(-inf, inf), combine_cells=False):
        """Return the Interspike intervals (ISIs) of cells.
        
        Deprecated Parameters:
            cell_type
        
        Note:
        
        Returns the ISIs of all cell types if cells = None.
        
        Results limited to those spikes occuring in the interval time_range.
        
        If combine_cells is False, then returns a list of lists of ISIs, one
        list for each neuron. If True, returns a single list of ISIs covering
        all neurons of the cell_type.
        """
        names = self.types
        spikes = self._get_spike_times()
        isis = []
        time_range = [self.time(time_range[0]), self.time(time_range[1])]
        t_min = min(time_range)
        t_max = max(time_range)
        if cells is None and cell_type is not None:
            cells = cell_type
            warn(DeprecationWarning('cell_type is deprecated; use cells instead.'))
        for i in self.get_indices(cells):
            isi_current = []
            current_spikes = spikes[i]
            for j in range(1, len(current_spikes)):
                if current_spikes[j - 1] >= t_min and current_spikes[j] <= t_max:
                    isi_current.append(current_spikes[j] - current_spikes[j - 1])
            isis.append(array.array('d', isi_current))
        if combine_cells == True:
            isi_all = []
            for isi in isis:
                isi_all += isi
            isis = isi_all
        return isis
    

    def cv(self, cells=None, cell_type=None, time_range=(-inf, inf), combine_cells=False):
        """Return the CVs of cells.
        
        Deprecated Parameters:
            cell_type
        
        Note:
        
        Returns the CVs of all cell types if cells = None.
        
        Results limited to those spikes occuring in the interval time_range.
        
        If combine_cells is False, then returns a list of CVs, one for each
        neuron. If true, returns the average cv of all neurons of cell_type.
        """
        isis = self.isi(cells, cell_type, time_range, False)
        cvs = [analysis.cv(i) for i in isis]
        if combine_cells:
            cvs = sum(cvs) / float(len(cvs))
        return cvs        
    


    
def _split_name(s):
    """If s ends in a number, splits the variable part from the id part."""
    if s[0] not in ascii_letters:
        raise Exception('Invalid variable name')
    if s[-1] not in digits:
        return s, None
    for j in range(len(s), 0, -1):
        if s[j - 1] not in digits:
            return s[0 : j], int(s[j : ])
    return s[0], int(s[1 : ])





class Csv(SimData):
    """Format for saving a simulation as comma separated variables (csv).
    
    Note that SN Net cannot load these files back in because they contain no
    identifying information about the variables.
    """
    def save(self, sim, output_folder, run_id):
        file = open(os.path.join(output_folder, str(run_id) + '.csv'), 'w')
        for i in range(0, len(sim._data_t)):
            file.write(str(sim._data_t[i]) + ', ' + ', '.join([str(y) for y in sim._data_y[i]]) + '\n')
        file.close()



class Full(SimData):
    """Simulation filetype.
    
    Stores all of the simulation data.
    
    Warning: This is potentially not compatible between different versions
    of SN Net.
    
    Attributes:
        t -- list of simulation time points
        run_time -- how long the simulation was ran for
        types -- list of types of each cell
        spike_times -- list of when cells spiked
        field_potentials -- dict of field potentials
    """
    def __init__(self):
        self._names = None
    
    def state_variables(self):
        return self._sim.state_variables()
    
    def _extract_data(self, sim):
        self._algebraic_history = sim._algebraic_history
        self._sim = sim
        
    def _get_run_time(self):
        """Return total simulation run time."""
        return self._sim.t
        
    def _get_times(self):
        """Return a list of the data point times."""
        return self._sim._data_t
    
    def index(self, var, neuron_index = None):
        """Returns index of variable var from neuron with index neuron_index.
        
        Note:
        
        If neuron_index is None, returns the index of the global variable var.
        """
        return self._sim.index(var, neuron_index)

    def potential(self, i):
        """Return the membrane potential time course for neuron i."""
        return self._sim.potential(i)

    def get_values(self, var, neuron_index = None):
        """Returns a list of the values of var over time.
        
        neuron_index = which neuron's values to return
        
        Note:
        
        If neuron_index is None, then returns the dynamics for the global
        variable var.
        """
        return self._sim.get_values(var, neuron_index)
        
    def _get_spike_times(self):
        """Return a list of lists of spike times for each cell."""
        return self._sim.spike_times
        
    def _get_names(self):
        """Return a list of the cell types of the neurons."""

        if self._names is None:
            self._names = []
            for neuron in self._sim._neurons:
                self._names.append(neuron._name)
        return self._names
    
    def _get_field_potentials(self):
        return self._sim.field_potentials




class Full2(SimData):
    """Simulation filetype.
    
    Stores all of the simulation data.
    
    Initially loads only a subset of the data (spike times, types,
    field potentials, t); automatically loads more if needed. This
    gives faster loading for common operations like make_rasters.
    
    Attributes:
        t -- list of simulation time points
        run_time -- how long the simulation was ran for
        types -- list of types of each cell
        spike_times -- list of when cells spiked
        field_potentials -- dict of field potentials
    """
    def __init__(self):
        self._names = None
    
    def state_variables(self):
        return self._sim.state_variables()
    
    def _extract_data(self, sim):
        self._data_t = sim._data_t
        self._spike_times = sim.spike_times
        self._field_potentials = sim.field_potentials
        self._names = []
        self._run_time = sim.t
        for neuron in sim._neurons:
            self._names.append(neuron._name)
        _save(self, self._output_folder, self._run_id)
        self._sim = sim
        self._algebraic_history = sim._algebraic_history
        return 'ab'

    def __getattr__(self, var):
        if var in ('_sim', '_algebraic_history'):
            try:
                data = load(self._filename, 2)
                self._sim = data._sim
                self._algebraic_history = data._algebraic_history
                if var == '_sim':
                    return data._sim
                else:
                    return self._algebraic_history
            except:
                raise AttributeError(var)
        else:
            return SimData.__getattr__(self, var)

    def _get_run_time(self):
        """Return total simulation run time."""
        return self._run_time
        
    def _get_times(self):
        """Return a list of the data point times."""
        return self._data_t
    
    def index(self, var, neuron_index = None):
        """Returns index of variable var from neuron with index neuron_index.
        
        Note:
        
        If neuron_index is None, returns the index of the global variable var.
        """
        return self._sim.index(var, neuron_index)

    def potential(self, i):
        """Return the membrane potential time course for neuron i."""
        return self._sim.potential(i)

    def get_values(self, var, neuron_index = None):
        """Returns a list of the values of var over time.
        
        neuron_index = which neuron's values to return
        
        Note:
        
        If neuron_index is None, then returns the dynamics for the global
        variable var.
        """
        return self._sim.get_values(var, neuron_index)
        
    def _get_spike_times(self):
        """Return a list of lists of spike times for each cell."""
        return self._spike_times
        
    def _get_names(self):
        """Return a list of the cell types of the neurons."""
        return self._names
    
    def _get_field_potentials(self):
        return self._field_potentials
        



class Spikes(SimData):
    """Simulation filetype.
    
    Only stores the time of each spike.

    Attributes:
        run_time -- how long the simulation was ran for
        types -- list of types of each cell
        spike_times -- list of when cells spiked
    """
    def _extract_data(self, sim):
        self._spike_times = sim.spike_times
        self._names = sim.types
        self._run_time = sim.t

    def _get_spike_times(self):
        """Return a list of lists of spike times for each cell."""
        return self._spike_times
        
    def _get_names(self):
        """Return a list of the cell types of the neurons."""
        return self._names
        
    def _get_run_time(self):
        """Return total simulation run time."""
        return self._run_time


class Enhanced(Spikes):
    """Simulation filetype.
    
    Same as Spikes, except also stores field potentials and times.

    Attributes:
        t -- list of simulation time points
        run_time -- how long the simulation was ran for
        types -- list of types of each cell
        spike_times -- list of when cells spiked
        field_potentials -- dict of field potentials
    """
    def _extract_data(self, sim):
        Spikes._extract_data(self, sim)
        self._t = sim._data_t
        self._field_potentials = sim.field_potentials

    def _get_times(self):
        """Return a list of the data point times."""
        return self._t
    
    def _get_field_potentials(self):
        return self._field_potentials


_length_long = len(array.array('L', [0]).tostring())

def load(filename, item=1):
    """Load the simulation stored in filename.
    
    Parameters:
        filename -- the file to load
        item -- which pickled item in the file to load
    
    If filename refers to a directory, then loads the file 0.sdat from that directory.
    """
    if os.path.isdir(filename):
        filename = os.path.join(filename, '0.sdat')
    f = open(filename, 'rb')
    header = f.read(len(_snnet_data_header))
    if header != _snnet_data_header:
        for i in xrange(item):
            f.seek(0)
            data = cPickle.load(f)
    else:
        version = f.read(4)
        version += '\x00' * (_length_long - 4)
        version = array.array('L', version)
        if sys.byteorder != 'little':
            version.byteswap()
        version = version[0]
        max_supported_version = 1
        if version > max_supported_version:
            raise Exception('File saved in version format %d. Max supported is %d.'
                            % (version, max_supported_version))

        # advance to the desired item
        for i in xrange(item - 1):
            len_data_len = ord(f.read(1))
            data_length = int(f.read(len_data_len))
            f.seek(f.tell() + data_length)

        # read the item    
        len_data_len = ord(f.read(1))
        data_length = int(f.read(len_data_len))
        data = cPickle.loads(zlib.decompress(f.read(data_length)))
    
    f.close()
    data._filename = filename
    return data
    
    
    
    

#
# helper functions
#


def _save(data, output_folder, run_id=None, mode=None):
    """Pickle the SimData data.
    
    Parameters:
        data -- the SimData
        output_folder -- the folder to store the data
        run_id -- the simulation number
        mode -- file open mode
    
    Note:
    
    If mode is None, will use 'wb'
    
    The data will be stored in ouput_folder/run_id.pkl
    
    If run_id is None, treats output_folder as the path.
    """
    # pickle with protocol 2, supported in python 2.3+
    if run_id is None:
        filename = output_folder
    else:
        filename = os.path.join(output_folder, str(run_id) + '.sdat')
    if mode is None:
        mode = 'wb'
    
    
    save_version = array.array('L', [1]).tostring()
    if sys.byteorder != 'little':
        save_version.byteswap()
    if len(save_version) > 4:
        save_version = save_version[0 : 4]
    
    f = open(filename, mode)
    if mode == 'wb':
        f.write(_snnet_data_header)
        f.write(save_version)
    compressed_data = zlib.compress(cPickle.dumps(data, 2))
    data_length = '%d' % len(compressed_data)
    f.write('%c%s' % (len(data_length), data_length))
    f.write(compressed_data)
    f.close()
