"""
Visualization tools for snnet 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_dir
from parallelManagement import parallel_run_over_queue
import os
import platform

def default_filename_rule(filename):
    """Data stored in a/b/c.sdat should have image a_b_c.pdf"""
    if filename[-4 :] == '.pkl':
        filename = filename[0 : -4]
    elif filename[-5 :] == '.sdat':
        filename = filename[0 : -5]
    tails = []
    while filename != '':
        filename, tail = os.path.split(filename)
        tails.append(tail)
    return '_'.join(tails[::-1]) + '.pdf'

    
def make_rasters(data_dir, image_dir='.', plot_types=None, shading=None, shading_color=[.5, 1, 1],
                 show_field_potentials=True,
                 filename_rule=default_filename_rule, image_processing=None,
                 num_parallel=1):
    """Make raster diagrams for all sdat files in a directory (and its subdirectories).

    Parameters:
        data_dir -- the directory with the data
        image_dir -- the directory to put the image files
        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
        filename_rule -- a function that takes a path in data_dir, returns a filename for the image
        image_processing -- a function to run on the figure before it is saved (or None)
        num_parallel -- maximum number of files to work on simultaneously
                
    Note:
    
    If plot_types is None, then plots for all types of neurons.
    """
    
    # graphics must either be done one at a time, or in separate processes
    # for now
    if platform.system() not in ['Darwin', 'Linux']:
        num_parallel = 1
        
    run_f_on_dir(lambda data, filename: _make_rasters(data, filename, image_dir, plot_types, shading,
                                                             shading_color, show_field_potentials, filename_rule,
                                                             image_processing),
                        data_dir, enter_subdirs = True, num_parallel = num_parallel)


def _make_rasters(data, filename, image_dir, plot_types, shading, shading_color, show_field_potentials,
                  filename_rule, image_processing):
    """"Helper function for make_rasters."""
    figure()
    raster(data, plot_types, shading, shading_color, show_field_potentials)
    if image_processing is not None:
        image_processing()
    savefig(os.path.join(image_dir, filename_rule(filename)))
    close()
    


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]) in (int, float):
        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)
    
    # compute how many of each type of plotted neuron we have
    type_count = [neuron_types.count(t) for t in plot_types]
    
    num_rows = len(plot_types)
    
    # add shading, set holds
    for j in range(num_rows):
        subplot(num_rows, num_cols, num_cols * j + 1)
        y_hi = type_count[j]
        if shading is not None:
            for i, t in enumerate(shading):
                if len(t) != 2:
                    raise Exception('Time range ' + repr(t) + ' does not have length 2.')
                _shade_time_period(data.time(t[0]), data.time(t[1]), 0, y_hi, shading_color)
        hold(True)
    
    # 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
    lines = [[] 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]:
                lines[j].append([(t, y + .1), (t, y + .9)])
            type_count[j] += 1
        
    # plot the lines
    for i in range(num_rows):
        subplot(num_rows, num_cols, num_cols * i + 1)
        line_segments = LineCollection(lines[i], linewidth = .5, colors = [0, 0, 0, 1])
        gca().add_collection(line_segments)
    
    # label axes, adjust limits, adjust tick size
    for j in range(num_rows):
        subplot(num_rows, num_cols, num_cols * j + 1)
        ylabel(_clean_name(plot_types[j]))
        ylim(0, type_count[j])
        xlim(0, data.run_time)
        xtick = gca().get_xmajorticklabels()
        ytick = gca().get_ymajorticklabels()       
        for x in xtick:
            x.set_size('small')
        for y in ytick:
            y.set_size('small')
        draw()
    
    # plot field potentials if available
    if has_field_potentials:
        for j in range(len(plot_types)):
            subplot(num_rows, num_cols, num_cols * j + 2)
            plot(data.t, field_potentials[plot_types[j]], color = 'black')
            y_lo, y_hi = ylim()
            hold(True)
            for i, t in enumerate(shading):
                _shade_time_period(data.time(t[0]), data.time(t[1]), y_lo, y_hi, shading_color)
            ylim(y_lo, y_hi)
            xlim(0, data.run_time)
            xtick = gca().get_xmajorticklabels()
            ytick = gca().get_ymajorticklabels()       
            for x in xtick:
                x.set_size('8')
            for y in ytick:
                y.set_size('8')
            draw()
            
    subplot(num_rows, num_cols, 1)
    title('Activity by Neuron')
    if has_field_potentials:
        subplot(num_rows, num_cols, 2)
        title('Field Potentials')
    

def _shade_time_period(t_lo, t_hi, y_lo, y_hi, shading_color):
    """Draws a filled rectangle with no edge."""
    fill([t_lo, t_lo, t_hi, t_hi], [y_lo, y_hi, y_hi, y_lo], facecolor = shading_color, edgecolor = 'none')

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