def test_weights(Simulator, nl):
    name = 'test_weights'
    n1, n2 = 100, 50

    def func(t):
        return np.array([np.sin(4 * t), np.cos(12 * t)])

    transform = np.array([[0.6, -0.4]])

    m = nengo.Model(name, seed=3902)
    u = nengo.Node(output=func)
    a = nengo.Ensemble(nl(n1), dimensions=2, radius=1.5)
    b = nengo.Ensemble(nl(n2), dimensions=1)
    bp = nengo.Probe(b)

    nengo.Connection(u, a)
    nengo.Connection(a, b, transform=transform,
                     weight_solver=nengo.decoders.lstsq_L2nz)

    sim = Simulator(m)
    sim.run(2.)

    t = sim.trange()
    x = func(t).T
    y = np.dot(x, transform.T)
    z = filtfilt(sim.data[bp], 10, axis=0)
    assert allclose(t, y.flatten(), z.flatten(),
                    plotter=Plotter(Simulator, nl),
                    filename='test_connection.' + name + '.pdf',
                    atol=0.1, rtol=0, buf=100, delay=10)

def test_weights(Simulator, nl):
    name = 'test_weights'
    n1, n2 = 100, 50

    def func(t):
        return np.array([np.sin(4 * t), np.cos(12 * t)])

    transform = np.array([[0.6, -0.4]])

    m = nengo.Network(label=name, seed=3902)
    with m:
        m.config[nengo.Ensemble].neuron_type = nl()
        u = nengo.Node(output=func)
        a = nengo.Ensemble(n1, dimensions=2, radius=1.5)
        b = nengo.Ensemble(n2, dimensions=1)
        bp = nengo.Probe(b)

        nengo.Connection(u, a)
        nengo.Connection(a, b, transform=transform,
                         solver=LstsqL2(weights=True))

    sim = Simulator(m)
    sim.run(2.)

    t = sim.trange()
    x = func(t).T
    y = np.dot(x, transform.T)
    z = filtfilt(sim.data[bp], 10, axis=0)
    assert allclose(t, y.flatten(), z.flatten(),
                    plotter=Plotter(Simulator, nl),
                    filename='test_connection.' + name + '.pdf',
                    atol=0.1, rtol=0, buf=100, delay=10)

def test_weights(Simulator, nl, plt, seed):
    n1, n2 = 100, 50

    def func(t):
        return np.array([np.sin(4 * t), np.cos(12 * t)])

    transform = np.array([[0.6, -0.4]])

    m = nengo.Network(label='test_weights', seed=seed)
    with m:
        m.config[nengo.Ensemble].neuron_type = nl()
        u = nengo.Node(output=func)
        a = nengo.Ensemble(n1, dimensions=2, radius=1.5)
        b = nengo.Ensemble(n2, dimensions=1)
        bp = nengo.Probe(b)

        nengo.Connection(u, a)
        nengo.Connection(a, b, transform=transform,
                         solver=LstsqL2(weights=True))

    sim = Simulator(m)
    sim.run(1.)

    t = sim.trange()
    x = func(t).T
    y = np.dot(x, transform.T)
    z = nengo.synapses.filtfilt(sim.data[bp], 0.005, dt=sim.dt)
    assert allclose(t, y, z, atol=0.1, buf=0.1, delay=0.01, plt=plt)

def test_lowpass(Simulator, plt, seed):
    dt = 1e-3
    tau = 0.03

    t, x, yhat = run_synapse(Simulator, seed, Lowpass(tau), dt=dt)
    y = Lowpass(tau).filt(x, dt=dt, y0=0)

    assert allclose(t, y, yhat, delay=dt, plt=plt)

def test_lowpass(Simulator, plt):
    dt = 1e-3
    tau = 0.03

    t, x, yhat = run_synapse(Simulator, nengo.synapses.Lowpass(tau), dt=dt)
    y = filt(x, tau / dt)

    assert allclose(t, y, yhat, delay=dt, plt=plt)

def test_decoders(Simulator, plt, seed):
    dt = 1e-3
    tau = 0.01

    t, x, yhat = run_synapse(Simulator, seed, Lowpass(tau), dt=dt, n_neurons=100)

    y = filt(x, tau, dt=dt)
    assert allclose(t, y, yhat, delay=dt, plt=plt)

def test_alpha(Simulator, plt, seed):
    dt = 1e-3
    tau = 0.03
    num, den = [1], [tau**2, 2*tau, 1]

    t, x, yhat = run_synapse(Simulator, seed, Alpha(tau), dt=dt)
    y = LinearFilter(num, den).filt(x, dt=dt, y0=0)

    assert allclose(t, y, yhat, delay=dt, atol=5e-6, plt=plt)

def test_lowpass(Simulator):
    dt = 1e-3
    tau = 0.03

    t, x, yhat = run_synapse(Simulator, nengo.synapses.Lowpass(tau), dt=dt)
    y = filt(x, tau / dt)

    assert allclose(t, y.flatten(), yhat.flatten(), delay=1,
                    plotter=Plotter(Simulator),
                    filename='test_synapse.test_lowpass.pdf')

def test_decoders(Simulator, nl):
    dt = 1e-3
    tau = 0.01

    t, x, yhat = run_synapse(
        Simulator, nengo.synapses.Lowpass(tau), dt=dt, n_neurons=100)

    y = filt(x, tau / dt)
    assert allclose(t, y.flatten(), yhat.flatten(), delay=1,
                    plotter=Plotter(Simulator, nl),
                    filename='test_synapse.test_decoders.pdf')

def test_scalar(Simulator, nl, plt, seed):
    """A network that represents sin(t)."""
    N = 40
    f = lambda t: np.sin(6.3 * t)

    m = nengo.Network(label='test_scalar', seed=seed)
    with m:
        m.config[nengo.Ensemble].neuron_type = nl()
        input = nengo.Node(output=f)
        A = nengo.Ensemble(N, 1, label='A')
        nengo.Connection(input, A)
        in_p = nengo.Probe(input, 'output')
        A_p = nengo.Probe(A, 'decoded_output', synapse=0.02)

    sim = Simulator(m)
    sim.run(1.0)
    t = sim.trange()
    target = f(t)

    assert allclose(t, target, sim.data[in_p], rtol=1e-3, atol=1e-5)
    assert allclose(t, target, sim.data[A_p], atol=0.1, delay=0.03, plt=plt)

def test_vector(Simulator, nl, plt, seed):
    """A network that represents sin(t), cos(t), cos(t)**2."""
    N = 100
    f = lambda t: [np.sin(6.3*t), np.cos(6.3*t), np.cos(6.3*t)**2]

    m = nengo.Network(label='test_vector', seed=seed)
    with m:
        m.config[nengo.Ensemble].neuron_type = nl()
        input = nengo.Node(output=f)
        A = nengo.Ensemble(N * 3, 3, radius=1.5)
        nengo.Connection(input, A)
        in_p = nengo.Probe(input)
        A_p = nengo.Probe(A, synapse=0.03)

    sim = Simulator(m)
    sim.run(1.0)
    t = sim.trange()
    target = np.vstack(f(t)).T

    assert allclose(t, target, sim.data[in_p], rtol=1e-3, atol=1e-5)
    assert allclose(t, target, sim.data[A_p],
                    plt=plt, atol=0.1, delay=0.03, buf=0.1)

def test_linearfilter(Simulator, plt, seed):
    dt = 1e-3

    # The following num, den are for a 4th order analog Butterworth filter,
    # generated with `scipy.signal.butter(4, 1. / 0.03, analog=True)`
    num = np.array([1234567.90123457])
    den = np.array([1.0, 87.104197658425107, 3793.5706248589954,
                    96782.441842694592, 1234567.9012345686])

    t, x, yhat = run_synapse(Simulator, seed, LinearFilter(num, den), dt=dt)
    y = filt(x, LinearFilter(num, den), dt=dt)

    assert allclose(t, y, yhat, delay=dt, plt=plt)

def test_linearfilter(Simulator, plt, seed):
    dt = 1e-3

    # The following num, den are for a 4th order analog Butterworth filter,
    # generated with `scipy.signal.butter(4, 0.2, analog=False)`
    num = np.array([0.00482434, 0.01929737, 0.02894606, 0.01929737, 0.00482434])
    den = np.array([1.0, -2.36951301, 2.31398841, -1.05466541, 0.18737949])

    synapse = LinearFilter(num, den, analog=False)
    t, x, yhat = run_synapse(Simulator, seed, synapse, dt=dt)
    y = filt(x, synapse, dt=dt)

    assert allclose(t, y, yhat, delay=dt, plt=plt)

def test_alpha(Simulator, plt):
    dt = 1e-3
    tau = 0.03
    b, a = [0.00054336, 0.00053142], [1, -1.9344322, 0.93550699]
    # ^^^ these coefficients found for tau=0.03 and dt=1e-3
    #   scipy.signal.cont2discrete(([1], [tau**2, 2*tau, 1]), dt)

    # b = [0.00054336283526056767, 0.00053142123234546667]
    # a = [1, -1.9344322009640118, 0.93550698503161778]
    # ^^^ these coefficients found by the exact algorithm used in Builder

    t, x, yhat = run_synapse(Simulator, nengo.synapses.Alpha(tau), dt=dt)
    y = lti(x, (b, a))

    assert allclose(t, y, yhat, delay=dt, atol=5e-6, plt=plt)

def test_triangle(Simulator, plt, seed):
    dt = 1e-3
    tau = 0.03

    t, x, ysim = run_synapse(Simulator, seed, Triangle(tau), dt=dt)
    yfilt = Triangle(tau).filt(x, dt=dt, y0=0)

    # compare with convolved filter
    n_taps = int(round(tau / dt)) + 1
    num = np.arange(n_taps, 0, -1, dtype=float)
    num /= num.sum()
    y = np.convolve(x.ravel(), num)[:len(t)]
    y.shape = (-1, 1)

    assert np.allclose(y, yfilt, rtol=0)
    assert allclose(t, y, ysim, delay=dt, rtol=0, plt=plt)

def test_general(Simulator, plt):
    import scipy.signal

    dt = 1e-3
    order = 4
    tau = 0.03

    num, den = scipy.signal.butter(order, 1. / tau, analog=True)
    num, den = num.real, den.real
    numi, deni, dt = scipy.signal.cont2discrete((num, den), dt)

    t, x, yhat = run_synapse(
        Simulator, nengo.synapses.LinearFilter(num, den), dt=dt)
    y = lti(x, (numi, deni))

    assert allclose(t, y, yhat, plt=plt)

def test_solvers(Simulator, nl_nodirect):

    N = 100
    decoder_solvers = [lstsq, lstsq_noise, lstsq_L2, lstsq_L2nz, lstsq_L1]
    weight_solvers = [lstsq_L1, lstsq_drop]

    dt = 1e-3
    tfinal = 4

    def input_function(t):
        return np.interp(t, [1, 3], [-1, 1], left=-1, right=1)

    model = nengo.Network('test_solvers', seed=290)
    with model:
        u = nengo.Node(output=input_function)
        # up = nengo.Probe(u)

        a = nengo.Ensemble(nl_nodirect(N), dimensions=1)
        ap = nengo.Probe(a)
        nengo.Connection(u, a)

        probes = []
        names = []
        for solver, arg in ([(s, 'decoder_solver') for s in decoder_solvers] +
                            [(s, 'weight_solver') for s in weight_solvers]):
            b = nengo.Ensemble(nl_nodirect(N), dimensions=1, seed=99)
            nengo.Connection(a, b, **{arg: solver})
            probes.append(nengo.Probe(b))
            names.append(solver.__name__ + (" (%s)" % arg[0]))

    sim = nengo.Simulator(model, dt=dt)
    sim.run(tfinal)
    t = sim.trange()

    # ref = sim.data[up]
    ref = filtfilt(sim.data[ap], 20)
    outputs = np.array([sim.data[probe] for probe in probes])
    outputs_f = filtfilt(outputs, 20, axis=1)

    close = allclose(t, ref, outputs_f,
                     plotter=Plotter(Simulator, nl_nodirect),
                     filename='test_decoders.test_solvers.pdf',
                     labels=names,
                     atol=0.05, rtol=0, buf=100, delay=7)
    for name, c in zip(names, close):
        assert c, "Solver '%s' does not meet tolerances" % name

def test_compare_solvers(Simulator, nl_nodirect):

    N = 70
    decoder_solvers = [
        Lstsq(), LstsqNoise(), LstsqL2(), LstsqL2nz(), LstsqL1()]
    weight_solvers = [LstsqL1(weights=True), LstsqDrop(weights=True)]

    dt = 1e-3
    tfinal = 4

    def input_function(t):
        return np.interp(t, [1, 3], [-1, 1], left=-1, right=1)

    model = nengo.Network('test_solvers', seed=290)
    with model:
        model.config[nengo.Ensemble].neuron_type = nl_nodirect()
        u = nengo.Node(output=input_function)
        a = nengo.Ensemble(N, dimensions=1)
        nengo.Connection(u, a)
        ap = nengo.Probe(a)

        probes = []
        names = []
        for solver in decoder_solvers + weight_solvers:
            b = nengo.Ensemble(N, dimensions=1, seed=99)
            nengo.Connection(a, b, solver=solver)
            probes.append(nengo.Probe(b))
            names.append("%s(%s)" % (
                solver.__class__.__name__, 'w' if solver.weights else 'd'))

    sim = Simulator(model, dt=dt)
    sim.run(tfinal)
    t = sim.trange()

    # ref = sim.data[up]
    ref = filtfilt(sim.data[ap], 20)
    outputs = np.array([sim.data[probe][:, 0] for probe in probes])
    outputs_f = filtfilt(outputs, 20, axis=1)

    close = allclose(t, ref, outputs_f,
                     plotter=Plotter(Simulator, nl_nodirect),
                     filename='test_solvers.test_solvers.pdf',
                     labels=names,
                     atol=0.05, rtol=0, buf=100, delay=7)
    for name, c in zip(names, close):
        assert c, "Solver '%s' does not meet tolerances" % name

def test_compare_solvers(Simulator, plt, seed):
    pytest.importorskip('sklearn')

    N = 70
    decoder_solvers = [
        Lstsq(), LstsqNoise(), LstsqL2(), LstsqL2nz(), LstsqL1()]
    weight_solvers = [LstsqL1(weights=True), LstsqDrop(weights=True)]

    tfinal = 4

    def input_function(t):
        return np.interp(t, [1, 3], [-1, 1], left=-1, right=1)

    model = nengo.Network(seed=seed)
    with model:
        u = nengo.Node(output=input_function)
        a = nengo.Ensemble(N, dimensions=1)
        nengo.Connection(u, a)
        ap = nengo.Probe(a)

        probes = []
        names = []
        for solver in decoder_solvers + weight_solvers:
            b = nengo.Ensemble(N, dimensions=1, seed=seed + 1)
            nengo.Connection(a, b, solver=solver)
            probes.append(nengo.Probe(b))
            names.append("%s(%s)" % (
                solver.__class__.__name__, 'w' if solver.weights else 'd'))

    sim = Simulator(model)
    sim.run(tfinal)
    t = sim.trange()

    # ref = sim.data[up]
    ref = nengo.synapses.filtfilt(sim.data[ap], 0.02, dt=sim.dt)
    outputs = np.array([sim.data[probe][:, 0] for probe in probes]).T
    outputs_f = nengo.synapses.filtfilt(outputs, 0.02, dt=sim.dt)

    close = allclose(t, ref, outputs_f,
                     atol=0.07, rtol=0, buf=0.1, delay=0.007,
                     plt=plt, labels=names, individual_results=True)

    for name, c in zip(names, close):
        assert c, "Solver '%s' does not meet tolerances" % name

def test_function_points(Simulator, seed, rng, plt):
    x = rng.uniform(-1, 1, size=(1000, 1))
    y = -x

    with nengo.Network(seed=seed) as model:
        u = nengo.Node(nengo.processes.WhiteSignal(1.0, high=10, rms=0.4))
        a = nengo.Ensemble(100, 1)
        v = nengo.Node(size_in=1)
        nengo.Connection(u, a, synapse=None)
        nengo.Connection(a, v, eval_points=x, function=y)

        up = nengo.Probe(u, synapse=nengo.Alpha(0.01))
        vp = nengo.Probe(v, synapse=nengo.Alpha(0.01))

    with Simulator(model, seed=seed) as sim:
        sim.run(1.0)

    assert allclose(sim.trange(), -sim.data[up], sim.data[vp],
                    buf=0.01, delay=0.005, atol=5e-2, rtol=3e-2, plt=plt)