# Numerical Simulation of a Cortical Neuron Network
# Based on Izhikevich's MATLAB code, adapted for Python

import numpy as np
import matplotlib.pyplot as plt

Ne = 800
Ni = 200
re = np.random.rand(Ne, 1)
ri = np.random.rand(Ni, 1)
a = np.vstack((0.02*np.ones((Ne,1)),0.02+0.08*ri)).flatten()
b = np.vstack((0.2*np.ones((Ne,1)),0.25-0.05*ri)).flatten()
c = np.vstack((-65+15*re**2,-65*np.ones((Ni,1)))).flatten()
d = np.vstack((8-6*re**2,2*np.ones((Ni,1)))).flatten()
S = np.hstack((0.5*np.random.rand(Ne+Ni,Ne),-1*np.random.rand(Ne+Ni,Ni)))
v = -65*np.ones(Ne+Ni)
u = b*v
firings = []
T = 1000

for t in range(1,T+1):
    I = np.concatenate((5*np.random.randn(Ne), 2*np.random.randn(Ni)))
    fired = np.where(v>=30)[0]
    if fired.size > 0:
        firings.extend([[t,idx] for idx in fired])
        v[fired] = c[fired]
        u[fired] += d[fired]
        I += np.sum(S[:, fired], axis=1)
    v += 0.5*(0.04*v**2 + 5*v +140 - u + I)
    v += 0.5*(0.04*v**2 + 5*v +140 - u + I)
    u += a*(b*v - u)

f = np.array(firings,dtype=float)
plt.figure(figsize=(10,6))
plt.plot(f[:,0],f[:,1],'.',markersize=2,color='black')
plt.xlabel('Time (ms)')
plt.ylabel('Neuron Number')
plt.title('Neuron Network Simulation, Size = 1000')
plt.show()