#Cwiczenie 4
from neuron import h
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))

plt.figure(figsize=(12.8, 6.4))
for d_lambda in [1, 0.3, 0.2, 0.1, 0.05]:  
    # 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

    # biophysics
    for sec in h.allsec (): 
        sec.Ra =  100
        sec.cm = 1

    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/(d_lambda*length_constant)+0.9)/2)*2+1
        print(sec, sec.nseg) 

    soma.insert("hh")

    apical.insert("pas")

    basilar.insert("pas")

    for seg in chain (apical, basilar): #iterator over two sections
        seg.pas.g = 0.0002 
        seg.pas.e = -65 

    axon.insert("hh")

# --------------------- Instrumentation ---------------------

    syn = h. AlphaSynapse (0.5, sec = soma) 
    syn.onset = 0.5
    syn.gmax = 0.05
    syn.e = 0

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

    h.dt = 0.025
    v_init = -65

    h.finitialize(v_init)
    h.fcurrent()         
  
    tstop = 1.5
    h.continuerun(tstop)
    rvp = h.RangeVarPlot('v', axon(1), apical(1)) 
    my_plot = rvp.plot(plt, label=f'dl={d_lambda}')

plt.xlabel("position (um)")
plt.ylabel("membrane potential (mV)")
plt.title("Space plot t = 1.5 ms")
plt.legend()
plt.show()