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

# CONSTANTS
sim_time = 1000.0 #ms
sim_res = 1.0 #ms  IS THIS SAME NEST IZ Model's 'h' ?
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, 
    'precision': 17,
    "record_from":["V_m","U_m"]})

spikedetector = nest.Create("spike_detector", params={
    'withgid': True, #'withgid' used to tell spike detector to record the source id from which it received the event (i.e. the id of our neuron)
    'withtime': True, #'withtime' used to tell spike detector to record the time received event from source id
    'to_memory': True}) #False, # default is True


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

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

#---------------------------------------------------------------------------------
#Test-7  NEST Izhikevich ODEs in Python
# ref: https://github.com/nest/nest-simulator/blob/48c599f4e377627e75e5602d62cd481776ede1e1/models/izhikevich.cpp#L235
# Rewrite lines 224-239 from C++ to Python
h = 1. #Time::get_resolution().get_ms();
V_min= -1.8e+308
v7=c*np.ones(len(timeVec)+1)
u7=Um0*np.zeros(len(timeVec)+1)
I_e7=np.zeros(len(timeVec)+1)
spiketimes7=np.zeros(len(timeVec))
I_e7=I_e*(np.ones((int(sim_time)), dtype=np.int))

v7[0] = c +  (0.04 * Vm0+5.0) * Vm0 + 140.0 - Um0 + I_e7[0]  #v[i]+=0.5*((0.04*v[i]+5)*v[i]+140-u[i]+I[i]); i=all neurons @t
u7[0] = Um0 + a * ( 0.2 * Vm0 - Um0 ); #u[i]+=a[i]*(0.2*v[i]-u[i]); i=all neurons @t

for t in range(1,999): #  for ( long lag = from; lag < to; ++lag )
#    I_syn = B_.spikes_.get_value( lag );
    v7[t] = v7[t-1] + h * 0.5 * ( 0.04 * v7[t-1] * v7[t-1] + 5.0 * v7[t-1] + 140.0 - u7[t-1] + I_e7[t]) #+ S.I_ + I_syn );
    v7[t] = v7[t] + h * 0.5 * ( 0.04 * v7[t] * v7[t] + 5.0 * v7[t] + 140.0 - u7[t-1] + I_e7[t]) #+ S.I_ + I_syn );
    u7[t] = u7[t-1] + h * a * ( b * v7[t] - u7[t-1] );
    # lower bound of membrane potential
#    v7 = ( v7 < V_min ? V_min : v7 ); # [i for i in v7 if V_min <= i]
    if V_min > v7[t]:
        v7[t]=V_min
    # threshold crossing
    if ( v7[t] >= V_th ):
        v7[t] = c
        u7[t] = u7[t] + d
        #compute spike time
        spiketimes7[t]=1  # set_spiketime( Time::step( origin.get_steps() + lag + 1 ) );


#---------------------------------------------------------------------------------
#Test-8 Exact Iz C++ in Python
#  ref:https://www.izhikevich.org/publications/spnet.cpp, lines 784-872
#  ref: C:\Users\AR\Google Drive\(Drive) CSI 899 and Research\Reproducing IZ Model\reproducing-polychronization-master\reproducing-polychronization-master\code\original_model\poly_spnet.cpp
#  ref:https://www.izhikevich.org/publications/spnet.cpp, lines 784-872
# ref: C:\Users\AR\Google Drive\(Drive) CSI 899 and Research\Reproducing IZ Model\reproducing-polychronization-master\reproducing-polychronization-master\code\original_model\poly_spnet_bitwise_reproduction.cpp
#  lines starting at 819
v8=c*np.ones(len(timeVec)+1)
u8=Um0*np.zeros(len(timeVec)+1)
I_e8=np.zeros(len(timeVec)+1)
spiketimes8=np.zeros(len(timeVec))
I_eVec8=I_e*(np.ones((int(sim_time)), dtype=np.int))
I_e8=I_eVec8

v8[0] = c +  (0.04 * Vm0+5.0) * Vm0 + 140.0 - Um0 + I_e8[0]  #v[i]+=0.5*((0.04*v[i]+5)*v[i]+140-u[i]+I[i]); i=all neurons @t
u8[0] = Um0 + a * ( 0.2 * Vm0 - Um0 ); #u[i]+=a[i]*(0.2*v[i]-u[i]); i=all neurons @t

for t in range(1,999):     #  for (t=0;t<1000;t++)
    if ( v8[t-1] >= V_th ):     # fired = find(v>=30);
        v8[t-1] = c             # v[i] = -65; i=all neurons @t
        u8[t-1] = u8[t-1] + d        # u[i]+=d[i]; i=all neurons @t
        spiketimes8[t-1]=1      # firings[N_firings  ][0]=t;
    v8[t] = v8[t-1] + 0.5 * ( ( 0.04 * v8[t-1]+5.0 ) * v8[t-1] + 140.0 - u8[t-1] + I_e8[t] ) #v[i]+=0.5*((0.04*v[i]+5)*v[i]+140-u[i]+I[i]); i=all neurons @t
    v8[t] = v8[t] + 0.5 * ( ( 0.04 * v8[t]+5.0 ) * v8[t] + 140.0 - u8[t-1] + I_e8[t] )       #v[i]+=0.5*((0.04*v[i]+5)*v[i]+140-u[i]+I[i]); i=all neurons @t
    u8[t] = u8[t-1] + a * ( 0.2 * v8[t] - u8[t-1] ); #u[i]+=a[i]*(0.2*v[i]-u[i]); i=all neurons @t



#---------------------------------------------------------------------------------
# Compare Potentials & Spikes from NEST Model, Python (NEST ODE), Python (Paper ODE)

#---------------------------------------------------------------------------------
# TABLES

#TABLE-Spike Counts
spikeCounts = [np.count_nonzero(x1 == -65.0) for x1 in [v1,v7,v8]] #sequence.indices()
print('SPIKES:\n', spikeCounts[0], "   NEST Model\n", spikeCounts[1], "   Python (NEST ODE)\n",spikeCounts[2],"   Python (Paper ODE)") #handles=[spNEST, spODE1, spODE2]

#TABLE-Potential Values during 1st 10 ms
vResults = [(x2, v1[x2], v7[x2], v8[x2]) for x2 in range(0, 10)] #0 to 10 ms
print(tabulate(vResults, headers=["1st 10ms", "NEST Model", "Python (NEST ODE)","Python (Paper ODE)"]))

# TABLE-Spike Times, in order
tResults=[(x3, nest.GetStatus(spikedetector)[0]["events"]['times'][x3], list(compress(timeVec,(v7==-65)))[x3], list(compress(timeVec,(v8==-65)))[x3]) for x3 in range(1, min(spikeCounts))]
print(tabulate(tResults, headers=["Spiking Order", "NEST Model Sp (ms)", "Python (NEST ODE) Sp (ms)","Python (Paper ODE) Sp (ms)"]))

#---------------------------------------------------------------------------------
#PLOTS

#PLOT--POTENTIAL
fig, axs = plt.subplots(num=100,nrows=3, ncols=1)
fig.suptitle('Izhikevich Neuron:\n' +
    'NEST Model, Python (NEST ODE), Python (Paper ODE)\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 Model')
axs[0].set_xlabel('time (ms)')
axs[0].set_ylabel('V_m (mV)')
axs[0].legend(loc='upper right')
axs[1].plot(timeVec,v7[:-1], 'r', label = 'Python (NEST ODE)')
axs[1].set_xlabel('time (ms)')
axs[1].set_ylabel('V_m (mV)')
axs[1].legend(loc='upper right')
axs[2].plot(timeVec,v8[:-1], 'g', label = 'Python (Paper ODE)')
axs[2].set_xlabel('time (ms)')
axs[2].set_ylabel('V_m (mV)')
axs[2].legend(loc='upper right')

#plt.show()


#PLOT--SPIKE TIMES
dspd = nest.GetStatus(spikedetector)[0]["events"]
sp1Ind=np.where(dspd["senders"] == 1)
tspd1 = dspd["times"][sp1Ind]

plt.figure(101) # 
plt.title('Spikes: NEST Model, Python (NEST ODE), Python (Paper ODE)') 
spNEST = plt.scatter(tspd1, (dspd["senders"][sp1Ind]), color='b', s=50.0, alpha=1.0, marker='*', label='NEST spikes')
spODE1 = plt.scatter(timeVec,spiketimes8==1,color='r', s=30.0, alpha=0.5, marker='v', label = 'Python (NEST ODE) Spikes') #v4[:-1]==-65
spODE2 = plt.scatter(timeVec,spiketimes8==1,color='g', s=200.0, alpha=1.0, marker='|', label = 'Python (Paper ODE) Spikes') #v4[:-1]==-65
plt.xlabel('time (ms)')
plt.ylabel('spikes')
plt.ylim([0.5, 1.5])
plt.legend(handles=[spNEST, spODE1, spODE2])

plt.show()