#Cwiczenie 3
from cProfile import label
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

    v = h.Vector().record(soma(0.5)._ref_v)  # Membrane potential vector
    ax05 = h.Vector().record(axon(0.5)._ref_v)  # Membrane potential vector
    ap05 = h.Vector().record(apical(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)
    
    plt.subplot(1, 3, 1)
    plt.plot(t, ap05)
    plt.ylim(-90, 80)
    plt.xlabel("t (ms)")
    plt.ylabel("V (mV)")
    plt.title("apical")
    
    plt.subplot(1, 3, 2)
    plt.plot(t, v)
    plt.ylim(-90, 80)
    plt.xlabel("t (ms)")
    plt.ylabel("V (mV)")
    plt.title("soma")
    
    plt.subplot(1, 3, 3)
    plt.plot(t, ax05, label =f"d_lambda = {d_lambda}")
    plt.ylim(-90, 80)
    plt.xlabel("t (ms)")
    plt.ylabel("V (mV)")
    plt.title("axon")
    plt.legend()
plt.show()