from __future__ import division
# Georgia State University
# Neuroscience Institue
# Ricardo Erazo
# traveling waves in HH endogenous oscillatory neurons.

import numpy as np
from matplotlib import pyplot as plt
from scipy.integrate import solve_ivp
from tempfile import TemporaryFile


nr_neurons = 200
ap1 = nr_neurons+1
delta1 = 1e-2
sigma1 = 2e-2
gsyn= 100
domain = np.arange(0,nr_neurons)*delta1
domain = np.array([domain])
Vt = 0.0
tau2 = 3
Ie = 0.001
Ispike = 0.1*Ie*np.exp(-np.arange(0,nr_neurons,1)*delta1/sigma1)

# tau2 = 0.1
# Ie = 0.5
# Ispike = 0.3*Ie*np.exp(-np.arange(0,nr_neurons,1)*delta1/sigma1)

Isyn = Ie*np.exp(-np.arange(0,nr_neurons,1)*delta1/sigma1)
#Isyn = np.zeros((1,np.size(Isyn)))
#Isyn = np.array(Isyn)
#Isyn[0] = 1;
#Isyn[1:nr_neurons]= 0


#Experimental parameters
Htheta= 0.041326
K2theta= -0.0075

#Cell parameters
gNa= 105
ENa= 0.045
gK2= 30
EK=-0.07
gH= 4
EH= -0.021
gl= 8
El= -0.046
Vrev= 0.01
# Vrev= -0.045 - 0.05

#time scales
Cm= 0.5
tau_h= 0.04050
tau_m= 0.1
tau_k= 2
tau_s= 0.1
tau1= Cm;
int_time = 50

# ODE
#initial conditions
## rest
v0= -0.04220452
h0= 0.99225122
m0= 0.29849783
n0= 0.01681514
s0 = 0
## shocked
v2 = -0.036892644544022
h2 = 0.867431828696752
m2 = 0.016203233894593
n2 = 0.175309528389848
s2 = 1

#First shocked neuron
def vecf(t, y0):
	mNass=1./(1.+np.exp(-150.*(y0[0]+0.0305)))
	mK2ss=1./(1.+np.exp(-83.*(y0[0]+K2theta)))

	dv = (-1/Cm) * (gNa*mNass*mNass*mNass*y0[1]*(y0[0]-ENa)+ gK2*y0[3]*y0[3]*(y0[0]-EK) + gH*y0[2]*y0[2]*(y0[0]-EH) + gl*(y0[0]-El) + 0.006 +gsyn*y0[4]*(y0[0]-Vrev))
	dh = (1/(1+ np.exp(500*(y0[0]+0.0325))) - y0[1])/tau_h
	dm = (1/(1+2*np.exp(180*(y0[0]+Htheta))+np.exp(500*(y0[0]+Htheta))) -y0[2])/tau_m
	dn = (mK2ss - y0[3])/tau_k
	# for i in np.arange(0,np.size(ts)):
	# 	print i
	if (t> ts):
		ds = np.exp(-(t-ts)/tau2)
	else:
		ds = 0
	# ds = (1/(1+np.exp(-5000*(y0[0]-Vrev)))-y0[4])/0.01
	f = dv, dh, dm, dn, ds
	return f

def Vt_cross(t,y0): return y0[0]-Vt
# Vt_cross.terminal= False
Vt_cross.terminal= True
Vt_cross.direction = 1

#initial conditions
y0 = np.array((v0,h0,m0,n0, s0))
y2 = np.array((v2,h2,m2,n2, s2))
#spike_times = np.zeros((1,nr_neurons))
statevar = np.zeros((nr_neurons,5))
for i in np.arange(0,nr_neurons):
	statevar[i,:] = y0
t1= 0
# ts = np.zeros((nr_neurons,1))
ts=0
neuron = solve_ivp(vecf, [t1, t1+30], y2, events = Vt_cross)
ts = neuron.t_events[0]
ts=ts[0]
neuron1 = solve_ivp(vecf, [t1, t1+30], y0, events = Vt_cross)

# speed = delta1/np.diff(np.transpose(ts))
# c = speed[0]
# acceleration = np.diff(speed)/np.diff(np.transpose(ts[0:-1]))
# a = acceleration[0,:]
# plt.figure()
# plt.ion()
# plt.plot(delta1*np.arange(0,np.size(c)),c)