from neuron import h
from neuron.units import ms, mV
import matplotlib.pyplot as plt
import numpy as np
h.load_file("stdrun.hoc")


soma = h.Section(name = "soma")
soma.L = 10
soma.diam = 10/np.pi
soma.insert("hh")
soma.insert("kext_in")

iclamp = h.IClamp(soma(0.5))
iclamp.delay = 2
iclamp.dur = 500        
iclamp.amp = 0.002


def initialize_potassium():
    h.ko0_k_ion = 10         # seawater, 4 x default value (2.5)
    h.ki0_k_ion = 4*54.4     # 4 x default value, preserves ek 
initialize_potassium()

v = h.Vector().record(soma(0.5)._ref_v)  # Membrane potential vector
ko = h.Vector().record(soma(0.5)._ref_ko) # Extracellular potassium vector
ki = h.Vector().record(soma(0.5)._ref_ki) # Extracellular potassium vector
ek = h.Vector().record(soma(0.5)._ref_ek) # Potassium Nernst potential vector
t = h.Vector().record(h._ref_t)  # Time stamp vector

h.finitialize(-65 * mV)
h.continuerun(500 * ms)

plt.figure(figsize=(6.4, 4.8))

plt.subplot(3,1,1)
plt.plot(t, v, t, ek)
plt.ylabel("V, EK (mV)")
plt.gca().legend(('V', 'EK'))

plt.subplot(3,1,2)
plt.plot(t, ko)
plt.ylim(8, 30)
plt.xlabel("t (ms)")
plt.ylabel("[K+]o (mM)")

plt.subplot(3,1,3)
plt.plot(t, ki)
plt.ylim(210, 220)
plt.xlabel("t (ms)")
plt.ylabel("[K+]i (mM)")
plt.show()