import numpy
import syde556
reload(syde556)


x, X = syde556.generate_signal(T=1, dt=0.001, rms=0.3, limit=10, seed=3)

N_i = 20
encoders_i = numpy.random.choice([1, -1], N_i)
intercepts_i = numpy.random.uniform(-0.95, 0.95, N_i)
max_rates_i = numpy.random.uniform(250, 300, N_i)
alpha_i, bias_i = syde556.find_alpha_and_bias(intercepts_i, max_rates_i)

N_j = 19
encoders_j = numpy.random.choice([1, -1], N_j)
intercepts_j = numpy.random.uniform(-0.95, 0.95, N_j)
max_rates_j = numpy.random.uniform(250, 300, N_j)
alpha_j, bias_j = syde556.find_alpha_and_bias(intercepts_j, max_rates_j)


A_i = syde556.activity(x, encoders_i, alpha_i, bias_i)
d_i = syde556.decoders(A_i, x, noise=0.1, dx=1.0/len(x))

A_j = syde556.activity(x, encoders_j, alpha_j, bias_j)
d_j = syde556.decoders(A_j, x, noise=0.1, dx=1.0/len(x))

xhat_i = numpy.dot(syde556.add_noise(A_i, 0.1), d_i)
xhat_j = numpy.dot(syde556.add_noise(A_j, 0.1), d_j)

t = numpy.arange(len(x))*0.001
figure(figsize=(8,4))
subplot(1,2,1)
plot(t, x, label='$x$')
plot(t, xhat_i, label='$\hat{x}$')
legend()
xlabel('time (seconds)')
ylabel('value')
ylim(-1.5, 1.0)
title('first population')

subplot(1,2,2)
plot(t, x, label='$y$')
plot(t, xhat_j, label='$\hat{y}$')
legend()
xlabel('time (seconds)')
ylabel('value')
ylim(-1.5, 1.0)
title('second population')

show()

weights = numpy.zeros((N_i, N_j))
for j in range(N_j):
    J_desired = alpha_j[j]*encoders_j[j]*x+bias_j[j]
    weights[:,j] = syde556.decoders(A_i, J_desired, noise=0.1, dx=1.0/len(x))

J_j = numpy.dot(A_i, weights)
A_j = syde556.lif(J_j)

xhat_j = numpy.dot(syde556.add_noise(A_j, 0.1), d_j)

figure(figsize=(8,4))
subplot(1,2,1)
plot(t, x, label='$x$')
plot(t, xhat_i, label='$\hat{x}$')
legend()
xlabel('time (seconds)')
ylabel('value')
ylim(-1.5, 1.0)
title('first population')

subplot(1,2,2)
plot(t, x, label='$y$')
plot(t, xhat_j, label='$\hat{y}$')
legend()
xlabel('time (seconds)')
ylabel('value')
ylim(-1.5, 1.0)
title('second population fed from first')

show()

weights = numpy.outer(d_i, alpha_j*encoders_j)

J_j = numpy.dot(A_i, weights)+bias_j
A_j = syde556.lif(J_j)

xhat_j = numpy.dot(syde556.add_noise(A_j, 0.1), d_j)

figure(figsize=(8,4))
subplot(1,2,1)
plot(t, x, label='$x$')
plot(t, xhat_i, label='$\hat{x}$')
legend()
xlabel('time (seconds)')
ylabel('value')
ylim(-1.5, 1.0)
title('first population')

subplot(1,2,2)
plot(t, x, label='$y$')
plot(t, xhat_j, label='$\hat{y}$')
legend()
xlabel('time (seconds)')
ylabel('value')
ylim(-1.5, 1.0)
title('second population fed from first')

show()

J_j = numpy.outer(numpy.dot(A_i, d_i), alpha_j*encoders_j)+bias_j

import syde556

dt = 0.001
x = numpy.linspace(-1, 1, 1000)
A = syde556.Ensemble(neurons=25, dimensions=1)
B = syde556.Ensemble(neurons=23, dimensions=1)
decoder_A = A.compute_decoder(x, 2*x)
decoder_B = B.compute_decoder(x, x)

x, X = syde556.generate_signal(1.0, dt, rms=0.3, limit=5, seed=2)

spikes_A, x_A = A.simulate_spikes(x, decoder_A, tau=0.01)
spikes_B, x_B = B.simulate_spikes(x_A, decoder_B, tau=0.01)

figure()
t = numpy.arange(len(x))*dt
imshow(spikes_A.T, extent=(0,1.0,-1,1), aspect='auto', interpolation='none', cmap='binary')
plot(t, x_A)
plot(t, x)
title('population A')

figure()
imshow(spikes_B.T, extent=(0,1.0,-1,1), aspect='auto', interpolation='none', cmap='binary')
plot(t, x_B)
plot(t, 2*x)
title('population B')

show()

import syde556

dt = 0.001
x = numpy.linspace(-1, 1, 1000)
A = syde556.Ensemble(neurons=25, dimensions=1)
B = syde556.Ensemble(neurons=23, dimensions=1)
decoder_A = A.compute_decoder(x, x**2)
decoder_B = B.compute_decoder(x, x)

x, X = syde556.generate_signal(1.0, dt, rms=0.3, limit=5, seed=2)

spikes_A, x_A = A.simulate_spikes(x, decoder_A, tau=0.01)
spikes_B, x_B = B.simulate_spikes(x_A, decoder_B, tau=0.01)

figure()
t = numpy.arange(len(x))*dt
imshow(spikes_A.T, extent=(0,1.0,-1,1), aspect='auto', interpolation='none', cmap='binary')
plot(t, x_A)
plot(t, x)
title('population A')

figure()
imshow(spikes_B.T, extent=(0,1.0,-1,1), aspect='auto', interpolation='none', cmap='binary')
plot(t, x_B)
plot(t, x**2)
title('population B')

show()

import syde556

dt = 0.001
x = numpy.linspace(-1, 1, 1000)
A = syde556.Ensemble(neurons=50, dimensions=1)
B = syde556.Ensemble(neurons=50, dimensions=1)
C = syde556.Ensemble(neurons=50, dimensions=1)

decoder_A = A.compute_decoder(x, x, noise=0.2)
decoder_B = B.compute_decoder(x, x, noise=0.2)
decoder_C = C.compute_decoder(x, x, noise=0.2)

x, X = syde556.generate_signal(1.0, dt, rms=0.2, limit=5, seed=2)
y, X = syde556.generate_signal(1.0, dt, rms=0.2, limit=5, seed=6)

spikes_A, xhat = A.simulate_spikes(x, decoder_A, tau=0.01)
spikes_B, yhat = B.simulate_spikes(y, decoder_B, tau=0.01)

spikes_C, zhat = C.simulate_spikes(xhat+yhat, decoder_C, tau=0.01)

figure()
t = numpy.arange(len(x))*dt
plot(t, x, label='$x$')
plot(t, y, label='$y$')
plot(t, x+y, label='$x+y$')
plot(t, zhat, label='$\hat{z}$')
legend(loc='best')

show()

import syde556
reload(syde556)

dt = 0.001

A = syde556.Ensemble(neurons=50, dimensions=2)
B = syde556.Ensemble(neurons=50, dimensions=2)
C = syde556.Ensemble(neurons=50, dimensions=2)

x = numpy.random.normal(0, 0.5, size=(2,1000))
decoder_A = A.compute_decoder(x, x, noise=0.2)

decoder_B = B.compute_decoder(x, x, noise=0.2)
decoder_C = C.compute_decoder(x, x, noise=0.2)

x = numpy.zeros((2,1000))+numpy.array([0.25, 0.1])[:,None]
y = numpy.zeros((2,1000))+numpy.array([0.1, 0.25])[:,None]

#x = numpy.array([syde556.generate_signal(1.0, dt, rms=0.2, limit=2, seed=2)[0],
#                 syde556.generate_signal(1.0, dt, rms=0.2, limit=2, seed=3)[0]])
#y = numpy.array([syde556.generate_signal(1.0, dt, rms=0.2, limit=2, seed=4)[0],
#                 syde556.generate_signal(1.0, dt, rms=0.2, limit=2, seed=5)[0]])

spikes_A, xhat = A.simulate_spikes(x, decoder_A, tau=0.1)
spikes_B, yhat = B.simulate_spikes(y, decoder_B, tau=0.1)

spikes_C, zhat = C.simulate_spikes(xhat+yhat, decoder_C, tau=0.1)

figure()
plot(xhat[0], xhat[1], label='x')
plot(yhat[0], yhat[1], label='y')
plot(zhat[0], zhat[1], label='z')
legend(loc='best')

figure()
title('first dimension')
t = numpy.arange(1000)*dt
plot(t, xhat[0], label='x')
plot(t, yhat[0], label='y')
plot(t, zhat[0], label='z')
legend(loc='best')

figure()
title('second dimension')
t = numpy.arange(1000)*dt
plot(t, xhat[1], label='x')
plot(t, yhat[1], label='y')
plot(t, zhat[1], label='z')

legend(loc='best')


show()