def cooccurr_ix(l, th=0.001):
     ix = []
     for i in range(len(l)):
         for j in range(i+1, len(l)):
             if abs(l[i] - l[j]) < th:
                 ix.append((i, j))
     return ix

        def step_conv3(t, x):
            x = x.reshape(shape_in)
            nk, ni, nj = shape_in[-3:]
            f = filters.shape[0]
            sk, si, sj = filters.shape[-3:]
            si2 = (si - 1) / 2
            sj2 = (sj - 1) / 2
            sk2 = (sk - 1) / 2

            y = np.zeros(shape_out)
            for k in range(nk):
                for i in range(ni):
                    for j in range(nj):
                        i0, i1 = i - si2, i + si2 + 1
                        j0, j1 = j - sj2, j + sj2 + 1
                        k0, k1 = k - sk2, k + sk2 + 1
                        sli = slice(max(-i0, 0), min(ni + si - i1, si))
                        slj = slice(max(-j0, 0), min(nj + sj - j1, sj))
                        slk = slice(max(-k0, 0), min(nk + sk - k1, sk))
                        w = (filters[:, i, j, :, sli, slj] if local_filters else
                             filters[:, :,slk, sli, slj])
                        xkij = x[:,max(k0, 0):min(k1, nk), max(i0, 0):min(i1, ni), max(j0, 0):min(j1, nj)]
                        y[:,k, i, j] = np.dot(xkij.ravel(), w.reshape(f, -1).T)

            if biases is not None:
                y += biases

            return y.ravel()

def get_syllables(n_syllables, minfreq, maxfreq, rng=np.random):
    allpaths = []
    for gdir in ['ges-de-ccv', 'ges-de-cv', 'ges-de-cvc', 'ges-de-v']:
        allpaths.extend([ges_path(gdir, gfile)
                         for gfile in os.listdir(ges_path(gdir))])
    indices = rng.permutation(len(allpaths))
    return ([allpaths[indices[i]] for i in range(n_syllables)],
            [rng.uniform(minfreq, maxfreq) for i in range(n_syllables)])

def cd_encoders_biases(n_encoders, trainX, trainY, rng=np.random, mask=None,
                       norm_min=0.05, norm_tries=10):
    """Constrained difference (CD) method for encoders from data [1]_.

    Parameters
    ==========
    n_encoders : int
        Number of encoders to generate.
    trainX : (n_samples, n_dimensions) array-like
        Training features.
    trainY : (n_samples,) array-like
        Training labels.

    Returns
    =======
    encoders : (n_encoders, n_dimensions) array
        Generated encoders.
    biases : (n_encoders,) array
        Generated biases. These are biases assuming `f = G[E * X + b]`,
        and are therefore more like Nengo's `intercepts`.

    References
    ==========
    .. [1] McDonnell, M. D., Tissera, M. D., Vladusich, T., Van Schaik, A.,
       Tapson, J., & Schwenker, F. (2015). Fast, simple and accurate
       handwritten digit classification by training shallow neural network
       classifiers with the "Extreme learning machine" algorithm. PLoS ONE,
       10(8), 1-20. doi:10.1371/journal.pone.0134254
    """
    assert trainX.shape[0] == trainY.size
    trainX = trainX.reshape(trainX.shape[0], -1)
    trainY = trainY.ravel()
    d = trainX.shape[1]
    classes = np.unique(trainY)
    assert mask is None or mask.shape == (n_encoders, d)

    inds = [(trainY == label).nonzero()[0] for label in classes]
    train_norm = npext.norm(trainX, axis=1).mean()

    encoders = np.zeros((n_encoders, d))
    biases = np.zeros(n_encoders)
    for k in range(n_encoders):
        for _ in range(norm_tries):
            i, j = rng.choice(len(classes), size=2, replace=False)
            a, b = trainX[rng.choice(inds[i])], trainX[rng.choice(inds[j])]
            dab = a - b
            if mask is not None:
                dab *= mask[k]
            ndab = npext.norm(dab)**2
            if ndab >= norm_min * train_norm:
                break
        else:
            raise ValueError("Cannot find valid encoder")

        encoders[k] = (2. / ndab) * dab
        biases[k] = np.dot(a + b, dab) / ndab

    return encoders, biases

def test_lif_step(upsample, n_elements):
    """Test the lif nonlinearity, comparing one step with the Numpy version."""
    rng = np.random

    dt = 1e-3
    n_neurons = [12345, 23456, 34567]
    J = RA([rng.normal(scale=1.2, size=n) for n in n_neurons])
    V = RA([rng.uniform(low=0, high=1, size=n) for n in n_neurons])
    W = RA([rng.uniform(low=-5 * dt, high=5 * dt, size=n) for n in n_neurons])
    OS = RA([np.zeros(n) for n in n_neurons])

    ref = 2e-3
    taus = list(rng.uniform(low=15e-3, high=80e-3, size=len(n_neurons)))

    queue = cl.CommandQueue(ctx)
    clJ = CLRA(queue, J)
    clV = CLRA(queue, V)
    clW = CLRA(queue, W)
    clOS = CLRA(queue, OS)
    clTau = CLRA(queue, RA(taus))

    # simulate host
    nls = [LIF(tau_ref=ref, tau_rc=taus[i])
           for i, n in enumerate(n_neurons)]
    for i, nl in enumerate(nls):
        if upsample <= 1:
            nl.step_math(dt, J[i], OS[i], V[i], W[i])
        else:
            s = np.zeros_like(OS[i])
            for j in range(upsample):
                nl.step_math(dt / upsample, J[i], s, V[i], W[i])
                OS[i] = (1./dt) * ((OS[i] > 0) | (s > 0))

    # simulate device
    plan = plan_lif(queue, clJ, clV, clW, clV, clW, clOS, ref, clTau, dt,
                    n_elements=n_elements, upsample=upsample)
    plan()

    if 1:
        a, b = V, clV
        for i in range(len(a)):
            nc, _ = not_close(a[i], b[i]).nonzero()
            if len(nc) > 0:
                j = nc[0]
                print("i", i, "j", j)
                print("J", J[i][j], clJ[i][j])
                print("V", V[i][j], clV[i][j])
                print("W", W[i][j], clW[i][j])
                print("...", len(nc) - 1, "more")

    n_spikes = np.sum([np.sum(os) for os in OS])
    if n_spikes < 1.0:
        logger.warn("LIF spiking mechanism was not tested!")
    assert ra.allclose(J, clJ.to_host())
    assert ra.allclose(V, clV.to_host())
    assert ra.allclose(W, clW.to_host())
    assert ra.allclose(OS, clOS.to_host())

def test_one_short_segment_many_dots(planner):
    for ND in 2, 20, 100:
        check_from_shapes(
            planner,
            0.5, 0.6, 0.7,
            A_shapes=[(10, 1 + ii % 2) for ii in range(ND)],
            X_shapes=[(1 + ii % 2, 1) for ii in range(ND)],
            A_js=[range(ND)],
            X_js=[range(ND)])

def test_one_short_segment_many_longer_dots(planner):
    for ND in 2, 20, 100:
        check_from_shapes(
            planner,
            0.5, 0.6, 0.7,
            A_shapes=[(2000, ii + 1) for ii in range(ND)],
            X_shapes=[(ii + 1, 1) for ii in range(ND)],
            A_js=[range(ND)],
            X_js=[range(ND)])

def rasterplot(time, spikes, ax=None, **kwargs):
    """Generate a raster plot of the provided spike data

    Parameters
    ----------
    time : array
        Time data from the simulation
    spikes: array
        The spike data with columns for each neuron and 1s indicating spikes
    ax: matplotlib.axes.Axes
        The figure axes to plot into.

    Returns
    -------
    ax: matplotlib.axes.Axes
        The axes that were plotted into

    Examples
    --------
    >>> import nengo
    >>> model = nengo.Model("Raster")
    >>> A = nengo.Ensemble(nengo.LIF(20), dimensions=1)
    >>> A_spikes = nengo.Probe(A, "spikes")
    >>> sim = nengo.Simulator(model)
    >>> sim.run(1)
    >>> rasterplot(sim.trange(), sim.data[A_spikes])
    """

    if ax is None:
        ax = plt.gca()

    colors = kwargs.pop('colors', None)
    if colors is None:
        color_cycle = plt.rcParams['axes.color_cycle']
        colors = [color_cycle[ix % len(color_cycle)]
                  for ix in range(spikes.shape[1])]

    if hasattr(ax, 'eventplot'):
        spikes = [time[spikes[:, i] > 0].flatten()
                  for i in range(spikes.shape[1])]
        for ix in range(len(spikes)):
            if spikes[ix].shape == (0,):
                spikes[ix] = np.array([-1])
        ax.eventplot(spikes, colors=colors, **kwargs)
        ax.set_ylim(len(spikes) - 0.5, -0.5)
        if len(spikes) == 1:
            ax.set_ylim(0.4, 1.6)  # eventplot plots different for len==1
        ax.set_xlim(left=0, right=max(time))

    else:
        # Older Matplotlib, doesn't have eventplot
        for i in range(spikes.shape[1]):
            ax.plot(time[spikes[:, i] > 0],
                    np.ones_like(np.where(spikes[:, i] > 0)).T + i, ',',
                    color=colors[i], **kwargs)

    return ax

    def get_convolution_matrix(self):
        """Return the matrix that does a circular convolution by this vector.

        This should be such that ``A*B == dot(A.get_convolution_matrix, B.v)``.
        """
        D = len(self.v)
        T = []
        for i in range(D):
            T.append([self.v[(i - j) % D] for j in range(D)])
        return np.array(T)

def check_from_shapes(
    planner,
    alpha, beta, gamma,
    A_shapes, X_shapes,
    A_js,
    X_js,
):
    rng = np.random.RandomState(1234)
    A = RA([0.1 + rng.rand(*shp) for shp in A_shapes])
    X = RA([0.1 + rng.rand(*shp) for shp in X_shapes])
    Y = RA([0.1 + rng.rand(
        A_shapes[A_js[ii][0]][0],
        X_shapes[X_js[ii][0]][1])
        for ii in range(len(A_js))])
    A_js = RA(A_js)
    X_js = RA(X_js)
    # -- prepare initial conditions on device
    queue = cl.CommandQueue(ctx)
    clA = CLRA(queue, A)
    clX = CLRA(queue, X)
    clY = CLRA(queue, Y)
    clA_js = CLRA(queue, A_js)
    clX_js = CLRA(queue, X_js)
    assert allclose(A, clA)
    assert allclose(X, clX)
    assert allclose(Y, clY)
    assert allclose(A_js, clA_js)
    assert allclose(X_js, clX_js)

    # -- run cl computation
    prog = planner(
        queue, alpha, clA, clA_js, clX, clX_js, beta, clY, gamma=gamma)

    plans = prog.plans
    assert len(plans) == 1
    plans[0]()

    # -- ensure they match
    for i in range(len(A_js)):
        ref = gamma + beta * Y[i] + alpha * sum(
            [np.dot(A[aj], X[xj])
             for aj, xj in zip(A_js[i].ravel(), X_js[i].ravel())])
        sim = clY[i]
        if not np.allclose(ref, sim, atol=1e-3, rtol=1e-3):
            print('A_shapes',  A_shapes)
            print('X_shapes', X_shapes)
            if len(ref) > 20:
                print('ref', ref[:10], '...', ref[-10:])
                print('sim', sim[:10], '...', sim[-10:])
            else:
                print('ref', ref)
                print('sim', sim)
            assert 0

def _mmul_transforms(A_shape, B_shape, C_dim):
    transformA = np.zeros((C_dim, A_shape[0] * A_shape[1]))
    transformB = np.zeros((C_dim, B_shape[0] * B_shape[1]))

    for i in range(A_shape[0]):
        for j in range(A_shape[1]):
            for k in range(B_shape[1]):
                tmp = (j + k * A_shape[1] + i * B_shape[0] * B_shape[1])
                transformA[tmp * 2][j + i * A_shape[1]] = 1
                transformB[tmp * 2 + 1][k + j * B_shape[1]] = 1

    return transformA, transformB

def test_lif_speed(rng, heterogeneous):
    """Test the speed of the lif nonlinearity

    heterogeneous: if true, use a wide range of population sizes.
    """
    dt = 1e-3
    ref = 2e-3
    tau = 20e-3

    n_iters = 1000
    if heterogeneous:
        n_neurons = [1.0e5] * 50 + [1e3] * 500
    else:
        n_neurons = [1.1e5] * 50
    n_neurons = list(map(int, n_neurons))

    J = RA([rng.randn(n) for n in n_neurons], dtype=np.float32)
    V = RA([rng.uniform(low=0, high=1, size=n) for n in n_neurons],
           dtype=np.float32)
    W = RA([rng.uniform(low=-10 * dt, high=10 * dt, size=n)
            for n in n_neurons], dtype=np.float32)
    OS = RA([np.zeros(n) for n in n_neurons], dtype=np.float32)

    queue = cl.CommandQueue(
        ctx, properties=cl.command_queue_properties.PROFILING_ENABLE)

    clJ = CLRA(queue, J)
    clV = CLRA(queue, V)
    clW = CLRA(queue, W)
    clOS = CLRA(queue, OS)

    for i, blockify in enumerate([False, True]):
        plan = plan_lif(queue, dt, clJ, clV, clW, clOS, ref, tau,
                        blockify=blockify)

        with Timer() as timer:
            for j in range(n_iters):
                plan()

        print("plan %d: blockify = %s, dur = %0.3f"
              % (i, blockify, timer.duration))
    print "Original LIF impl"
    for i, blockify in enumerate([False, True]):
        plan = plan_lif_old(queue, dt, clJ, clV, clW, clOS, ref, tau,
                        blockify=blockify)

        with Timer() as timer:
            for j in range(n_iters):
                plan()

        print("plan %d: blockify = %s, dur = %0.3f"
              % (i, blockify, timer.duration))

def _conjgrad_iters(calcAx, b, x, maxiters=None, rtol=1e-6):
    """Solve the single-RHS linear system using conjugate gradient."""

    if maxiters is None:
        maxiters = b.shape[0]

    r = b - calcAx(x)
    p = r.copy()
    rsold = np.dot(r, r)

    for i in range(maxiters):
        Ap = calcAx(p)
        alpha = rsold / np.dot(p, Ap)
        x += alpha * p
        r -= alpha * Ap

        rsnew = np.dot(r, r)
        beta = rsnew / rsold

        if np.sqrt(rsnew) < rtol:
            break

        if beta < 1e-12:  # no perceptible change in p
            break

        # p = r + beta*p
        p *= beta
        p += r
        rsold = rsnew

    return x, i+1

def conjgrad_scipy(A, Y, sigma, tol=1e-4):
    """Solve the least-squares system using Scipy's conjugate gradient."""
    import scipy.sparse.linalg
    Y, m, n, d, matrix_in = _format_system(A, Y)

    damp = m * sigma**2
    calcAA = lambda x: np.dot(A.T, np.dot(A, x)) + damp * x
    G = scipy.sparse.linalg.LinearOperator(
        (n, n), matvec=calcAA, matmat=calcAA, dtype=A.dtype)
    B = np.dot(A.T, Y)

    X = np.zeros((n, d), dtype=B.dtype)
    infos = np.zeros(d, dtype='int')
    itns = np.zeros(d, dtype='int')
    for i in range(d):
        def callback(x):
            itns[i] += 1  # use the callback to count the number of iterations

        X[:, i], infos[i] = scipy.sparse.linalg.cg(
            G, B[:, i], tol=tol, callback=callback)

    info = {'rmses': npext.rms(Y - np.dot(A, X), axis=0),
            'iterations': itns,
            'info': infos}
    return X if matrix_in else X.flatten(), info

def zpk2tf(z, p, k):
    """Return polynomial transfer function representation from zeros
    and poles

    Parameters
    ----------
    z : ndarray
        Zeros of the transfer function.
    p : ndarray
        Poles of the transfer function.
    k : float
        System gain.

    Returns
    -------
    b : ndarray
        Numerator polynomial.
    a : ndarray
        Denominator polynomial.

    """
    z = atleast_1d(z)
    k = atleast_1d(k)
    if len(z.shape) > 1:
        temp = poly(z[0])
        b = zeros((z.shape[0], z.shape[1] + 1), temp.dtype.char)
        if len(k) == 1:
            k = [k[0]] * z.shape[0]
        for i in range(z.shape[0]):
            b[i] = k[i] * poly(z[i])
    else:
        b = k * poly(z)
    a = atleast_1d(poly(p))
    return b, a

def test_init():
    a = SemanticPointer([1, 2, 3, 4])
    assert len(a) == 4

    a = SemanticPointer([1, 2, 3, 4, 5])
    assert len(a) == 5

    a = SemanticPointer(list(range(100)))
    assert len(a) == 100

    a = SemanticPointer(27)
    assert len(a) == 27
    assert np.allclose(a.length(), 1)

    with pytest.raises(Exception):
        a = SemanticPointer(np.zeros(2, 2))

    with pytest.raises(Exception):
        a = SemanticPointer(-1)
    with pytest.raises(Exception):
        a = SemanticPointer(0)
    with pytest.raises(Exception):
        a = SemanticPointer(1.7)
    with pytest.raises(Exception):
        a = SemanticPointer(None)
    with pytest.raises(Exception):
        a = SemanticPointer(int)

    def __call__(self, A, Y, rng=None, E=None):
        tstart = time.time()
        Y, m, n, d, matrix_in = format_system(A, Y)

        # solve for coefficients using standard solver
        X, info0 = self.solver1(A, Y, rng=rng)
        X = self.mul_encoders(X, E)

        # drop weights close to zero, based on `drop` ratio
        Xabs = np.sort(np.abs(X.flat))
        threshold = Xabs[int(np.round(self.drop * Xabs.size))]
        X[np.abs(X) < threshold] = 0

        # retrain nonzero weights
        Y = self.mul_encoders(Y, E)
        for i in range(X.shape[1]):
            nonzero = X[:, i] != 0
            if nonzero.sum() > 0:
                X[nonzero, i], info1 = self.solver2(
                    A[:, nonzero], Y[:, i], rng=rng)

        t = time.time() - tstart
        info = {'rmses': rmses(A, X, Y), 'info0': info0, 'info1': info1,
                'time': t}
        return X if matrix_in else X.flatten(), info

def test_groupby(hashable, force_list, rng):
    if hashable:
        keys = list(range(1, 5))
    else:
        keys = [[0, 0], [0, 1], [1, 0], [1, 1]]

    keys = sorted(keys)

    # make groups and pairs
    groups = [rng.randn(rng.randint(5, 10)) for _ in keys]

    pairs = []
    for key, group in zip(keys, groups):
        pairs.extend((key, value) for value in group)

    # shuffle pairs
    pairs = [pairs[i] for i in rng.permutation(len(pairs))]

    # call groupby
    keygroups = groupby(pairs, lambda p: p[0], force_list=force_list)

    keys2 = sorted(map(lambda x: x[0], keygroups))
    assert keys2 == keys

    for key2, keygroup2 in keygroups:
        group = groups[keys.index(key2)]
        group2 = map(lambda x: x[1], keygroup2)
        assert sorted(group2) == sorted(group)

def test_direct_mode_with_single_neuron(Simulator, plt, seed):
    radius = 2
    dim = 5

    config = nengo.Config(nengo.Ensemble)
    config[nengo.Ensemble].neuron_type = nengo.Direct()
    with config:
        product = nengo.networks.Product(
            1, dim, radius, net=nengo.Network(seed=seed))

    func_A = lambda t: np.sqrt(radius)*np.sin(np.arange(1, dim+1)*2*np.pi*t)
    func_B = lambda t: np.sqrt(radius)*np.sin(np.arange(dim, 0, -1)*2*np.pi*t)
    with product:
        input_A = nengo.Node(func_A)
        input_B = nengo.Node(func_B)
        nengo.Connection(input_A, product.A)
        nengo.Connection(input_B, product.B)
        p = nengo.Probe(product.output, synapse=0.005)

    sim = Simulator(product)
    sim.run(1.0)

    t = sim.trange()
    AB = np.asarray(list(map(func_A, t))) * np.asarray(list(map(func_B, t)))
    delay = 0.013
    offset = np.where(t >= delay)[0]

    for i in range(dim):
        plt.subplot(dim+1, 1, i+1)
        plt.plot(t + delay, AB[:, i])
        plt.plot(t, sim.data[p][:, i])
    plt.xlim(right=t[-1])

    assert rmse(AB[:len(offset), :], sim.data[p][offset, :]) < 0.2

def test_sine_waves(Simulator, nl):
    radius = 2
    dim = 5
    product = nengo.networks.Product(
        200, dim, radius, neuron_type=nl(), seed=63)

    func_A = lambda t: radius*np.sin(np.arange(1, dim+1)*2*np.pi*t)
    func_B = lambda t: radius*np.sin(np.arange(dim, 0, -1)*2*np.pi*t)
    pstc = 0.003
    with product:
        input_A = nengo.Node(func_A)
        input_B = nengo.Node(func_B)
        nengo.Connection(input_A, product.A)
        nengo.Connection(input_B, product.B)
        p = nengo.Probe(product.output, synapse=pstc)

    sim = Simulator(product)
    sim.run(1.0)

    t = sim.trange()
    AB = np.asarray(list(map(func_A, t))) * np.asarray(list(map(func_B, t)))
    delay = 0.011
    offset = np.where(t > delay)[0]

    with Plotter(Simulator) as plt:
        for i in range(dim):
            plt.subplot(dim+1, 1, i+1)
            plt.plot(t + delay, AB[:, i], label="$A \cdot B$")
            plt.plot(t, sim.data[p][:, i], label="Output")
            plt.legend()
        plt.savefig('test_product.test_sine_waves.pdf')
        plt.close()

    assert rmse(AB[:len(offset), :], sim.data[p][offset, :]) < 0.3