"""
Mathematical tools for working with SN Net simulations.

This file contains functions of general use and functions that work on groups
of simulations.

Robert A McDougal, 2011
"""

from math import sqrt
import os
import sim_data
import Queue
from parallelManagement import parallel_run_over_queue

try:
    from numpy.matlib import mean, std
except ImportError:
    def mean(x):
        """Compute the mean."""
        return float(sum(x)) / len(x)
    def std(x, ddof=0):
        """Compute the standard deviation of x."""
        mu = mean(x)
        return sqrt(float(sum([(y - mu) ** 2 for y in x])) / (len(x) - ddof))



inf = float('inf')

def compute_spike_times(t, v, threshold=0):
    """Return a list of the times t when v crosses threshold"""
    spikes = []
    if len(v) != len(t):
        raise Exception('compute_spike_times(t, v) needs t and v to have the same length')
    for i in range(1, len(v)):
        if v[i] > threshold and v[i - 1] <= threshold:
            spikes.append(t[i])
    return spikes

def discard_nans(*items):
    """Return a copy of the list(s) items without any NaNs."""
    result = [[i for i in item if (i == i)] for item in items]
    if len(items) == 1:
        return result[0]
    else:
        return result

def _run_f_post(result, i, v):
    """Helper function for run_f_on_dir."""
    result[i] = v

def _launch_f_on_pkl(param):
    """Helper function for run_f_on_dir."""
    f, i, pkl, dir_name = param
    num_f_args = f.func_code.co_argcount
    filename = os.path.join(dir_name, pkl)
    data = sim_data.load(filename)
    if num_f_args == 1:
        return {'i': i, 'v': f(data)}
    else:
        return {'i': i, 'v': f(data, filename)}



def run_f_on_pkl_in_dir(f, dir_name, enter_subdirs=False, num_parallel=1, temp_dir='.'):
    """Runs a function on the unpickled data from each pkl file in dir_name.
    
    Parameters:
        f -- the function to run.
        dir_name -- the name of the directory
        enter_subdirs -- should we apply f on files in subdirectories
        num_parallel -- maximum number of evaluations to run simultaneously
        temp_dir -- a string or function evaluating to a string representing
                    a temporary directory

    Returns:
        A list of the return values. (Or lists of lists, etc... if entering subdirectories.)
    
    Note:
    
    f may take 1 or 2 arguments. If f takes only 1 argument, passes the unpickled data;
    if f takes 2, passes unpickled data as first argument, filename as second.
    
    Parallelization is with processes on Linux and OSX; otherwise, threads.
    
    If num_parallel > 1, then f must return a picklable object. (This is not
    usually a problem.)
    
    THIS FUNCTION IS DEPRECATED. USE run_f_on_dir INSTEAD.
    """
    warn(DeprecationWarning('run_f_on_pkl_in_dir is deprecated. Use run_f_on_dir instead.'))
    run_f_on_dir(f, dir_name, enter_subdirs, num_parallel, temp_dir)


def run_f_on_dir(f, dir_name, enter_subdirs=False, num_parallel=1, temp_dir='.'):
    """Runs a function on the simulation data from each sdat (or pkl) file in dir_name.
    
    Parameters:
        f -- the function to run.
        dir_name -- the name of the directory
        enter_subdirs -- should we apply f on files in subdirectories
        num_parallel -- maximum number of evaluations to run simultaneously
        temp_dir -- a string or function evaluating to a string representing
                    a temporary directory

    Returns:
        A list of the return values. (Or lists of lists, etc... if entering subdirectories.)
    
    Note:
    
    f may take 1 or 2 arguments. If f takes only 1 argument, passes the unpickled data;
    if f takes 2, passes unpickled data as first argument, filename as second.
    
    Parallelization is with processes on Linux and OSX; otherwise, threads.
    
    If num_parallel > 1, then f must return a picklable object. (This is not
    usually a problem.)
    """
    pkl_queue = Queue.Queue()
    post = []
    result = _build_pkl_queue(f, dir_name, enter_subdirs, pkl_queue, post)
    parallel_run_over_queue(_launch_f_on_pkl, pkl_queue,
                            num_parallel=num_parallel,
                            post_processing=post, temp_dir=temp_dir)
    return result




def _make_post(result):
    """Helper function for _build_pkl_queue."""
    return lambda r: _run_f_post(result, r['i'], r['v'])


def _build_pkl_queue(f, dir_name, enter_subdirs, pkl_queue, post_list):
    """Helper function for run_f_on_dir.
    
    Recursively builds the pkl_queue, but does not run anything.
    """
    files = os.listdir(dir_name)
    pkl_files = [file for file in files if (file[-4 : ] == '.pkl' or file[-5 : ] == '.sdat')]
    result = [0] * len(pkl_files)
    for i, pkl in enumerate(pkl_files):
        pkl_queue.put([f, i, pkl, dir_name])
        post_list.append(_make_post(result))

                
    if enter_subdirs:
        sub_dirs = [file for file in files if os.path.isdir(os.path.join(dir_name, file))]
        for sub_dir_name in sub_dirs:
            result.append(_build_pkl_queue(f, os.path.join(dir_name, sub_dir_name), True, pkl_queue, post_list))
    return result


# TODO: verify that each .pkl is a SimData
def num_sims(dir_name):
    """Return the number of simulations in dir_name."""
    return len([file for file in os.listdir(dir_name) if (file[-4 : ] == '.pkl' or file[-5 :] == '.sdat')])


def collect_isi(dir_name, cells=None, cell_type=None, time_range=(-inf, inf),
                combine_cells=True, combine_sims=True, num_parallel=1):
    """Collect the ISIs from all the simulations in dir_name.
    
    Deprecated Parameters:
        cell_type -- the types of cells to collect from
    """
    isi = run_f_on_dir(lambda d: d.isi(cells, cell_type, time_range, combine_cells),
                              dir_name, num_parallel=num_parallel)
    if not combine_sims:
        return isi
    
    return reduce(lambda x, y: x + y, isi)


def collect_cv(dir_name, cells=None, cell_type=None, time_range=(-inf, inf),
               combine_sims=True, num_parallel=1):
    """Collect the CVs from all the simulations in dir_name.
    
    Deprecated Parameters:
        cell_type -- the type of cells to collect from
    """
    cvs = run_f_on_dir(lambda d: d.cv(cells, cell_type, time_range, False),
                              dir_name, num_parallel=num_parallel)
    if not combine_sims:
        return cvs
    else:
        return reduce(lambda x, y: x + y, cvs)


def avg_active_over_time(dir_name, times, cells=None, cell_type=None,
                         cell_group=None, indices=None, num_parallel=1,
                         activity_window=100):
    """Compute average number of cells active at given times.
    
    Parameters:
        dir_name -- the directory of simulations
        times -- the times to check
        cells -- the cells to check
        num_parallel -- the maximum number of files to work with simultaneously
        activity_window -- how long ago a cell may have fired and still be
                           considered active

    Deprecated Parameters:
        cell_type -- the types of cells to count
        cell_group -- the group of cells to count
        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.
    """
    active = run_f_on_dir(lambda d: d.active(times, cells, cell_type, cell_group, indices, activity_window),
                                 dir_name, num_parallel=num_parallel)
    avg_active = [0] * len(times)
    num_sims = float(len(active))
    for i in range(len(times)):
        for a in active:
            avg_active[i] += a[i]
        avg_active[i] /= float(num_sims)
    
    return avg_active

def cv(items):
    """Compute the coefficient of variation (cv) of the input.
        
        cv(items) = std(items) / mean(items)
    """
    return std(items, ddof = 1) / mean(items)



try:
    from numpy.fft import fft

    def frequencies(t, y, scale=1000):
        """Compute the magnitudes of frequencies present in y.
        
        Assumes t is evenly sampled.
        
        Set scale to 1000 if t in ms and want frequency in hz.
        """
        mag = fft(y)
        scale = 2. / len(t)
        mag = [abs(m) * scale for m in mag[0 : int(len(t) / 2)]]
        t_span = t[-1] - t[0]
        freq = [1000. * i / t_span for i in range(int(len(t) / 2))]
        return freq, mag
except ImportError:
    def frequencies(t, y, scale=1000):
        """Compute the magnitudes of frequencies present in y.
        
        Assumes t is evenly sampled.
        
        Set scale to 1000 if t in ms and want frequency in hz.
        
        Not available since cannot load numpy.
        """
        raise Exception('frequencies requires numpy, which is not available.')
    
