#the code below implements a 4th order Runge_Kutta method to numerically solve the Fitzhugh-Nagumo model and plots the results

var('v r t')
a = 0.7
b = 0.8
c = 10
I = -0.58

dv = c*(v - (1/3)*v^3 + r + I)
dr = (-1/c)*(v - a + b*r)
pts = desolve_system_rk4(des=[dv, dr], vars=[v, r], ivar=t, step=0.01, end_points=[0, 200], ics=[0, 0, 0])
ptsx = [[i, j] for i, j, k in pts]
ptsy = [[i, k] for i, j, k in pts]

p = line(ptsx, color='blue', legend_label='membrane potential (v)') + line(ptsy, color='red', legend_label='recovery variable (r)')
p.show(axes_labels=['$v$', '$r$'])

#now we will graph just dv/dt 

var('v r t')
a = 0.7
b = 0.8
c = 10
I = -0.58

dv = c*(v - (1/3)*v^3 + r + I)
dr = (-1/c)*(v - a + b*r)
pts = desolve_system_rk4(des=[dv, dr], vars=[v, r], ivar=t, step=0.01, end_points=[0, 200], ics=[0, 0, 0])
ptsx = [[i, j] for i, j, k in pts]
ptsy = [[i, k] for i, j, k in pts]

p = line(ptsx, color='blue')
p.show(axes_labels=['$t$', '$v$'])

# Now find the v and r nullclines and plot them in a phase plane


v, r, t = var('v r t')
a = 0.08
b = 0.8
c = 0.7
F = [v-(v^3)/3 - r,a*(v - b*r + c)]
n = sqrt(F[0]^2 + F[1]^2)
F_unit = [F[0]/n, F[1]/n] #set all vectors in the vector field to be same length
p = plot_vector_field(F_unit, (v,-4,4), (r,-4,4), axes_labels=['$V(t)$','$R(t)$'], xmax = 4, xmin = -4, ymax = 4, ymin = -4, aspect_ratio=1)
p += implicit_plot(F[0], (v,-4,4), (r,-4,4), color="green")
p += implicit_plot(F[1], (v,-4,4), (r,-4,4), color="red")
p

# V nullcline is r = v - (v^3)/3
# R nullcline is v = 0.8*r - 0.7

#What are the eigenvalues of the linearization of this system at the origin?
#first find the Jacobian

# Define the system of equations
var('v r t')
a = 0.7
b = 0.8
c = 10
I = 0

f1 = c*(v - (1/3)*v^3 + r + I)
f2 = (-1/c)*(v - a + b*r)

# Compute the Jacobian matrix
Jacobian = matrix([[diff(f1, v), diff(f1, r)],
                   [diff(f2, v), diff(f2, r)]])

# Evaluate the Jacobian matrix at the origin
J_at_origin = Jacobian.subs({v: 0, r: 0})

print("Jacobian Matrix:", Jacobian)

#now find the precise eigenvalues based off of the Jacobian


A = matrix([[1-1.199^2, -1], [0.08, -0.064]])
d = det(A)


A.eigenvalues()

#graph the nullclines based on I, which ranges from no external stimulus to a strong external stimulus

for i in range(6):
    v, r, t = var('v r t')
    a = 0.08
    b = 0.8
    c = 0.7
    I=i
    F = [v-(v^3)/3 - r+I,a*(v - b*r + c)]

    P = desolve_system_rk4(F,[v, r],ics=[0,-3,-3],ivar=t,end_points=500,step=0.1)
    Q = [ [j,k] for i,j,k in P]
    p = line(Q, axes_labels=['$v(t)$','$r(t)$'], thickness=2)

    n = sqrt(F[0]^2 + F[1]^2)
    F_unit = [F[0]/n, F[1]/n] #set all vectors in the vector field to be same length
    p += plot_vector_field(F_unit, (v,-4,4), (r,-4,4), axes_labels=['$x(t)$','$y(t)$'], xmax = 4, xmin = -4, ymax = 4, ymin = -4, aspect_ratio=1)

    p += implicit_plot(F[0], (v,-4,4), (r,-4,4), color="green")
    p += implicit_plot(F[1], (v,-4,4), (r,-4,4), color="red")

    show(p)

#observe the corresponding neuron spike with there is no external stimulus to a very strong external stimulus

for i in range(6):
    v, r, t = var('v r t')
    a = 0.08
    b = 0.8
    c = 0.7
    I=i
    F = [v-(v^3)/3 - r+I,a*(v - b*r + c)]

    P = desolve_system_rk4(F,[v, r],ics=[0,-3,-3],ivar=t,end_points=500,step=0.1)
    Q = [ [i, j] for i,j,k in P]

    z = line(Q, color='purple', axes_labels=['$t$','$v(t), r(t)$'], legend_label='$v(t)$', legend_color='purple', fontsize=12)
    show(z)

#now we will increase the firing frequency by making the constant 'c' smaller

var('v r t')
a = 0.7
b = 0.8
c = 3
I = -0.58

dv = c*(v - (1/3)*v^3 + r + I)
dr = (-1/c)*(v - a + b*r)
pts = desolve_system_rk4(des=[dv, dr], vars=[v, r], ivar=t, step=0.01, end_points=[0, 200], ics=[0, 0, 0])
ptsx = [[i, j] for i, j, k in pts]
ptsy = [[i, k] for i, j, k in pts]

p = line(ptsx, color='blue')
p.show(axes_labels=['$t$', '$v$'])

#now we will increase the recovery time after activation by making the constant 'b' bigger

var('v r t')
a = 0.7
b = 1.5
c = 10
I = -0.58

dv = c*(v - (1/3)*v^3 + r + I)
dr = (-1/c)*(v - a + b*r)
pts = desolve_system_rk4(des=[dv, dr], vars=[v, r], ivar=t, step=0.01, end_points=[0, 200], ics=[0, 0, 0])
ptsx = [[i, j] for i, j, k in pts]
ptsy = [[i, k] for i, j, k in pts]

p = line(ptsx, color='blue')
p.show(axes_labels=['$t$', '$v$'])