"""
The simulation engine.

Robert A McDougal, 2011
"""

from dynamics import Dynamics
from neuronGroup import NeuronGroup
from string import ascii_letters
import tokenizer
import random
import analysis
from tokenizer import Variable, parse_string, tokens_to_string
from errors import SyntaxError

tokens = ('-', '+', '*', '/', '^', '=', "'", '(', ')')
whitespace = (' ', '\t')

def tokenize_snnet(s, ignore_case=False):
    """Load the snnet description s.
    
    Parameters:
        s -- the snnet source
        ignore_case -- should we ignore case
    
    Returns:
        algebraic, differential, diff_list, concentrations
            -- dicts of expressions and a list of the order of the state variables
            -- concentrations is a list of the variables that are concentrations

    Note:
    
    ignore_case simply drops capitalization, if any
    
    s may be either the name of a snnet file (tries this first) or a string
    giving a snnet description.
    """
    
    differential = {}
    algebraic = {}
    diff_list = []
    concentrations = []
    functions = {}
    fluxes = []
    is_file = True
    try:
        file = open(s)
    except:
        is_file = False

    if is_file:
        lines = []
        while True:
            line = file.readline()
            if not line:
                break
            lines.append(line)
        file.close()
    else:
        lines = s.split('\n')
        if len(lines) < 2:
            raise Exception('File not found: ' + s)
    
    
    for line in lines:
        # discard leading and trailing white space
        line = line.strip()
        
        if ignore_case:
            line = line.lower()
        
        # remove comments (if any)
        comment_pos = line.find('#')        
        if comment_pos >= 0:
            line = line[0 : comment_pos]
        
        if line[ : 14].lower() == 'concentration ' or line[ : 15].lower() == 'concentrations ':
            concentrations += [s.strip() for s in line[14 : ].split(',')]
        else:
            # split into equations, process
            equations = _split_into_equations(line)
            for equation in equations:
                _process_equation(equation, differential, algebraic, diff_list, functions, fluxes)
    space_list = []
    for c in concentrations:
        if ' ' in c:
            space_list.append(c)
        if c not in diff_list:
            differential[c] = parse_string('0')
            diff_list.append(c)
    if space_list:
        raise Exception('concentration state names cannot contain spaces: ' + ', '.join(space_list))
    
    return algebraic, differential, diff_list, functions, concentrations, fluxes


def _split_into_equations(s):
    """Splits s into a list of equations at commas.
    
    Note:
    
    Commas inspired parentheses are considered part of the equation, not
    equation separators.
    """
    equation_begin = 0
    paren_count = 0
    equations = []
    for i, c in enumerate(s):
        if c == '(':
            paren_count += 1
        elif c == ')':
            paren_count -= 1
            if paren_count < 0:
                raise SyntaxError(s, 'unbalanced parentheses')
        elif c == ',' and paren_count == 0:
            equation = s[equation_begin : i].strip()
            if equation:
                equations.append(equation)
            equation_begin = i + 1
    equation = s[equation_begin :].strip()
    if equation:
        equations.append(equation)
    return equations



def _process_equation(line, differential, algebraic, diff_list, functions, fluxes):
    """Processes a single equation, updating dictionaries and state variable list as appropriate."""
    if len(line) > 0:
        # all remaining lines should be of the form token[(...) | '] = token[s]
        i = line.find('=')
        if i < 1:
            # are we a flux?
            i = line.find(' flux ')
            if i < 1:
                raise SyntaxError(line, 'missing equal sign')
            flux_var = line[: i].strip()
            if ' ' in flux_var:
                raise SyntaxError(line, 'invalid flux variable: ' + flux_var)
            flux_rhs = line[i + 6 :].strip()
            fluxes.append({'var': flux_var, 'val': flux_rhs})
            return
        
        lhs = line[: i].strip()
        rhs_string = line[i + 1 : ].strip()
        if not rhs_string:
            raise SyntaxError(line, 'no value assigned to: ' + lhs)
        rhs = tokenizer.parse_string(line[i + 1 : ])
        
        # are we algebraic or differential or a function
        is_function = False
        if lhs[-1] == "'":
            is_algebraic = False
            base_name = lhs[0 : -1]
        elif lhs[-1] == ')':
            is_algebraic = True
            is_function = True
            if '(' not in lhs:
                raise SyntaxError(line, 'invalid variable name or function declaration')
            open_paren_location = lhs.index('(')
            base_name = lhs[: open_paren_location]
            var_tokens = tokenizer.tokenize(lhs[open_paren_location + 1 : -1])
            for i, t in enumerate(var_tokens):
                if (i % 2 == 0 and not isinstance(t, tokenizer.Variable)) or (i % 2 == 1 and t.content != ','):
                    raise SyntaxError(line, 'invalid lhs in function declaration')
            # get the names of the variables
            var_names = [t.content for t in var_tokens if isinstance(t, tokenizer.Variable)]
        else:
            is_algebraic = True
            base_name = lhs
        
        # verify valid identifier
        if base_name[0] not in ascii_letters:
            raise SyntaxError(line, 'invalid variable name: ' + base_name)
        
        for c in base_name:
            if c not in tokenizer.letters_and_numbers:
                raise SyntaxError(line, 'invalid variable name: ' + base_name)

        if base_name in algebraic and (not is_function) and is_algebraic:
            raise SyntaxError(line, 'variable already defined: ' + base_name)
        
        if is_function:
            if base_name in functions or base_name in tokenizer.functions:
                raise SyntaxError(line, 'function already defined: ' + base_name)
            functions[base_name] = {'var_names': var_names, 'expression': rhs, 'num_args': len(var_names)}
        elif is_algebraic:
            algebraic[base_name] = rhs
        else:
            differential[base_name] = rhs
            diff_list.append(base_name)
        



def _make_creator(sim, algebraic, differential, diff_list, functions, concentrations, fluxes, neuron_name, potential_var):
    return lambda n, name=None, verify_defined=True, container=None: sim._create_neurons(n, algebraic, differential, diff_list, functions, concentrations, fluxes, neuron_name, potential_var, name, verify_defined, container=container)


def _verify_initial(diff_list, algebraic, dynamics):
    undeclared = []
    for var in diff_list:
        if var not in algebraic:
            undeclared.append(var)
    if len(undeclared) > 0:
        raise Exception('Initial conditions not declared in ' + dynamics + ' for: '
                        + ', '.join(undeclared))


class Simulation(object):
    """An SN Net simulation."""
    
    def __init__(self, integrator, param_override=[], max_advance_time=500):
        self._neurons = []
        self.random = random.Random()
        self._param_override = param_override
        self._integrator = integrator
        self._noise = [None]
        self.t = 0
        self._leave_temp_files = False
        self._data_t = []
        self._data_y = []
        self._connections = []
        self._fluxes = []
        self._suffix = ''
        self._spike_times_cache = None
        self._field_potentials = None
        self._save_data = []
        self._cell_groups = {}
        self._time_points = {}
        self._cell_group_objects = {}
        self._total_volume = 1
        self._algebraic_history = []
        self._shared_dynamics = [Dynamics({}, {}, [], {}, self, None)]
        self._synapse_count = 0

        # replace this list for automatic file cleanup
        self._temp_files = []
        
        # the maximum amount to advance at one time
        # if this is large, compilers may use a lot of memory
        self.max_advance_time = max_advance_time
    
    def __dir__(self):
        return ['state_variables', 'save_data', 'integrator_method',
                'define_neuron', 'global_dynamics', 'index', 'get_values',
                'potential', 'set', 'save', 'run_to', 'undefined_variables',
                'advance', 'random', 't', '_shared_dynamics', '_neurons',
                '_param_override', '_integrator', '_noise', '_data_t',
                '_data_y', '_connections', '_suffix', '_field_potentials',
                '_save_data', '_cell_groups', '_temp_files', 'value', '_algebraic_history',
                'max_advance_time', 'group', 'time', 'name_time',
                'has_synaptic_delays', '_total_volume', 'collapse_probability']



    def has_synaptic_delays(self):
        """Check if any of the synapses are delayed."""
        for syns in self._connections:
            for syn in syns:
                if syn['rule'].has_delay():
                    return True
        return False
            

    def collapse_probability(self, s):
        try:
            if not isinstance(s, str) or '[' not in s:
                return float(s)
                    
            i = s.find('[')
            j = s.find(']')
            
            if j + 1 != len(s):
                raise SyntaxError(s, 'Invalid number')
            
            inside = s[i + 1 : j].strip()
            mean = s[ : i].strip()
            
            if len(mean) == 0:
                # uniform distribution
                parts = inside.split(':')
                if len(parts) != 2:
                    raise SyntaxError(s, 'Invalid number.')
                values = [float(v) for v in parts]
                return self.random.uniform(min(values), max(values))
            
            mean = float(mean)
            
            if inside[-1] == '%':
                std_dev = mean * float(inside[ : -1]) * .01
            else:
                std_dev = float(inside)
            
            return self.random.gauss(mean, std_dev)
        except:
            raise SyntaxError(s, 'Invalid number.')

    def time(self, t):
        """Return the time denoted by t.
        
        Note:
        
        t may be a numeric value, or a named 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 name_time(self, name, toffset=0):
        """Assign a name to a particular time.
        
        Parameters:

            name -- the name to assign
            toffset -- the offset from the current time
        """
        self._time_points[name] = self.t + toffset        
                
    def group(self, name):
        """Return a predefined named group."""
        if name in self._cell_group_objects:
            return self._cell_group_objects[name]
        raise Exception('No such cell group: ' + name)
        
    def state_variables(self):
        """Return a list of all known state variables."""
        var = []
        for dyn in self._shared_dynamics:
            var += dyn.diff_list
        for i, dyn in enumerate(self._neurons):
            var += ['%s%d' % (v, i) for v in dyn.diff_list]
        return var
    
    def __getattr__(self, var):
        """Allow access via dot notation."""
        if var == 'spike_times':
            return self._get_spike_times()
        elif var == 'field_potentials':
            return self._get_field_potentials()
        elif var == 'types':
            return self._get_types()
        raise AttributeError(var)
        
    def save_data(self, data=None):
        """Append data to list of things to save."""
        if data is not None:
            self._save_data.append(data)
        else:
            return self._save_data
    
    def integrator_method(self, integration_method=None):
        """Set or get the current integration method"""
        if integration_method is None:
            return self._integrator
        else:
            self._integrator = integration_method

    def define_neuron(self, neuron, potential_var=None, ignore_case=False):
        """Load a neuron for use in the simulation.
        
        Parameters:
            neuron -- the snnet containing the neural dynamics
            potential_var -- the name of the variable for membrane potential
            ignore_case -- should we drop capitalization?
        
        Note:
        
        neuron can be either the name of a file, or a string containing equations.
        
        potential_var is used for determining spike times and connecting
        synapses.
        
        If no potential_var is specified, then uses the first variable
        whose derivative is specified.
        """
        algebraic, differential, diff_list, functions, concentrations, fluxes = tokenize_snnet(neuron, ignore_case=ignore_case)
        if len(diff_list) == 0:
            raise Exception('snnet file has no differential equations.')
        if potential_var is None:
            potential_var = diff_list[0]
        elif potential_var not in diff_list:
            raise Exception('No such state variable: ' + potential_var)
        
        # make sure initial conditions have been specified for each differential equation
        _verify_initial(diff_list, algebraic, neuron)
        
        neuron_creator = _make_creator(self, algebraic, differential, diff_list, functions, concentrations, fluxes, neuron, potential_var)
        return neuron_creator
    
    def _get_types(self):
        """Return a list of the types of each neuron."""
        return [n._name for n in self._neurons]
    
    # TODO: Verify variable names in global dynamics are distinct.
    # TODO: modify to support volume being defined in terms of other things
    def global_dynamics(self, d, ignore_case=False):
        """Loads shared (global) dynamics.
        
        Parameters:
            d -- dynamics
            ignore_case -- should we drop capitalization in the snnet?
        
        Note:
        
        d can be either a filename or a string of equations.
        
        If multiple global dynamics are loaded, the variable names must not
        overlap.
        """
        algebraic, differential, diff_list, functions, concentrations, fluxes = tokenize_snnet(d, ignore_case=ignore_case)
        if fluxes:
            raise Exception('global dynamics cannot have outward fluxes')
        _verify_initial(diff_list, algebraic, d)
        new_dynamics = Dynamics(algebraic, differential, diff_list, functions, self, d, concentrations=concentrations)
        self._shared_dynamics.append(new_dynamics)
        self._apply_overrides(new_dynamics)
        if 'volume' in diff_list:
            raise Exception('change in total simulation volume is not supported')
        if 'volume' in algebraic:
            self._total_volume = float(tokens_to_string(algebraic['volume']))
        return new_dynamics
        
    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.
        """
        if neuron_index is None:
            has_var = False
            j = 0
            for dyn in self._shared_dynamics:
                if not(has_var):
                    if var in dyn.diff_list:
                        j += dyn.diff_list.index(var)
                        has_var = True
                        break
                    else:
                        j += len(dyn.diff_list)
            for n in self._neurons:
                j += len(n.diff_list)
        else:
            if var not in self._neurons[neuron_index].diff_list:
                raise Exception('Variable ' + var + ' not found')
            j = self._neurons[neuron_index].diff_list.index(var)
            for i in range(neuron_index):
                j += len(self._neurons[i].diff_list)
        return j
    
    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.
        """
        j = self.index(var, neuron_index)
        return [row[j] for row in self._data_y]
    
    def potential(self, neuron_index):
        """Return the membrane potential time course of neuron neuron_index."""
        return self.get_values(self._neurons[neuron_index]._potential_var, neuron_index)
    
    def value(self, v):
        """Return the value of v.
        
        If v is a number, returns that.
        If v is an overriden parameter, returns the value of the overriden parameter.
        
        If gna is overriden as gna = 2, then value('gna [10%]') returns '2 [10%]'
        """
        if isinstance(v, str):
            if v in self._param_override:
                return self._param_override[v]
            elif '[' in v:
                split = -1 - v[ : : -1].find('[')
                lhs = v[ : split].strip()
                rhs = v[split : ]
                if lhs in self._param_override:
                    return str(self._param_override[lhs]) + ' ' + rhs            
        return v
    
    def _get_field_potentials(self):
        """Return a dict of field potentials."""
        if self._field_potentials is None:
            field_potentials = {}
            type_counts = {}
            for j, n in enumerate(self._neurons):
                if n._name not in field_potentials:
                    field_potentials[n._name] = self.potential(j)
                    type_counts[n._name] = 1
                else:
                    type_counts[n._name] += 1
                    old_fp = field_potentials[n._name]
                    new_v = self.potential(j)
                    field_potentials[n._name] = [old_fp[i] + new_v[i] for i in range(len(new_v))]
            for t in type_counts:
                old_fp = field_potentials[t]
                num = float(type_counts[t])
                field_potentials[t] = [old_fp[i] / num for i in range(len(old_fp))]
            self._field_potentials = field_potentials
        return self._field_potentials
    
    def _get_spike_times(self):
        """Return a list of lists of spike times for each cell."""
        
        if self._spike_times_cache is None:
            self._spike_times_cache = []
            for i, neuron in enumerate(self._neurons):
                self._spike_times_cache.append(analysis.compute_spike_times(self._data_t, self.get_values(neuron._potential_var, i)))
        return self._spike_times_cache

    def _apply_overrides(self, dyn):
        """Apply simulation overrides to dynamics dyn.
        
        Used for doing parameter sweeps."""
        for var in self._param_override:
            dyn.set(var, self._param_override[var])
    
    def set(self, var, val=None, std_dev_frac=None, uniform=None):
        """Sets parameters for global dynamics.

        Has no effect on neuron parameters.
        """
        for dynamics in self._shared_dynamics:
            dynamics.set(var, val, std_dev_frac, uniform)

    def _create_neurons(self, num, algebraic, differential, diff_list, functions,
                        concentrations, fluxes, name, potential_var,
                        group_name=None, verify_defined=True, container=None):
        num = self.value(num)
        for i in range(num):
            neuron = Dynamics(algebraic, differential, diff_list, functions,
                              self, name, potential_var, index=len(self._neurons),
                              concentrations=concentrations, container=container)
            self._neurons.append(neuron)
            self._connections.append([])
            self._fluxes.append([])
            self._apply_overrides(neuron)
            self._apply_defaults(neuron)
            neuron._verify_defined = verify_defined
        i = len(self._neurons)
        group = NeuronGroup(range(i - num, i), self, group_name)
        for flux in fluxes:
            group.flux(flux['var'], flux['val'])
        return group

    def _apply_defaults(self, neuron):
        """Apply default values to neuron if none specified in the snnet.
        
        Defaults:
            capacitance = 1
        """
        if 'capacitance' not in neuron.algebraic:
            neuron.algebraic['capacitance'] = parse_string('1')
            
    def save(self, output_folder, run_id, save_type):
        """Save the simulation.
        
        Parameters:
            output_folder -- the folder to save the simulation to
            run_id -- the name of the file within the folder
            save_type -- a SimData subclass describing which data to save
        """
        # create an instance of the requested SimData and pass it the data
        save_type().save(self, output_folder, run_id)

    def run_to(self, t_stop):
        """Run the simulation until time t_stop.
        
        Parameters:
            t_stop -- stop time
        
        See also: advance
        """
        if t_stop < self.t:
            raise Exception('Already past stop time')
        self.advance(t_stop - self.t)
    
    # TODO:  add default connectivity rule support
    def _get_default_rule(self, i, j):
        raise Exception('Default connection rule support not yet implemented')
        
    def undefined_variables(self):
        """Returns a list of undefined variables (or empty list)."""
        shared_var = ['t']
        for dyn in self._shared_dynamics:
            shared_var += dyn.algebraic.keys()

        undefined = []
        for dyn in self._neurons + self._shared_dynamics:
            if not dyn._verify_defined:
                continue
            local_var = dyn.algebraic.keys()
            local_var.append('effective_volume')
            if dyn._container is not None:
                container_vars = self._neurons[dyn._container].algebraic.keys()
            else:
                container_vars = shared_var
            for v in container_vars:
                local_var.append('external_' + v)
            for eqn_type in (dyn.differential, dyn.algebraic):
                for v in eqn_type:
                    for token in eqn_type[v]:
                        if isinstance(token, Variable) and token.content not in local_var and token.content not in shared_var:
                            var_name = dyn._name + ':' + token.content
                            if var_name not in undefined:
                                undefined.append(var_name)
        return undefined
                        
    
    def advance(self, run_time):
        """Run the simulation for the period of time run_time.
        
        Parameters:
            run_time -- how long to run simulation for
        
        See also: run_to
        """
        # old precomputed data is no longer valid
        self._spike_times_cache = None
        self._field_potentials = None

        run_time = self.value(run_time)
        remaining_run_time = run_time

        if run_time < 0:
            raise Exception('advance amount must be nonnegative.')
        
        # check for undefined variables before running
        missing_vars = self.undefined_variables()
        
        if missing_vars:
            raise Exception('Undefined variables: ' + ', '.join(missing_vars))
        
        # save the algebraic equations (and parameters) at this point
        neuron_algebras = [0] * len(self._neurons)
        neuron_algebras = [dict(dyn.algebraic) for dyn in self._neurons]
        global_algebras = [dict(dyn.algebraic) for dyn in self._shared_dynamics]
        self._algebraic_history.append({'t': self.t, 'neurons': neuron_algebras, 'global': global_algebras})
        
        while remaining_run_time > 0:
            run_time = min(remaining_run_time, self.max_advance_time)
        
            t_new, y_new = self._integrator.integrate(self, run_time, temp_files=self._temp_files,
                                                      make_local_copy=self._leave_temp_files)
            
            # discard first row if duplicate of old data
            if self._data_t:
                if self._data_t[-1] == t_new[0]:
                    t_new = t_new[1 :]
                    y_new = y_new[1 :]
            
            # adjust simulation time
            self.t += run_time
            
            # save the state variables
            final_state = y_new[-1]
            i = 0
            for neuron in self._neurons:
                for j, var in enumerate(neuron.diff_list):
                    _adjust_for_concentrations(neuron, y_new, var, i, j)
                    neuron.set(var, final_state[i + j])
                i += len(neuron.diff_list)
            for dyn in self._shared_dynamics:
                for j, var in enumerate(dyn.diff_list):
                    _adjust_for_concentrations(dyn, y_new, var, i, j, is_shared=True)
                    dyn.set(var, final_state[i + j])
                i += len(dyn.diff_list)

            # store the new data
            self._data_t += t_new
            self._data_y += y_new
            
            remaining_run_time -= run_time


def _total_volume(dynamics, y_new):
    """Returns the timecourse of the volume of the given dynamics.

    This does NOT subtract the volume of subcompartments.
    """
    if dynamics.index is None:
        return [dynamics.simulation._total_volume] * len(y_new)
    elif 'volume' not in dynamics.diff_list:
        return [float(tokens_to_string(dynamics.algebraic['volume']))] * len(y_new)
    else:
        v_index = i + dynamics.diff_list.index('volume')
        return [row[v_index] for row in y_new]  


def _adjust_for_concentrations(neuron, y_new, var, i, j, is_shared=False):
    """Helper function for advance.
    
    Turns masses back into concentrations, if needed.
    """
    if var not in neuron.concentrations:
        return
    
    effective_volume = _total_volume(neuron, y_new)

    sim = neuron.simulation
    if is_shared:
        subtract_volumes = [n.index for n in neuron.simulation._neurons if n._container is None]
    else:           
        subtract_volumes = reduce(lambda x, y: x + y, neuron._compartments.values(), [])
    
    for k in subtract_volumes:
        local_volume = _total_volume(sim._neurons[k], y_new)
        effective_volume = [ev - lv for ev, lv in zip(effective_volume, local_volume)]
    
    for row, volume in zip(y_new, effective_volume):
        row[i + j] /= volume

