import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import norm

from alphacsc.simulate import get_atoms
from alphacsc.utils import construct_X_multi, _get_D
from alphacsc.learn_d_z_multi import learn_d_z_multi


# Generate synchronous D
n_times_atom, n_times = 20, 500
n_chan = 5
n_atoms = 2
n_trials = 10

v0 = get_atoms('triangle', n_times_atom)  # temporal atoms
v1 = get_atoms('square', n_times_atom)

u0 = np.random.uniform(size=n_chan)  # spatial maps
u1 = np.random.uniform(size=n_chan)

# Build D and scale atoms
atoms = []
for (u, v) in [(u0, v0), (u1, v1)]:
    scale = max(1., norm(np.concatenate([u, v])))
    u /= scale
    v /= scale
    atoms.append(np.outer(u, v))

D = np.array(atoms)
uv = np.array([np.r_[u0, v0], np.r_[u1, v1]])

starts = list()

# add atoms
rng = np.random.RandomState(27)
for _ in range(n_atoms):
    starts.append(rng.randint(low=0, high=n_times - n_times_atom + 1,
                  size=(n_trials,)))
# add activations
Z = np.zeros((n_atoms, n_trials, n_times))
for i in range(n_trials):
    for k_idx, (d, start) in enumerate(zip(D, starts)):
        Z[k_idx, i, starts[k_idx][i]] = rng.uniform()


X = construct_X_multi(Z, D)

pobjs, uv_hats = list(), list()
for random_state in range(7):
    pobj, times, uv_hat, Z_hat = learn_d_z_multi(X, n_atoms, n_times_atom,
                                                 random_state=random_state,
                                                 n_jobs=1, reg=0.01)
    pobjs.append(pobj[-1])
    uv_hats.append(uv_hat)

best_state = np.argmin(pobjs)

plt.close('all')
plt.plot(pobj)

plt.figure("u (Spatial maps)")
plt.plot(uv_hats[best_state][:, :n_chan].T, 'r')
plt.plot(uv[:, :n_chan].T, 'k--')

plt.figure("v (Temporal atoms)")
plt.plot(uv_hats[best_state][:, n_chan:].T, 'r')
plt.plot(uv[:, n_chan:].T, 'k--')

plt.figure("D")
D_hat = _get_D(uv_hats[best_state], n_chan)
for i, d_hat in enumerate(D_hat):
    plt.subplot(2, 1, i + 1)
    plt.plot(d_hat.T, 'r')
    plt.plot(D[i].T, 'k--')

X_hat = construct_X_multi(Z_hat, D_hat)

plt.figure("X")
plt.plot(X[0, 0])
plt.plot(X_hat[0, 0])
plt.show()

@@ -1152,16 +1152,695 @@
 orm()%0A%0A%0A
+fig, axes = plt.subplots(nrows=2, num='atoms', figsize=(10, 8))%0A%0A%0Adef callback(X, uv_hat, Z_hat, reg):%0A    if axes%5B0%5D.lines == %5B%5D:%0A        axes%5B0%5D.plot(uv_hat%5B:, :n_chan%5D.T)%0A        axes%5B1%5D.plot(uv_hat%5B:, n_chan:%5D.T)%0A        axes%5B0%5D.grid(True)%0A        axes%5B1%5D.grid(True)%0A        axes%5B0%5D.set_title('spatial atom')%0A        axes%5B1%5D.set_title('temporal atom')%0A    else:%0A        for line_0, line_1, uv in zip(axes%5B0%5D.lines, axes%5B1%5D.lines, uv_hat):%0A            line_0.set_ydata(uv%5B:n_chan%5D)%0A            line_1.set_ydata(uv%5Bn_chan:%5D)%0A    for ax in axes:%0A        ax.relim()  # make sure all the data fits%0A        ax.autoscale_view(True, True, True)%0A    plt.draw()%0A    plt.pause(0.001)%0A%0A%0A
 X = cons
@@ -2073,16 +2073,84 @@
 _state,%0A
+                                                 callback=callback,%0A
         

add plotting callback