from neuron import h, gui
from matplotlib import pyplot
h.load_file('stdlib.hoc')


class BallStick:
    def __init__(self, _id):
        self._id = _id

        # topology
        self.soma = h.Section(name='soma', cell=self)
        self.dend = h.Section(name='dend', cell=self)
        self.dend.connect(self.soma)

        # geometry
        self.dend.L = 500
        self.dend.diam = 3
        self.soma.L = 20
        self.soma.diam = 20

        self.all = [self.soma, self.dend]

        # biophysics
        self.soma.insert('hh')
        for sec in self.all:
            sec.Ra = 100  # Ohm * cm
        
        # synapse
        self.syn = h.ExpSyn(self.dend(0.5))
        self.syn.tau = 3

        self._discretize()

    def _discretize(self):
        for sec in self.all:
            sec.nseg = int((sec.L / (0.1 * h.lambda_f(100)) + 0.9) / 2.) * 2 + 1

    def __str__(self):
        return 'BallStick[{}]'.format(self._id)


"""
cell will fire with current injection to center of soma
delay = 1
dur = 0.1
amp = 5
"""

if __name__ == '__main__':
    # create the cells
    cell0 = BallStick(0)
    cell1 = BallStick(1)

    # connect the cells
    nc = h.NetCon(cell0.soma(0.5)._ref_v, cell1.syn, sec=cell0.soma)
    nc.threshold = 0
    nc.weight[0] = 0.01
    nc.delay = 1

    # make the 0th cell fire
    ic = h.IClamp(cell0.soma(0.5))
    ic.delay = 1
    ic.dur = 0.1
    ic.amp = 5

    # record the cells
    t = h.Vector()
    v0 = h.Vector()
    v1 = h.Vector()
    t.record(h._ref_t)
    v0.record(cell0.soma(0.5)._ref_v)
    v1.record(cell1.soma(0.5)._ref_v)

    # initialize and run
    h.finitialize(-65)
    h.continuerun(15)

    # plot it
    pyplot.plot(t, v0)
    pyplot.plot(t, v1)
    pyplot.show()