import numpy
import time
import random
from itertools import repeat

# units
second = 1.
ms = 0.001
mV = 0.001

# Parameters
N = 4000
Ne = int(N * 0.8)
Ni = N - Ne
Nplot = 4
dt = 0.1 * ms
T = 1 * second
numsteps = int(T / dt)
p = 0.02
Vr = 0 * mV
Vt = 10 * mV
we = 6. / 10. # excitatory synaptic weight (voltage)
wi = 67. / 10. # inhibitory synaptic weight
refrac = 5 * ms

# State variable S=[v;ge;gi] and variable used in Euler step
# dS=[v';ge';gi'] used in the main loop below
S = numpy.zeros((3, N))
dS = numpy.zeros((3, N))
v, ge, gi = S
v__tmp, ge__tmp, gi__tmp = dS

# Last spike times, stores the most recent time that a neuron
# has spiked, which is used for refractory periods
LS = numpy.zeros(N)
LS[:] = -2 * refrac

# Initialisation of state variables
S[0, :] = (numpy.random.randn(N) * 5 - 5) * mV
S[1, :] = numpy.random.randn(N) * 1.5 + 4
S[2, :] = numpy.random.randn(N) * 12 + 20

# Weight matrix
# Generate random connectivity matrix (note: no weights)
W = []
for _ in xrange(N):
    k = numpy.random.binomial(N, p, 1)[0]
    a = random.sample(xrange(N), k)
    a.sort()
    a = numpy.array(a)
    W.append(a)

Vrec = [[] for _ in range(Nplot)]
spikesrec = []

t = 0.
Nspikes = 0
start = time.time()
for _ in xrange(numsteps):
    # EULER UPDATE CODE, this is the update code generated by Brian for the equations:
    #dv/dt = (-v+ge*(Ee-v)+gi*(Ei-v))*(1./taum) : volt
    #dge/dt = -ge*(1./taue) : 1
    #dgi/dt = -gi*(1./taui) : 1 
    v__tmp[:] = (-v + ge * (0.06 - v) + gi * (-0.02 - v)) * (1. / 0.02)
    ge__tmp[:] = -ge * (1. / 0.005)
    gi__tmp[:] = -gi * (1. / 0.01)
    S += dt * dS

    # spike propagation, essentially the same as CUBA except I've used the names v,ge,gi
    # here instead of S[0,:] etc.
    spikes = (v > Vt).nonzero()[0]
    for i in spikes:
        if i < Ne:
            ge[W[i]] += we
        else:
            gi[W[i]] += wi

    # refractory period, if the last spike of neuron i was within refrac of the current time,
    # fix the value to Vr - note that this is how Brian used to do this, but now we have a
    # slightly different mechanism, but this is very close.
    LS[spikes] = t
    v[LS > t - refrac] = Vr

    # Monitoring, actually takes about 1/3 of the time of the simulation when I ran it,
    # for speed comparisons this bit should be commented out
    Nspikes += len(spikes)
    spikesrec += zip(spikes, repeat(t))
    for i in range(Nplot):
        Vrec[i].append(S[0, i])

    t += dt

print 'Time taken', time.time() - start
print 'Spikes', Nspikes

try:
    import pylab
    if len(spikesrec):
        sn, st = zip(*spikesrec)
        pylab.subplot(121)
        pylab.plot(st, sn, '.')
        pylab.subplot(122)
        for i in range(Nplot):
            pylab.plot(Vrec[i])
        pylab.show()
except ImportError:
    print 'Cannot plot, no pylab'
