"""A collection of neurons.

Robert A McDougal, 2011
"""

from errors import CompareError
from parameterList import ParameterList
from tokenizer import parse_string

class NeuronGroup(object):
    """A collection of neurons."""
    
    def __init__(self, indices, simulation, name=None):
        """Initialize the NeuronGroup.
        
        Parameters:
            indices -- the indices of the neurons in simulation
            simulation -- the simulation these cells belong to
            name -- a name for this group of cells (used in analysis functions)
        
        Note:
        
        A neuron group contains a given neuron at most once. That is,
        repeats in indices are ignored.
        """
        self._list = list(set(indices))
        self._simulation = simulation
        if name is None:
            name = len(simulation._cell_groups)
        simulation._cell_groups[name] = self._list
        simulation._cell_group_objects[name] = self
        self._name = name

    
    def __setattr__(self, name, value):
        """Set parameters using dot notation."""
        if name in ('_list', '_simulation', '_name', '_param_lists'):
            object.__setattr__(self, name, value)
        else:
            self.set(name, value)

    def __getattr__(self, var):
        """Get parameter list."""
        if var in ('_list', '_simulation', '_name', '_param_lists'):
            return object.__getattr__(self, var)
        if var == 'presynaptics':
            pre_list = []
            for n in self._list:
                if n >= 0:
                    pre_list += [p['pre_id'] for p in self._simulation._connections[n]]
            return NeuronGroup(pre_list, self._simulation)
        if var == 'postsynaptics':
            post_list = []
            for n in self._list:
                for i, n2_pre in enumerate(self._simulation._connections):
                    found = False
                    if n in [p['pre_id'] for p in n2_pre]:
                        post_list.append(i)
            return NeuronGroup(post_list, self._simulation)
        values = []
        attr_exists = False
        for i in self._list:
            if i < 0:
                val = None
            else:
                val = self._simulation._neurons[i]._get_val(var)
                if val is not None:
                    attr_exists = True
            values.append(val)
        
        # if none of the neurons have the requested parameter raise AttributeError
        if not attr_exists:
            raise AttributeError(var)
        return ParameterList(values, self._simulation)

    def __dir__(self):
        """List the variables and methods."""
        state_vars = []
        for i in self._list:
            state_vars += self._simulation._neurons[i].diff_list
        return list(set(['set', 'name', 'subset', 'len', 'connect', 'presynaptics', 'postsynaptics',
                         'flux', 'add_compartment'] + state_vars))

    # comparisons ignore name
    
    def __eq__(first, second):
        """Check for NeuronGroup equality:  =="""
        if first._simulation != second._simulation:
            raise CompareError('Cannot compare groups from different simulations.')
        return set(first._list) == set(second._list)
        
    def __ne__(first, second):
        """Check for NeuronGroup inequality:  !="""
        if first._simulation != second._simulation:
            raise CompareError('Cannot compare groups from different simulations.')
        return set(first._list) != set(second._list)

    def __lt__(first, second):
        """Check for set containment:  <"""
        if first._simulation != second._simulation:
            raise CompareError('Cannot compare groups from different simulations.')
        return set(first._list) < set(second._list)

    def __le__(first, second):
        """Check for set containment or equality:  <="""
        if first._simulation != second._simulation:
            raise CompareError('Cannot compare groups from different simulations.')
        return set(first._list) <= set(second._list)


    def __gt__(first, second):
        """Check for set containment:  >"""
        if first._simulation != second._simulation:
            raise CompareError('Cannot compare groups from different simulations.')
        return set(first._list) > set(second._list)

    def __ge__(first, second):
        """Check for set containment or equality:  >= """
        if first._simulation != second._simulation:
            raise CompareError('Cannot compare groups from different simulations.')
        return set(first._list) >= set(second._list)


    # set combinations

    def __invert__(self):
        """Return a NeuronGroup of all the neurons not in this NeuronGroup.
        
        Example:  not_group = ~group
        """
        return NeuronGroup(list(set(range(len(self._simulation._neurons))) - set(self._list)), self._simulation)

    def __and__(first, second):
        """Return the intersection of two neuron groups."""
        if first._simulation != second._simulation:
            raise CompareError('Cannot intersect groups from different simulations.')
        return NeuronGroup(list(set(first._list) & set(second._list)), first._simulation)

    def __or__(first, second):
        """Union."""
        return first + second
        
    def __add__(first, second):
        """Return the union of two NeuronGroup objects."""
        if first._simulation is not second._simulation:
            raise CompareError('Cannot add lists from different simulations')
        return NeuronGroup(first._list + second._list, first._simulation)    
    
    def __sub__(first, second):
        """Return the set difference of two NeuronGroup objects."""
        if first._simulation is not second._simulation:
            raise CompareError('Cannot subtract lists from different simulations')
        new_list = filter(lambda v: v not in second._list, first._list)
        return NeuronGroup(new_list, first._simulation)

    def __xor__(first, second):
        """Return the symmetric difference of two neuron groups."""
        if first._simulation != second._simulation:
            raise CompareError('Cannot take symmetric difference of groups from different simulations.')
        return NeuronGroup(list(set(first._list) ^ set(second._list)), first._simulation)

    # misc
    
    def set(self, var, val=None, std_dev_frac=None, uniform=None):
        """Set parameters.
        
        Parameters:
            var -- the name of the parameter to set
            val -- the value to set it to
        
        Note:
        
        val may be a single value or a list of the same size as the NeuronGroup.
        
        Setting a parameter to None does not change its value.
        """
        val = self._simulation.value(val)
        std_dev_frac = self._simulation.value(std_dev_frac)
        uniform = self._simulation.value(uniform)
        if (val is None) and (std_dev_frac is None) and (uniform is None):
            raise Exception('Nothing to set to')
        if not(isinstance(val, list)):
            val = [val] * len(self._list)
        if len(val) != len(self._list):
            raise Exception('Set value must either be a constant or a list of the same length as the NeuronGroup.')
        for i, neuron in enumerate(self._list):
            if val[i] is not None:
                if neuron < 0:
                    raise Exception('Cannot set noise sources')
                self._simulation._neurons[neuron].set(var, val[i], std_dev_frac, uniform)

    
    def name(self, name):
        """Change the name of the current cell group."""
        del self._simulation._cell_groups[self._name]
        self._name = name
        self._simulation._cell_groups[name] = self._list
        return self
    
    def subset(self, n, name = None):
        """Return a random subset of the group of size n.
        
        Optionally accepts a name for the new group of cells.
        """
        n = self._simulation.value(n)
        return NeuronGroup(self._simulation.random.sample(self._list, n), self._simulation, name)
    
    def len(self):
        """Return the number of neurons in the group."""
        return len(self._list)
    
    def __len__(self):
        """Return the number of neurons in the group.
        
        Call via:
            len(group)
        """
        return self.len()
    
    def __repr__(self):
        """Return a string representing the contents."""
        return 'NeuronGroup([' + ', '.join([repr(i) for i in self._list]) + '], sim, name = %r)' % self._name    
    
    def add_compartment(self, snnet, num=1, potential_var=None, ignore_case=False, verify_defined=True, name=None, **params):
        """Add dynamics to the inside of each of the members of the group."""
        organelle_maker = self._simulation.define_neuron(snnet, potential_var, ignore_case = ignore_case)
        organelles = NeuronGroup([], self._simulation, name)
        for i in self._list:
            # create the organelles
            group = organelle_maker(num, verify_defined=verify_defined, container=i)
            organelles += group
            
            # set containment relationship
            for j in group._list:
                self._simulation._neurons[i].add_compartment(snnet, j)
            
            # adjust parameters as needed
            for v in params:
                if v not in param_overrides:
                    group.set(v, params[v])

    def flux(self, var, rate, external_var=None):
        """Define a flux between an compartments and their containers.
        
        Note:
        
        If external_var is omitted or None, will be taken equal to var.
        
        If var does not exist on the inside, it will be created.
        If external_var does not exist on the outside, it will be created.
        """
        if external_var is None:
            external_var = var
        for i in self._list:
            if i < 0:
                raise Exception('Cannot make fluxes using noise sources.')
            neuron_i = self._simulation._neurons[i]
            if (var in neuron_i.diff_list) and (var not in neuron_i.concentrations):
                raise Exception('State variable ' + var + ' must be declared a concentration to add a flux.')
            if var not in self._simulation._neurons[i].diff_list:
                organelle = self._simulation._neurons[i]
                organelle.diff_list.append(var)
                organelle.differential[var] = parse_string('0')
                if var not in organelle.concentrations:
                    organelle.concentrations.append(var)
            self._simulation._fluxes[i].append({'var': var, 'rate': rate, 'external_var': external_var})
            container_id = self._simulation._neurons[i]._container
            if container_id is not None:
                container = self._simulation._neurons[container_id]
                if external_var not in container.diff_list:
                    container.diff_list.append(external_var)
                    container.differential[external_var] = parse_string('0')
            else:
                for d in self._simulation._shared_dynamics:
                    if external_var in d.diff_list:
                        break
                else:
                    # _shared_dynamics[0] gets all newly added global variables
                    container = self._simulation._shared_dynamics[0]
                    container.diff_list.append(external_var)
                    container.differential[external_var] = parse_string('0')
                    container.concentrations.append(external_var)

    def connect(self, pre_synaptic, p=1, n=None, allow_self_connect=False, synapse=None):
        """Let the current cell group receive from pre_synaptic cells group.
        
        Parameters:
            p -- connection probability
            n -- connect to this fixed number of cells
            allow_self_connect -- may cells be connected to themselves
            synapse -- synaptic connection rule
        
        Note:
        
        It is an error to specify both p and n.
        
        p and n can be specified as numbers or as strings referring to
        the names of overriden parameters (e.g. in a sweep or via run's param).
        
        If synapse is a list of synapses, will choose one uniformly at random
        for each connection.
        """
        if isinstance(synapse, type):
            synapse = synapse() 
        if not isinstance(synapse, list):
            synapse = [synapse]
        p = self._simulation.value(p)
        n = self._simulation.value(n)
        if n is not None:
            n = int(n)
            if pre_synaptic.len() < n:
                raise Exception('Insufficient number of presynaptic cells.')
            if p != 1:
                raise Exception('Cannot specify both probability and a fixed number of cells.')
            for i in self._list:
                if pre_synaptic.len() == n and not(allow_self_connect) and i in pre_synaptic._list:
                    raise Exception('Insufficient number of presynaptic cells.')
                connect_list = self._simulation.random.sample(pre_synaptic._list, n)
                if not(allow_self_connect):
                    while i in connect_list:
                        connect_list = self._simulation.random.sample(pre_synaptic._list, n)
                for j in connect_list:
                    if len(synapse) > 1:
                        syn = self._simulation.random.choice(synapse)
                    else:
                        syn = synapse[0]
                    if syn is None:
                        rule = self._simulation._get_default_rule(i, j)
                    else:
                        rule = syn
                    for s in reduce(lambda x, y: x + y, self._simulation._connections):
                        if s['pre_id'] == j and s['rule'] == rule:
                            syn_id = s['syn_id']
                            break
                    else:
                        syn_id = self._simulation._synapse_count
                        self._simulation._synapse_count += 1
                    self._simulation._connections[i].append({'pre_id': j, 'rule': rule, 'syn_id': syn_id})

        else:
            for i in self._list:
                if i < 0:
                    raise Exception('Noise sources cannot receive input')
                for j in pre_synaptic._list:
                    if (i is not j) or allow_self_connect:
                        if p >= 1 or self._simulation.random.random() <= p:
                            if len(synapse) > 1:
                                syn = self._simulation.random.choice(synapse)
                            else:
                                syn = synapse[0]
                            if syn is None:
                                rule = self._simulation._get_default_rule(i, j)
                            else:
                                rule = syn
                            for s in reduce(lambda x, y: x + y, self._simulation._connections):
                                if s['pre_id'] == j and s['rule'] == rule:
                                    syn_id = s['syn_id']
                                    break
                            else:
                                syn_id = self._simulation._synapse_count
                                self._simulation._synapse_count += 1
                            self._simulation._connections[i].append({'pre_id': j, 'rule': rule, 'syn_id': syn_id})
    
    
    def __getitem__(self, item):
        """Use [] notation to get subgroups and compartments."""
        if isinstance(item, str):
            new_list = []
            for i in self._list:
                try:
                    new_list += self._simulation[i]._compartments[item]
                except:
                    pass            
        elif isinstance(item, tuple):
            new_list = []
            for i in item:
                indices = self._list[i]
                if type(indices) == list:
                    new_list += indices
                else:
                    new_list.append(self._list[i])
        elif isinstance(item, slice):
            new_list = self._list[item]
        elif isinstance(item, int) or isinstance(item, float):
            new_list = [self._list[int(item)]]
        else:
            raise Exception('Invalid subset argument: ' + repr(item))
        return NeuronGroup(new_list, self._simulation)
