%matplotlib inline

from brian2 import *

n = 100
freq = 8 * Hz
I_theta = 10.0

a = 0.02
b = 0.2
c = -50.0
d = 2.0

eqs = '''
    dv/dt = (0.04 * v**2 + 5 * v + 140 - u + I_const + theta) / (7*ms) : 1
    du/dt = a * (b * v - u) / (7*ms) : 1
    I_const : 1
    theta = I_theta * sin(2 * pi * freq * t) : 1 (shared)
'''

neurons = NeuronGroup(n, model=eqs, threshold='v > 30', reset='v = c; u = u + d')
neurons.v = -87.0 # 'rand()'
neurons.u = 0.0
neurons.I_const = '40*i/n'

S = SpikeMonitor(neurons)
trace = StateMonitor(neurons, ('v', 'theta'), record=[20,80])

settling_time = 500*ms
run(settling_time)

S = SpikeMonitor(neurons)
trace = StateMonitor(neurons, ('v', 'theta'), record=[20,80])

trace_time = 1000*ms
run(trace_time)

f, ax = subplots(3, 1, sharex=True, figsize=(13,8))
sax, vax1, vax2 = ax

sax.plot(S.t/ms, S.i, '.k')
sax.set(xlabel='Time (ms)', ylabel='Neuron index')

vax1.plot(trace.t/ms, trace.v[0], label=r'$V_m$')
vax1.plot(trace.t/ms, 2*(trace.theta[0]+1), label='theta')
vax1.set(xlabel='Time (ms)', ylabel=r'$V_m$ [Neuron 20]')
vax1.legend(loc='upper right')

vax2.plot(trace.t/ms, trace.v[1])
vax2.plot(trace.t/ms, 2*(trace.theta[0]+1))
vax2.set(xlabel='Time (ms)', ylabel=r'$V_m$ [Neuron 80]')

sax.set_xlim(settling_time / ms, settling_time / ms + trace_time / ms);

posterdir = '/Users/joe/projects/poster'
imagefn = 'mehta_burst_demo.pdf'

savepath = os.path.join(posterdir, imagefn)
f.savefig(savepath)