"""
Visualization tools for SN Net simulations.

Robert A McDougal, 2011
"""

from matplotlib.pyplot import *
from matplotlib.collections import LineCollection
from sim_data import load
from analysis import run_f_on_pkl_in_dir
from parallelManagement import parallel_run_over_queue
import os
import platform



def raster(data, plot_types = None, shading = None,
           shading_color = [.5, 1, 1], show_field_potentials = True):
    """Plot a raster diagram of activity.
    
    Parameters:
        data -- the SimData (or filename containing SimData) to plot
        plot_types -- the types of neurons to plot (defined by filename)
        shading -- list of time intervals to shade (or None)
        shading_color -- rgb color triplet for shading
        show_field_potentials -- do we plot the field potentials
                                 ignored if f.p. data not in data
    """
    
    if isinstance(data, str):
        data = load(data)
    
    spike_times = data.spike_times
    neuron_types = data.types
    
    if shading is None:
        shading = []
    if len(shading) == 2 and type(shading[0]) is int:
        shading = [shading]

    # check if field potential data is available
    if show_field_potentials:
        try:
            field_potentials = data.field_potentials
            has_field_potentials = True
            num_cols = 2
        except:
            has_field_potentials = False
            num_cols = 1
    else:
        has_field_potentials = False
        num_cols = 1
    
    # if no list of neuron types to plot is passed, then plot all of them
    if plot_types is None:
        # find the various types of neurons used (only list once each)
        plot_types = []
        for t in neuron_types:
            if t not in plot_types:
                plot_types.append(t)
   
    num_rows = len(plot_types)
    gnuplot_script = ''
    
    #
    # compute the data
    #

    # type count will be the number of times seen a given type of neuron (0)
    type_count = [0] * len(plot_types)
    
    # loop through neurons, compute spike lines
    plot_data = ['' for t in plot_types]
    for i, current_type in enumerate(neuron_types):
        if current_type in plot_types:
            j = plot_types.index(current_type)
            y = type_count[j]
            for t in spike_times[i]:
                plot_data[j] += '%g %g %g %g %g\n' % (t, y + .1, y + .1, y + .9, y + .9)
            type_count[j] += 1
    
    data_files = []
    current = 0
    
    data_files.append('tmp_snnet_raster.dat')    
    f = file('tmp_snnet_raster.dat', 'w')
    f.write(plot_data[current])
    f.close()
    
    #export_to_csv('tmp_snnet_field_potentials.dat', data.t, data.field_potentials[plot_types[current]], separator = ' ')
    
    os.system('gnuplot tmp_snnet_raster.gnuplot')
    os.system('ps2pdf raster.ps raster.pdf')


def export_to_csv(filename, *args, **kwargs):
    """Save lists to csv file.
    
    Parameters:
        filename -- output csv file
        list1, list2, ... -- lists to output to the file
        separator -- separator to use between data points.
        
    Note:
    
    If no lists are passed, no file will be created.
    
    separator defaults to ', '
    
    Example usage:
    
    export_to_csv('results', x, y)
    export_to_csv('results', x, y, separator = ' ')
    """
    
    separator = kwargs.pop('separator', ', ')
    newline = kwargs.pop('newline', '\n')
    if kwargs:
        raise Exception('Unknown option: ' + kwargs.values())
    
    if len(args) == 0:
        return
    
    list_length = len(args[0])
    
    for a in args:
        if len(a) != list_length:
            raise Exception('all lists must be the same length')
    
    file = open(filename, 'w')
    for i in range(list_length):
        row = []
        for a in args:
            row.append(str(a[i]))
        file.write(separator.join(row) + newline)
    file.close()


def _clean_name(name):
    """Removes .snnet (if present) from name strings."""
    if name[-6 : ] == '.snnet':
        return name[0 : -6]
    else:
        return name
