import sim_ballena as ballena
import matplotlib.pyplot as plt

neu = [ ballena.Lif().tau(0.5).v_rest(-80), 
       ballena.Lif().tau(2).v_rest(-75), 
       ballena.Lif().tau(5).v_rest(-70) ]

net = ( ballena.Network(neu)
       .synapses_in( [(0,0),(1,1),(2,2)] )
       .synapses_net([]).weights([5,5,5])
       .outputs([0,1,2])
       .mode(['VOLTAGES']) )

instance = ballena.Instance( [(2,0),(3,1),(4,2)] )
res      = net.simulate( instance, 10 )

t = res.time()
v = res.v_trace()

for idx in range(3):
    plt.plot( v[idx], label=f'neu {idx}' )

plt.xlabel('time')
plt.ylabel('voltage')
plt.legend()
plt.show()