from scipy.special import erf
from scipy import stats
import nest
import matplotlib.pyplot as plt
import numpy as np

# CONSTANTS
sim_time = 1000.0 #ms
sim_res = 1. #ms
timeVec=list(range(0,(int(sim_time)-1))) #ms
a=0.02 
b=0.2 
c=-65.0 
d=8.0 
Vm0=-65.0 
Um0=b*Vm0 
V_th=30.0 # *** threshold for spiking ****
I_e=5.0 
I_eVec=I_e*(np.ones((int(sim_time)-1), dtype=np.int))

nest.ResetKernel()
nest.SetKernelStatus({'resolution': sim_res}) 

#---------------------------------------------------------------------------------
#Test-1, NEST "Izhikevich" Model
test1='Nest Iz'
neuron1 = nest.Create('izhikevich', params={
    'a': a, 
    'b': b,
    'c': c, 
    'd': d,
    'I_e': I_e, 
    'consistent_integration': False,
    'U_m': Um0,
    'V_m': Vm0,    
    'V_th': V_th})

mm = nest.Create("multimeter", params={ 
    "withtime":True, 
    'withgid': True, 
    "record_from":["V_m"]})

nest.Connect(mm, neuron1)
nest.Simulate(sim_time)

v1=np.array(nest.GetStatus(mm)[0]["events"]["V_m"],dtype=np.float64) #V_m vector for plot

#---------------------------------------------------------------------------------
#Test-2, Izhikevich ODEs Explicit
# paper ref: https://www.izhikevich.org/publications/spikes.pdf,  # code ref: https://www.izhikevich.org/publications/net.m
test2='Iz ODEs'
v2 = c*np.ones(len(timeVec)+1)    #v2[0] = Vm0 = -65
u2 = Um0*np.zeros(len(timeVec)+1) #u2[0] = Um0 = -13
I_e2=np.zeros(len(timeVec))

for t in range(0,len(range(0,len(timeVec)))):
    I_e2[t]=I_eVec[t]
    if (v2[t]<30): #calaculates next step's v,u using this steps values calc from previous step
        v2[t+1]=v2[t] + 0.5*(0.04*v2[t]**2 + 5*v2[t] + 140 - u2[t] + I_e2[t])
        v2[t+1]=v2[t+1] + 0.5*(0.04*v2[t+1]**2 + 5*v2[t+1] + 140 - u2[t] + I_e2[t])
        u2[t+1]=u2[t] + a*(b*v2[t+1] - u2[t])
    else: #adjusts this step's v --> V_th and resets v, u for next step
        v2[t]=30
        v2[t+1]=c 
        u2[t+1]=u2[t]+d                                                                                                                                                                             

#---------------------------------------------------------------------------------
#PLOT
fig, axs = plt.subplots(2)
fig.suptitle('NEST vs ODE Izhikevich Neuron Model Potentials\n' + 
    r'Params:  a: %s  b: %s  c: %s  d: %s  V_th %s  I_e: %s' %(a,b,c,d,V_th,I_e))
axs[0].plot(timeVec,v1, 'b-', label = 'NEST')
axs[0].set_xlabel('time (ms)')
axs[0].set_ylabel('membrane potential mV')
axs[0].legend(loc='upper right')
axs[1].plot(timeVec,v2[:-1], 'r', label = 'ODE')
axs[1].set_xlabel('time (ms)')
axs[1].set_ylabel('membrane potential mV')
axs[1].legend(loc='upper right')
plt.show()