# Set up imports and matplotlib options
%pylab inline
import numpy as np
import matplotlib.pyplot as plt
import quantities as pq
from matplotlib import rc
rc('font',**{'family':'sans-serif',
             'sans-serif': 'Helvetica, Arial',
             'weight': 'medium',
             'size': 18})
rc('text', usetex=False)
rc('lines', linewidth=3)

mp = 0.115
vp = 35.0
mo = 0.265
vo = 0.0

v = (mp * vp + mo * vo) / (mp + mo)
print v

vp = 40.0
v = (mp * vp + mo * vo) / (mp + mo)
print v

vp = np.arange(51.)
v = (mp * vp + mo * vo) / (mp + mo)
print v

plt.plot(vp, v)

vp = 35.0
mp = np.arange(0., 2.05, 0.05)
v = (mp * vp + mo * vo) / (mp + mo)
plt.plot(mp, v)

mp = 0.115 * pq.kg
vp = 35.0 * (pq.m / pq.s)
mo = 0.265 * pq.kg
vo = 0.0 * (pq.m / pq.s)
v = (mp * vp + mo * vo) / (mp + mo)
print v

vo = 0.0  # It's just zero, who cares...
v = (mp * vp + mo * vo) / (mp + mo)

mp = np.arange(0., 2.05, 0.05) * pq.kg
vp = 35.0 * (pq.m / pq.s)
mo = 0.265 * pq.kg
vo = 0.0 * (pq.m / pq.s)
v = (mp * vp + mo * vo) / (mp + mo)
plt.plot(mp, v)

plt.xlabel("Mass of puck in " + mp.dimensionality.string)
plt.ylabel("Velocity of pucktopus in " + v.dimensionality.string)
plt.plot(mp, v)

bags = pq.UnitQuantity('milk bags', pq.L * (4. / 3.))
person = pq.UnitQuantity('person', pq.dimensionless * 1)

need = 200 * (pq.mL / person)

need *= 250 * person
print need

print need.rescale(pq.L)

print need.rescale(bags)

print np.ceil(need.rescale(bags))

print need.item()
print need.rescale(bags).item()

def generate_spikes(filename, neurons=20, time=20 * pq.s, event=11 * pq.s):
    with open(filename, 'w') as f:
        # Choose random rates
        rates = np.random.randint(2, 10, size=neurons) * pq.Hz
        for rate in rates:
          # Randomly generate spikes
          spikes = np.random.uniform(0, time, size=time * rate)
          # Add spikes around the event
          spikes = np.append(spikes, np.random.normal(event, size=rate * 5))
          # Sort and save to text file
          np.savetxt(f, (np.sort(spikes),), fmt="%f", delimiter=', ')

# Uncomment to call this function
# generate_spikes("spikes.csv")

def load_spikes(filename):
    with open(filename, 'r') as f:
        lines = f.readlines()
    return [np.array(map(float, line.split(','))) for line in lines]

spiketrains = [st * pq.s for st in load_spikes("spikes.csv")]

def raster_plot(spiketrains, event=None, window=None):
    plt.clf()  # Start a fresh plot
    if event is not None:
        event.units = spiketrains[0].units
        plt.axvline(event, lw=2, color='r')
        plt.text(event, len(spiketrains) + 0.01 * len(spiketrains), "Event",
          color='r', horizontalalignment='center', fontsize=12)

        if window is not None:
            window.units = spiketrains[0].units
            plt.axvspan(event + window[0], event + window[1],
                        facecolor='#9696FF', zorder=0, ec='none')
    
    lastspike = 0.0 * pq.s
    firstspike = np.amin(spiketrains[0])
    for j, st in enumerate(spiketrains):
        lastspike = max(lastspike, np.amax(st))
        firstspike = min(firstspike, np.amin(st))
        plt.scatter(st, j * np.ones(st.shape), marker='|', color='k')
    plt.ylim(-1, len(spiketrains))
    plt.xlim(firstspike, lastspike)
    plt.xlabel("Time in " + spiketrains[0].dimensionality.string)
    plt.tight_layout()
    plt.show()

raster_plot(spiketrains)

event = 11 * pq.s
raster_plot(spiketrains, event=event)

window = (-0.5, 2) * pq.s
raster_plot(spiketrains, event=event, window=window)

def time_slice(spikes, tstart, tend):
    return spikes[np.logical_and(tstart <= spikes, spikes < tend)]

perievent = [time_slice(spikes, event + window[0], event + window[1])
             for spikes in spiketrains]
raster_plot(perievent, event=event)

event = 7000 * pq.ms
window = (-2500, 1000) * pq.ms
raster_plot(spiketrains, event=event, window=window)

event = 5.0 * pq.s
window = (-1000, 1000) * pq.ms
raster_plot(spiketrains, event=event, window=window)

def bin_spikes(spikes, tstart, tend, binsize):
    binsize.units = tstart.units
    bins = np.arange(tstart, tend + binsize, binsize)
    return np.histogram(spikes, bins=bins)[0]

binsize = 20 * pq.ms
binned = np.vstack([bin_spikes(st, 0 * pq.s, 20 * pq.s, binsize) for st in spiketrains])

def binned_plot(bins, window=None):
    plt.clf()
    plt.pcolormesh(bins, cmap="gray_r")
    if window is not None:
        plt.axvspan(window[0], window[1], facecolor='b', alpha=0.5, ec='none')
    plt.xlabel("Time bin")
    plt.xlim(0, bins.shape[1])

    # Discrete color bars take a bit of code
    bounds = np.arange(0.5, np.amax(bins) + 1)
    norm = mpl.colors.BoundaryNorm(bounds, len(bounds) - 1)
    plt.tight_layout()
    plt.colorbar(format="%d", spacing='proportional',
      norm=norm, boundaries=bounds)

binned_plot(binned, (525, 650))

peribinned = np.vstack([bin_spikes(st, event + window[0],
                    event + window[1], binsize) for st in spiketrains])
binned_plot(peribinned)

binsize = 50 * pq.ms
binned = np.vstack([bin_spikes(st, 0 * pq.s, 20 * pq.s, binsize) for st in spiketrains])
binwindow = (event + window).rescale(binsize.units) / binsize
binned_plot(binned, binwindow)

binsize = 0.05 * pq.s
binned = np.vstack([bin_spikes(st, 0 * pq.s, 20 * pq.s, binsize) for st in spiketrains])
binwindow = (event + window).rescale(binsize.units) / binsize
binned_plot(binned, binwindow)