#!/usr/bin/env python
# Hodgkin-Huxley model
# saltatory conduction
#

import numpy as np
import os
from neuron import h
import matplotlib.pyplot as plt
import matplotlib.animation as animation

segLength = 5.0
NodeSegs = 1
AxonSegs = 59
npnts = 10001

axonDiam = 5.0
Ra = 100.0 # axial registance
h.celsius = 35.0
fNa = 40.0  # density factor for Na channel
fK  = 40.0  # density factor for K channel

dx = 1.0/(npnts-1)
ulen = NodeSegs + AxonSegs

tstop = 10.0
h.dt  = 0.01

recstart = 0.0
recend = tstop
ntraces = int((recend-recstart)/h.dt) + 1

print("npnts = %d " % (npnts))
print("total axon length = %f mm" % (segLength*(npnts-1)/1000.0) )
print("ntraces (time points) = %d " % (ntraces))

# soma
soma = h.Section()
soma.L = 20.0
soma.diam = 20.0
soma.nseg = 1
soma.Ra = Ra

# axon intial segment
ais = h.Section()
ais.nseg = 3
ais.L = ais.nseg * segLength
ais.diam = 5.0
ais.Ra = Ra

# Hodgkin-Huxley default settings
gnabar_init = 0.12
gkbar_init = 0.036
gl_init = 0.0003
el_init = -54.3

axon = h.Section()

axon.nseg = npnts
axon.L = segLength * npnts
axon.diam = axonDiam
axon.Ra = Ra


soma.insert('hh')
soma.insert('pas')

ais.insert('hh')
ais.insert('pas')

for seg in ais:
    seg.hh.gnabar = fNa*gnabar_init
    seg.hh.gkbar = gkbar_init
    seg.hh.gl = gl_init

axon.insert('hh')
axon.insert('pas')

ais.connect(soma, 0, 0)
axon.connect(ais, 1, 0)

segs = [seg for seg in axon]

for i in range(axon.nseg):
    r = i % ulen
    segs[i].ena    = 115.0 - 65.0
    segs[i].ek     = -12.0 - 65.0
    segs[i].hh.el  = -0.05 - 65.0
    if r >= ulen-NodeSegs:
        segs[i].hh.gnabar = fNa * gnabar_init
        segs[i].hh.gkbar  = fK  * gkbar_init
        segs[i].hh.gl     = gl_init
    else:
#        print('myelinated ', i)
        segs[i].hh.gnabar = 0.0
        segs[i].hh.gkbar  = 0.0
        segs[i].hh.gl     = 0.0
        segs[i].cm = 0.002

# stimulation

stim1 = h.IClamp(0.5, sec=soma)
stim1.amp   = 5.00
stim1.delay = 1.0
stim1.dur   = 0.2

"""
#second stimulation for collision experiment
stim2 = h.IClamp(1.0, sec=axon)
stim2.amp   = 5.00
stim2.delay = 1.0
stim2.dur   = 1.0
"""

v_init = -70.0
h.finitialize(v_init)

fs = open("a.dat", "wb")
a = np.zeros(14)
b = np.zeros(npnts)
bb = np.zeros((ntraces,npnts))
j = 0
v0 = v_init
v1 = v_init
t0 = 0.0
t1 = 0.0
while h.t <= tstop:
    a[0] = h.t
    a[1] = soma(0.5).v
    a[2] = axon(0.0).v
    a[3] = axon(0.1).v
    a[4] = axon(0.2).v
    a[5] = axon(0.3).v
    a[6] = axon(0.4).v
    a[7] = axon(0.5).v
    a[8] = axon(0.6).v
    a[9] = axon(0.7).v
    a[10] = axon(0.8).v
    a[11] = axon(0.9).v
    a[12] = axon(0.0).v
    a[13] = axon(1.0).v

    if a[2] > v0:
        v0 = a[2]
        t0 = h.t
    if a[13] > v1:
        v1 = a[13]
        t1 = h.t
    fs.write(a)
    if h.t>= recstart and h.t <= recend:
        for k in range(npnts):
            b[k] = axon(dx*k).v
        bb[j] = b
        j = j+1
    h.fadvance()
fs.close()
os.system('~/bin/gc a.dat &')

print('t0 = %f' % (t0) )
print('t1 = %f' % (t1) )
print('conduction velocity = %f m/s' % (segLength*(npnts-1)/1000.0/(t1-t0)))

"""
fs2 = open("b.dat", "wb")
a2 = np.zeros(ntraces+1)
for i in range(npnts):
    a2[0] = dx*i
    for j in range(ntraces):
        a2[j+1] = bb[j,i]
    fs2.write(a2)
fs2.close()
os.system('~/bin/gc b.dat &')
"""

#
#  movie
#
##fig = plt.figure(figsize=(20,5))
##ax = plt.subplot()

fig, ax = plt.subplots()

xdata = np.linspace(0.0, npnts*segLength/1000.0, npnts)
ydata = np.zeros(npnts)
line, = ax.plot(xdata, ydata)
ax.set_xlim(0.0, npnts*segLength/1000.0)
ax.set_ylim(-100.0, 50.0)
nframes = 5

def update(i):
    ydata = bb[nframes*i]
    line.set_ydata(ydata)
    return line,

ani = animation.FuncAnimation(fig, update, frames = int(ntraces/nframes), interval=30)
plt.show()