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

# --------------------- Model specification ---------------------

# topology
soma = h.Section(name = "soma")
apical = h.Section(name = "apical")
basilar = h.Section(name = "basilar")
axon = h.Section(name = "axon")

apical.connect(soma)   #0-end of apical section ia attached to 1-end of a soma section
basilar.connect(soma(0))
axon.connect(soma(0))

# geometry
soma.L = 30       #if no units are specified, NEURON assumes um
soma.diam = 30
soma.nseg = 1

apical.L = 600
apical.diam = 1
apical.nseg = 23

basilar.L = 200
basilar.diam = 2
basilar.nseg = 5

axon.L = 1000
axon.diam = 1
axon.nseg = 37

#To calculate number of segments use formula:
#(pg.29, Chapter 5 of the NEURON book):
#This makes use of the function lambda () included in standard NEURON library stdlib.hoc
#forall {nseg = int((L/(0.1*lambda_f(100))+0.9)/2)*2+1}
#forall {print nseg}
#func lambda_f() { // currently accessed section, $1 == frequency
#return 1e5*sqrt(diam/(4*PI*$1*Ra*cm))

#In Python:                      
#for sec in h.allsec(): 
#    length_constant = 1e5*numpy.sqrt(sec.diam/(4*numpy.pi*100*sec.Ra*sec.cm))
#    sec.nseg = int((sec.L/(0.1*length_constant)+0.9)/2)*2+1
#    print(sec, sec.nseg)


# biophysics
for sec in h.allsec(): 
    sec.Ra =  100    # Axial resistance in Ohm * cm
    sec.cm = 1       # Membrane capacitance in micro Farads / cm^2

soma.insert("hh")

apical.insert("pas")

basilar.insert("pas")

for seg in chain (apical, basilar): #iterator over two sections
    seg.pas.g = 0.0002  # pas conductance in S/cm2 
    seg.pas.e = -65     # pas reversal potential in mV

axon.insert("hh")

# --------------------- Instrumentation ---------------------
# Synaptic input
syn = h.AlphaSynapse (0.5, sec = soma) 
syn.onset = 0.5     #in ms
syn.gmax = 0.05     #in micro S
syn.tau = 0.1       #in ms
syn.e = 0           #in mV

# Recordings
v = h.Vector().record(soma(0.5)._ref_v)  # Membrane potential vector
t = h.Vector().record(h._ref_t)  # Time stamp vector

# --------------------- Simulation control ---------------------

h.dt = 0.025
tstop = 10
v_init = -65

h.finitialize(v_init)
h.fcurrent()
h.continuerun(tstop)

# --------------------- Soma membrane voltage plot ---------------------

plt.figure()
plt.plot(t, v)
plt.ylim(-90, 40)
plt.xlabel("t (ms)")
plt.ylabel("V (mV)")
plt.show()