import numpy as np
import matplotlib.pyplot as plt
from plotstyle import *

def hompoisson(rate, trials, duration) :
    spikes = []
    for k in range(trials) :
        times = []
        t = 0.0
        while t < duration :
            t += np.random.exponential(1/rate)
            times.append(t)
        spikes.append(times)
    return spikes

def inhompoisson(rate, trials, dt) :
    spikes = []
    p = rate*dt
    for k in range(trials) :
        x = np.random.rand(len(rate))
        times = dt*np.nonzero(x<p)[0]
        spikes.append(times)
    return spikes


def pifspikes(input, trials, dt, D=0.1) :
    vreset = 0.0
    vthresh = 1.0
    tau = 1.0
    spikes = []
    for k in range(trials) :
        times = []
        v = vreset
        noise = np.sqrt(2.0*D)*np.random.randn(len(input))/np.sqrt(dt)
        for k in range(len(noise)) :
            v += (input[k]+noise[k])*dt/tau
            if v >= vthresh :
                v = vreset
                times.append(k*dt)
        spikes.append(times)
    return spikes

# parameter:
rate = 20.0
drate = 50.0
trials = 10
duration = 2.0
dt = 0.001
tau = 0.1;

# homogeneous spike trains:
homspikes = hompoisson(rate, trials, duration)

# OU noise:
rng = np.random.RandomState(54637281)
time = np.arange(0.0, duration, dt)
x = np.zeros(time.shape)+rate
n = rng.randn(len(time))*drate*tau/np.sqrt(dt)+rate
for k in range(1,len(x)) :
    x[k] = x[k-1] + (n[k]-x[k-1])*dt/tau
x[x<0.0] = 0.0

# inhomogeneous spike trains:
#inhspikes = inhompoisson(x, trials, dt)
# pif spike trains:
inhspikes = pifspikes(x, trials, dt, D=0.3)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=cm_size(figure_width, 0.5*figure_width))
fig.subplots_adjust(**adjust_fs(fig, left=4.0, right=1.0, top=1.2))

ax1.set_title('stationary')
ax1.set_xlim(0.0, duration)
ax1.set_ylim(-0.5, trials-0.5)
ax1.set_xlabel('Time [s]')
ax1.set_ylabel('Trial')
ax1.eventplot(homspikes, colors=[lsA['color']], linelength=0.8, lw=1)

ax2.set_title('non-stationary')
ax2.set_xlim(0.0, duration)
ax2.set_ylim(-0.5, trials-0.5)
ax2.set_xlabel('Time [s]')
ax2.set_ylabel('Trial')
ax2.eventplot(inhspikes, colors=[lsA['color']], linelength=0.8, lw=1)

plt.savefig('rasterexamples.pdf')
plt.close()