from brian2 import *

v_th = 15 / 1000
tau = 50 * ms
P = PoissonGroup(5, 50 * Hz)
MP = SpikeMonitor(P, record=True)
eqs = '''
dv/dt = -v/tau : 1
w_pool : 1
is_weight_release: boolean
is_add_weight :boolean
w_to_update : 1
hold_w : 1
delta_w : 1
hold_i : 1
'''

syn_eq = '''
w : 1
'''

on_pre_eq = {'pre': '''
w_to_update = w
hold_i = i
is_weight_release = True
'''}

on_post_eq = {'weight_release_pathway': '''
hold_w = w
w -= (hold_w*(hold_w - w_to_update)/(5 + hold_w))*(hold_w > w_to_update)*(hold_i != i)
delta_w += (hold_w*(hold_w - w_to_update)/(5 + hold_w))*(hold_w > w_to_update)*(hold_i != i)
''', 'add_weight_pathway': '''
hold_w = w
w += delta_w*(hold_i == i)
w += 0.8*w_pool*(hold_i == i)
delta_w -= delta_w*(hold_i == i)
w_pool -= 0.8*w_pool*(hold_i == i)
'''}

G = NeuronGroup(1, model=eqs, threshold='v>v_th', method='exact', events={
    'weight_release': 'is_weight_release',
    'add_weight': 'is_add_weight'
})
G.run_on_event('weight_release', '''is_weight_release = False
is_add_weight = True''')

G.run_on_event('add_weight', '''
is_add_weight = False
''')

S = Synapses(P, G, model=syn_eq, on_pre=on_pre_eq, on_post=on_post_eq, on_event={
    'weight_release_pathway': 'weight_release',
    'add_weight_pathway': 'add_weight'
})

S.connect()

MW = StateMonitor(S, 'w', record=True)  # StateMonitor for monitoring weights.
MW_POOL = StateMonitor(G, 'w_pool', record=True)  # StateMonitor for monitoring w_pool.
S.w = 0
G.w_pool = 5

run(100*ms)

print(S.w)
print(G.w_pool)
print(MP.t)
print(MP.i)

plt.figure(100)
# Drawing Figure 2.
plt.plot(MW.t / ms, MW.w[0], label='Synapse 1', color='g')
plt.plot(MW.t / ms, MW.w[1], label='Synapse 2', color='r')
plt.plot(MW.t / ms, MW.w[2], label='Synapse 3', color='b')
plt.plot(MW.t / ms, MW.w[3], label='Synapse 4', linestyle='dashed', color='c')
plt.plot(MW.t / ms, MW.w[4], label='Synapse 5', linestyle='dashed', color='m')
plt.plot(MW_POOL.t / ms, MW_POOL.w_pool[0], color='k', label='POOL')
xlabel('Time (ms)')
ylabel('w')
legend()
plt.show()