#Cwiczenie 1
from neuron import h
from neuron.units import ms, mV
from itertools import chain
import matplotlib.pyplot as plt 
import plotly
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}
#In Python:                      
#for sec in h.allsec(): 
#    length_constant = 1e5*numpy.sqrt(sec.diam/(4*numpy.pi*100*sec.Ra*sec.cm))
#    nseg = int((sec.L/(0.1*length_constant)+0.9)/2)*2+1
#    print(sec, nseg)

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

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
av0 = h.Vector().record(axon(0)._ref_v)  # Membrane potential vector
av1 = h.Vector().record(axon(0.5)._ref_v)  # Membrane potential vector
av2 = h.Vector().record(axon(1)._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.figure()
#plt.plot(t, v)
# plt.ylim(-90, 40)
# plt.xlabel("t (ms)")
# plt.ylabel("V (mV)")
# plt.show()

#using NEURON's gui
#ps = h.PlotShape()  # False tells h.PlotShape not to use NEURON's gui

#using matplotlib
#ps = h.PlotShape(False)  # False tells h.PlotShape not to use NEURON's gui
#ps.plot(plt)
#plt.show()

#using plotly
#ps = h.PlotShape(False)
#ps.plot(plotly).show()

#--------- Space plot at several time points---------- 

# using plotly
h.finitialize(v_init)
rvp = h.RangeVarPlot('v', axon(1), apical(1)) 
my_plot = rvp.plot(plotly, name="t=0")
for tstop in [1, 2, 3, 4, 5, 6, 7]:
    h.continuerun(tstop)
    my_plot = rvp.plot(my_plot, name=f"t={tstop}")  #f before string replaces {} expression with its value.
my_plot = my_plot.update_layout({
    "xaxis_title": "position (um)",
    "yaxis_title": "membrane potential (mV)"         
})
my_plot.show()


# using matplotlib
"""h.finitialize(v_init)
plt.figure(figsize=(12.8, 6.4))
rvp = h.RangeVarPlot('v', axon(1), apical(1)) 
my_plot = rvp.plot(plt, label='t=0')
for tstop in [1, 2, 3, 4, 5, 6, 7]:
    h.continuerun(tstop)
    my_plot = rvp.plot(plt, label=f't={tstop}')

plt.xlabel("position (um)")
plt.ylabel("membrane potential (mV)")
plt.title("Space Plot")
plt.legend()
plt.show()
"""