import numpy as np
import pytest

import nengo
from nengo.builder import Builder
from nengo.builder.ensemble import get_activities
from nengo.builder.learning_rules import SimRLS
from nengo.builder.operator import Copy, Reset
from nengo.builder.signal import Signal
from nengo.dists import UniformHypersphere
from nengo.exceptions import BuildError, ValidationError
from nengo.learning_rules import (
    BCM,
    PES,
    RLS,
    LearningRuleType,
    LearningRuleTypeParam,
    Oja,
    Voja,
)
from nengo.processes import WhiteSignal
from nengo.synapses import Alpha, Lowpass
from nengo.utils.numpy import nrmse


def best_weights(weight_data):
    return np.argmax(np.sum(np.var(weight_data, axis=0), axis=0))


def _test_pes(
    Simulator,
    AnyNeuronType,
    plt,
    seed,
    allclose,
    pre_neurons=False,
    post_neurons=False,
    weight_solver=False,
    vin=np.array([0.5, -0.5]),
    vout=None,
    n=200,
    function=None,
    transform=np.array(1.0),
    rate=1e-3,
    post_slice=None,
):
    vout = np.array(vin) if vout is None else vout

    with nengo.Network(seed=seed) as model:
        model.config[nengo.Ensemble].neuron_type = AnyNeuronType()

        stim = nengo.Node(output=vin)
        target = nengo.Node(output=vout)
        pre = nengo.Ensemble(n, dimensions=stim.size_out)
        post = nengo.Ensemble(n, dimensions=stim.size_out)
        error = nengo.Ensemble(n, dimensions=target.size_out)

        nengo.Connection(stim, pre)

        pre = pre.neurons if pre_neurons else pre
        if post_slice is None:
            post_decoded = post
            post = post.neurons if post_neurons else post
        else:
            # if post_neurons and post_slice then we're doing the slicing in the
            # neuron space, so leave the decoded output un-sliced
            post_decoded = post if post_neurons else post[post_slice]
            post = post.neurons if post_neurons else post
            post = post[post_slice]

        conn = nengo.Connection(
            pre,
            post,
            function=function,
            transform=transform,
            learning_rule_type=PES(rate),
        )
        if weight_solver:
            conn.solver = nengo.solvers.LstsqL2(weights=True)

        nengo.Connection(target, error, transform=-1)
        nengo.Connection(post_decoded, error)
        nengo.Connection(error, conn.learning_rule)

        post_p = nengo.Probe(post_decoded, synapse=0.03)
        error_p = nengo.Probe(error, synapse=0.03)

        weights_p = nengo.Probe(conn, "weights", sample_every=0.01)

    with Simulator(model) as sim:
        sim.run(0.5)
    t = sim.trange()
    weights = sim.data[weights_p]

    plt.subplot(211)
    plt.plot(t, sim.data[post_p])
    plt.ylabel("Post decoded value")
    plt.subplot(212)
    plt.plot(t, sim.data[error_p])
    plt.ylabel("Error decoded value")
    plt.xlabel("Time (s)")

    tend = t > 0.4
    assert allclose(sim.data[post_p][tend], vout, atol=0.05)
    assert allclose(sim.data[error_p][tend], 0, atol=0.05)
    assert not allclose(
        weights[0], weights[-1], atol=1e-5, record_rmse=False, print_fail=0
    )


def test_pes_ens_ens(Simulator, NonDirectNeuronType, plt, seed, allclose):
    function = lambda x: [x[1], x[0]]
    _test_pes(Simulator, NonDirectNeuronType, plt, seed, allclose, function=function)


def test_pes_weight_solver(Simulator, plt, seed, allclose):
    function = lambda x: [x[1], x[0]]
    _test_pes(
        Simulator, nengo.LIF, plt, seed, allclose, function=function, weight_solver=True
    )


def test_pes_ens_slice(Simulator, plt, seed, allclose):
    vin = [0.5, -0.5]
    vout = [vin[0] ** 2 + vin[1] ** 2]
    function = lambda x: [x[0] - x[1]]
    _test_pes(
        Simulator,
        nengo.LIF,
        plt,
        seed,
        allclose,
        vin=vin,
        vout=vout,
        function=function,
        post_slice=slice(0, 1),
    )


def test_pes_neuron_neuron(Simulator, plt, seed, rng, allclose):
    n = 200
    initial_weights = rng.uniform(high=4e-4, size=(n, n))
    _test_pes(
        Simulator,
        nengo.LIF,
        plt,
        seed,
        allclose,
        pre_neurons=True,
        post_neurons=True,
        n=n,
        transform=initial_weights,
        rate=7e-4,
    )


def test_pes_neuron_neuron_slice(Simulator, plt, seed, rng, allclose):
    n = 200
    initial_weights = rng.uniform(high=4e-4, size=(n // 2, n))
    _test_pes(
        Simulator,
        nengo.LIF,
        plt,
        seed,
        allclose,
        pre_neurons=True,
        post_neurons=True,
        n=n,
        transform=initial_weights,
        rate=7e-4,
        post_slice=slice(0, n // 2),
    )


def test_pes_neuron_ens(Simulator, plt, seed, rng, allclose):
    n = 200
    initial_weights = rng.uniform(high=1e-4, size=(2, n))
    _test_pes(
        Simulator,
        nengo.LIF,
        plt,
        seed,
        allclose,
        pre_neurons=True,
        post_neurons=False,
        n=n,
        transform=initial_weights,
    )


def test_pes_ens_neurons(Simulator, plt, seed, allclose):
    n = 200
    initial_weights = np.ones((n, 2))
    _test_pes(
        Simulator,
        nengo.LIF,
        plt,
        seed,
        allclose,
        pre_neurons=False,
        post_neurons=True,
        n=n,
        transform=initial_weights,
        rate=1e-4,
    )


def test_pes_transform(Simulator, seed, allclose):
    """Test behaviour of PES when function and transform both defined."""
    n = 200
    # error must be with respect to transformed vector (conn.size_out)
    T = np.asarray([[0.5], [-0.5]])  # transform to output

    m = nengo.Network(seed=seed)
    with m:
        u = nengo.Node(output=[1])
        a = nengo.Ensemble(n, dimensions=1)
        b = nengo.Node(size_in=2)
        e = nengo.Node(size_in=1)

        nengo.Connection(u, a)
        learned_conn = nengo.Connection(
            a,
            b,
            function=lambda x: [0],
            transform=T,
            learning_rule_type=nengo.PES(learning_rate=1e-3),
        )
        assert T.shape[0] == learned_conn.size_out
        assert T.shape[1] == learned_conn.size_mid

        nengo.Connection(b[0], e, synapse=None)
        nengo.Connection(nengo.Node(output=-1), e)
        nengo.Connection(e, learned_conn.learning_rule, transform=T, synapse=None)

        p_b = nengo.Probe(b, synapse=0.05)

    with Simulator(m) as sim:
        sim.run(1.0)
    tend = sim.trange() > 0.7

    assert allclose(sim.data[p_b][tend], [1, -1], atol=1e-2)


def test_pes_multidim_error(Simulator, seed):
    """Test that PES works on error connections mapping from N to 1 dims.

    Note that the transform is applied before the learning rule, so the error
    signal should be 1-dimensional.
    """

    with nengo.Network(seed=seed) as net:
        err = nengo.Node(output=[0])
        ens1 = nengo.Ensemble(20, 3)
        ens2 = nengo.Ensemble(10, 1)

        # Case 1: ens -> ens, weights=False
        conn = nengo.Connection(
            ens1,
            ens2,
            transform=np.ones((1, 3)),
            solver=nengo.solvers.LstsqL2(weights=False),
            learning_rule_type={"pes": nengo.PES()},
        )
        nengo.Connection(err, conn.learning_rule["pes"])
        # Case 2: ens -> ens, weights=True
        conn = nengo.Connection(
            ens1,
            ens2,
            transform=np.ones((1, 3)),
            solver=nengo.solvers.LstsqL2(weights=True),
            learning_rule_type={"pes": nengo.PES()},
        )
        nengo.Connection(err, conn.learning_rule["pes"])
        # Case 3: neurons -> ens
        conn = nengo.Connection(
            ens1.neurons,
            ens2,
            transform=np.ones((1, ens1.n_neurons)),
            learning_rule_type={"pes": nengo.PES()},
        )
        nengo.Connection(err, conn.learning_rule["pes"])

    with Simulator(net) as sim:
        sim.run(0.01)


@pytest.mark.parametrize("pre_synapse", [0, Lowpass(tau=0.05), Alpha(tau=0.005)])
def test_pes_synapse(Simulator, seed, pre_synapse, allclose):
    rule = PES(pre_synapse=pre_synapse)

    with nengo.Network(seed=seed) as model:
        stim = nengo.Node(output=WhiteSignal(0.5, high=10))
        x = nengo.Ensemble(100, 1)

        nengo.Connection(stim, x, synapse=None)
        conn = nengo.Connection(x, x, learning_rule_type=rule)

        p_neurons = nengo.Probe(x.neurons, synapse=pre_synapse)
        p_pes = nengo.Probe(conn.learning_rule, "activities")

    with Simulator(model) as sim:
        sim.run(0.5)

    assert allclose(sim.data[p_neurons][1:, :], sim.data[p_pes][:-1, :])


@pytest.mark.parametrize("weights", [False, True])
def test_pes_recurrent_slice(Simulator, seed, weights, allclose):
    """Test that PES works on recurrent connections from N to 1 dims."""

    with nengo.Network(seed=seed) as net:
        err = nengo.Node(output=[-1])
        stim = nengo.Node(output=[0, 0])
        post = nengo.Ensemble(50, 2, radius=2)
        nengo.Connection(stim, post)

        conn = nengo.Connection(
            post,
            post[1],
            function=lambda x: 0.0,
            solver=nengo.solvers.LstsqL2(weights=weights),
            learning_rule_type=nengo.PES(learning_rate=5e-4),
        )

        nengo.Connection(err, conn.learning_rule)
        p = nengo.Probe(post, synapse=0.025)

    with Simulator(net) as sim:
        sim.run(0.2)

    # Learning rule should drive second dimension high, but not first
    assert allclose(sim.data[p][-10:, 0], 0, atol=0.2)
    assert np.all(sim.data[p][-10:, 1] > 0.8)


def test_pes_cycle(Simulator):
    """Test that PES works when connection output feeds back into error."""

    with nengo.Network() as net:
        a = nengo.Ensemble(10, 1)
        b = nengo.Node(size_in=1)
        c = nengo.Connection(a, b, synapse=None, learning_rule_type=nengo.PES())
        nengo.Connection(b, c.learning_rule, synapse=None)

    with Simulator(net):
        # just checking that this builds without error
        pass


def test_pes_adv_idx(Simulator):
    with nengo.Network() as net:
        pre = nengo.Ensemble(10, 1)
        post = nengo.Ensemble(10, 1)
        nengo.Connection(
            pre.neurons,
            post.neurons[[0, 2, 3]],
            learning_rule_type=nengo.PES(),
            transform=np.ones((3, pre.n_neurons)),
        )

    with pytest.raises(BuildError, match="does not support advanced indexing"):
        Simulator(net)


@pytest.mark.parametrize(
    "pre_neurons,post_neurons,weight_solver",
    [
        (True, True, False),
        (True, False, False),
        (False, True, False),
        (False, False, True),
        (False, False, False),
    ],
)
@pytest.mark.parametrize("pre_slice", (True, False))
@pytest.mark.parametrize("post_slice", (True, False))
def test_pes_pre_post_varieties(
    Simulator, pre_neurons, post_neurons, weight_solver, pre_slice, post_slice
):

    with nengo.Network() as net:
        pre = nengo.Ensemble(10, 12)
        post = nengo.Ensemble(20, 22)
        pre_size = pre.n_neurons if pre_neurons else pre.dimensions
        post_size = post.n_neurons if post_neurons else post.dimensions
        if pre_slice:
            pre_size //= 2
            pre_slice = slice(0, pre_size)
        else:
            pre_slice = slice(None)
        if post_slice:
            post_size //= 2
            post_slice = slice(0, post_size)
        else:
            post_slice = slice(None)

        nengo.Connection(
            (pre.neurons if pre_neurons else pre)[pre_slice],
            (post.neurons if post_neurons else post)[post_slice],
            solver=nengo.solvers.LstsqL2(weights=weight_solver),
            learning_rule_type=nengo.PES(),
            transform=np.ones((post_size, pre_size)),
        )

    apply_encoders = post_neurons or (
        not pre_neurons and not post_neurons and weight_solver
    )

    with Simulator(net) as sim:
        assert (
            any(op.tag == "PES:encode" for op in sim.model.operators) == apply_encoders
        )

        sim.step()


@pytest.mark.parametrize(
    "rule_type, solver",
    [
        (BCM(learning_rate=1e-8), False),
        (Oja(learning_rate=1e-5), False),
        ([Oja(learning_rate=1e-5), BCM(learning_rate=1e-8)], False),
        ([Oja(learning_rate=1e-5), BCM(learning_rate=1e-8)], True),
    ],
)
def test_unsupervised(Simulator, rule_type, solver, seed, rng, plt, allclose):
    n = 200

    m = nengo.Network(seed=seed)
    with m:
        u = nengo.Node(WhiteSignal(0.5, high=10), size_out=2)
        a = nengo.Ensemble(n, dimensions=2)
        b = nengo.Ensemble(n + 1, dimensions=2)
        nengo.Connection(u, a)

        if solver:
            conn = nengo.Connection(a, b, solver=nengo.solvers.LstsqL2(weights=True))
        else:
            initial_weights = rng.uniform(high=1e-3, size=(b.n_neurons, a.n_neurons))
            conn = nengo.Connection(a.neurons, b.neurons, transform=initial_weights)
        conn.learning_rule_type = rule_type

        inp_p = nengo.Probe(u)
        weights_p = nengo.Probe(conn, "weights", sample_every=0.01)

        ap = nengo.Probe(a, synapse=0.03)
        up = nengo.Probe(b, synapse=0.03)

    with Simulator(m, seed=seed + 1) as sim:
        sim.run(0.5)
    t = sim.trange()

    plt.subplot(2, 1, 1)
    plt.plot(t, sim.data[inp_p], label="Input")
    plt.plot(t, sim.data[ap], label="Pre")
    plt.plot(t, sim.data[up], label="Post")
    plt.legend(loc="best", fontsize="x-small")
    plt.subplot(2, 1, 2)
    best_ix = best_weights(sim.data[weights_p])
    plt.plot(sim.trange(sample_every=0.01), sim.data[weights_p][..., best_ix])
    plt.xlabel("Time (s)")
    plt.ylabel("Weights")

    assert not allclose(
        sim.data[weights_p][0], sim.data[weights_p][-1], record_rmse=False, print_fail=0
    )


def learning_net(learning_rule=nengo.PES, net=None, rng=np.random):
    net = nengo.Network() if net is None else net
    with net:
        if learning_rule is nengo.PES:
            learning_rule_type = learning_rule(learning_rate=1e-5)
        else:
            learning_rule_type = learning_rule()

        u = nengo.Node(output=1.0)
        pre = nengo.Ensemble(10, dimensions=1)
        post = nengo.Ensemble(10, dimensions=1)
        initial_weights = rng.uniform(high=1e-3, size=(pre.n_neurons, post.n_neurons))
        conn = nengo.Connection(
            pre.neurons,
            post.neurons,
            transform=initial_weights,
            learning_rule_type=learning_rule_type,
        )
        if learning_rule is nengo.PES:
            err = nengo.Ensemble(10, dimensions=1)
            nengo.Connection(u, err)
            nengo.Connection(err, conn.learning_rule)

        net.activity_p = nengo.Probe(pre.neurons, synapse=0.01)
        net.weights_p = nengo.Probe(conn, "weights", synapse=None, sample_every=0.01)
    return net


@pytest.mark.parametrize("learning_rule", [nengo.PES, nengo.BCM, nengo.Oja])
def test_dt_dependence(Simulator, plt, learning_rule, seed, rng, allclose):
    """Learning rules should work the same regardless of dt."""
    m = learning_net(learning_rule, nengo.Network(seed=seed), rng)

    trans_data = []
    # Using dts greater near tau_ref (0.002 by default) causes learning to
    # differ due to lowered presynaptic firing rate
    dts = (0.0001, 0.001)
    colors = ("b", "g", "r")
    ax1 = plt.subplot(2, 1, 1)
    ax2 = plt.subplot(2, 1, 2)
    for c, dt in zip(colors, dts):
        with Simulator(m, dt=dt) as sim:
            sim.run(0.1)
        trans_data.append(sim.data[m.weights_p])
        best_ix = best_weights(sim.data[m.weights_p])
        ax1.plot(
            sim.trange(sample_every=0.01), sim.data[m.weights_p][..., best_ix], c=c
        )
        ax2.plot(sim.trange(), sim.data[m.activity_p], c=c)

    ax1.set_xlim(right=sim.trange()[-1])
    ax1.set_ylabel("Connection weight")
    ax2.set_xlim(right=sim.trange()[-1])
    ax2.set_ylabel("Presynaptic activity")

    assert allclose(trans_data[0], trans_data[1], atol=3e-3)
    assert not allclose(
        sim.data[m.weights_p][0],
        sim.data[m.weights_p][-1],
        record_rmse=False,
        print_fail=0,
    )


@pytest.mark.parametrize("learning_rule", [nengo.PES, nengo.BCM, nengo.Oja])
def test_reset(Simulator, learning_rule, plt, seed, rng, allclose):
    """Make sure resetting learning rules resets all state."""
    m = learning_net(learning_rule, nengo.Network(seed=seed), rng)

    with Simulator(m) as sim:
        sim.run(0.1)
        sim.run(0.2)

        first_t = sim.trange()
        first_t_trans = sim.trange(sample_every=0.01)
        first_activity_p = np.array(sim.data[m.activity_p], copy=True)
        first_weights_p = np.array(sim.data[m.weights_p], copy=True)

        sim.reset()
        sim.run(0.3)

    plt.subplot(2, 1, 1)
    plt.ylabel("Neural activity")
    plt.plot(first_t, first_activity_p, c="b")
    plt.plot(sim.trange(), sim.data[m.activity_p], c="g")
    plt.subplot(2, 1, 2)
    plt.ylabel("Connection weight")
    best_ix = best_weights(first_weights_p)
    plt.plot(first_t_trans, first_weights_p[..., best_ix], c="b")
    plt.plot(sim.trange(sample_every=0.01), sim.data[m.weights_p][..., best_ix], c="g")

    assert allclose(sim.trange(), first_t)
    assert allclose(sim.trange(sample_every=0.01), first_t_trans)
    assert allclose(sim.data[m.activity_p], first_activity_p)
    assert allclose(sim.data[m.weights_p], first_weights_p)


def test_learningruletypeparam():
    """LearningRuleTypeParam must be one or many learning rules."""

    class Test:
        lrp = LearningRuleTypeParam("lrp", default=None)

    inst = Test()
    assert inst.lrp is None
    inst.lrp = Oja()
    assert isinstance(inst.lrp, Oja)
    inst.lrp = [Oja(), Oja()]
    for lr in inst.lrp:
        assert isinstance(lr, Oja)
    # Non-LR no good
    with pytest.raises(ValueError):
        inst.lrp = "a"
    # All elements in list must be LR
    with pytest.raises(ValueError):
        inst.lrp = [Oja(), "a", Oja()]


def test_learningrule_attr(seed):
    """Test learning_rule attribute on Connection"""

    def check_rule(rule, conn, rule_type):
        assert rule.connection is conn and rule.learning_rule_type is rule_type

    with nengo.Network(seed=seed):
        a, b = [nengo.Ensemble(10, 2) for i in range(2)]
        T = np.ones((10, 10))

        r1 = PES()
        c1 = nengo.Connection(a.neurons, b.neurons, learning_rule_type=r1, transform=1)
        check_rule(c1.learning_rule, c1, r1)

        r2 = [PES(), BCM()]
        c2 = nengo.Connection(a.neurons, b.neurons, learning_rule_type=r2, transform=T)
        assert isinstance(c2.learning_rule, list)
        for rule, rule_type in zip(c2.learning_rule, r2):
            check_rule(rule, c2, rule_type)

        r3 = dict(oja=Oja(), bcm=BCM())
        c3 = nengo.Connection(a.neurons, b.neurons, learning_rule_type=r3, transform=T)
        assert isinstance(c3.learning_rule, dict)
        assert set(c3.learning_rule) == set(r3)  # assert same keys
        for key, value in r3.items():
            check_rule(c3.learning_rule[key], c3, value)


def test_voja_encoders(Simulator, PositiveNeuronType, rng, seed, allclose):
    """Tests that voja changes active encoders to the input."""
    n = 200
    learned_vector = np.asarray([0.3, -0.4, 0.6])
    learned_vector /= np.linalg.norm(learned_vector)
    n_change = n // 2  # modify first half of the encoders

    # Set the first half to always fire with random encoders, and the
    # remainder to never fire due to their encoder's dot product with the input
    intercepts = np.asarray([-1] * n_change + [0.99] * (n - n_change))
    rand_encoders = UniformHypersphere(surface=True).sample(
        n_change, len(learned_vector), rng=rng
    )
    encoders = np.append(rand_encoders, [-learned_vector] * (n - n_change), axis=0)

    m = nengo.Network(seed=seed)
    with m:
        m.config[nengo.Ensemble].neuron_type = PositiveNeuronType()
        u = nengo.Node(output=learned_vector)
        x = nengo.Ensemble(
            n,
            dimensions=len(learned_vector),
            intercepts=intercepts,
            encoders=encoders,
            max_rates=nengo.dists.Uniform(300.0, 400.0),
            radius=2.0,
        )  # to test encoder scaling

        conn = nengo.Connection(
            u, x, synapse=None, learning_rule_type=Voja(learning_rate=1e-1)
        )
        p_enc = nengo.Probe(conn.learning_rule, "scaled_encoders")
        p_enc_ens = nengo.Probe(x, "scaled_encoders")

    with Simulator(m) as sim:
        sim.run(1.0)
    t = sim.trange()
    tend = t > 0.5

    # Voja's rule relies on knowing exactly how the encoders were scaled
    # during the build process, because it modifies the scaled_encoders signal
    # proportional to this factor. Therefore, we should check that its
    # assumption actually holds.
    encoder_scale = (sim.data[x].gain / x.radius)[:, np.newaxis]
    assert allclose(sim.data[x].encoders, sim.data[x].scaled_encoders / encoder_scale)

    # Check that the last half kept the same encoders throughout the simulation
    assert allclose(sim.data[p_enc][0, n_change:], sim.data[p_enc][:, n_change:])
    # and that they are also equal to their originally assigned value
    assert allclose(
        sim.data[p_enc][0, n_change:] / encoder_scale[n_change:], -learned_vector
    )

    # Check that the first half converged to the input
    assert allclose(
        sim.data[p_enc][tend, :n_change] / encoder_scale[:n_change],
        learned_vector,
        atol=0.01,
    )
    # Check that encoders probed from ensemble equal encoders probed from Voja
    assert allclose(sim.data[p_enc], sim.data[p_enc_ens])


def test_voja_modulate(Simulator, NonDirectNeuronType, seed, allclose):
    """Tests that voja's rule can be modulated on/off."""
    n = 200
    learned_vector = np.asarray([0.5])

    def control_signal(t):
        """Modulates the learning on/off."""
        return 0 if t < 0.5 else -1

    m = nengo.Network(seed=seed)
    with m:
        m.config[nengo.Ensemble].neuron_type = NonDirectNeuronType()
        control = nengo.Node(output=control_signal)
        u = nengo.Node(output=learned_vector)
        x = nengo.Ensemble(n, dimensions=len(learned_vector))

        conn = nengo.Connection(
            u, x, synapse=None, learning_rule_type=Voja(post_synapse=None)
        )
        nengo.Connection(control, conn.learning_rule, synapse=None)

        p_enc = nengo.Probe(conn.learning_rule, "scaled_encoders")

    with Simulator(m) as sim:
        sim.run(1.0)
    tend = sim.trange() > 0.5

    # Check that encoders stop changing after 0.5s
    assert allclose(sim.data[p_enc][tend], sim.data[p_enc][-1])

    # Check that encoders changed during first 0.5s
    i = np.where(tend)[0][0]  # first time point after changeover
    assert not allclose(
        sim.data[p_enc][0], sim.data[p_enc][i], record_rmse=False, print_fail=0
    )


def _test_rls_network(
    Simulator,
    seed,
    plt,
    tols,
    dims=1,
    lrate=0.01,
    neuron_type=nengo.LIFRate(),
    tau=None,
    t_train=0.5,
    t_test=0.25,
):
    # Input is a scalar sinusoid with given frequency
    n_neurons = 100
    freq = 5

    # Learn a linear transformation within t_train seconds
    transform = np.random.RandomState(seed=seed).randn(dims, 1)
    lr = RLS(learning_rate=lrate, pre_synapse=tau)

    with nengo.Network(seed=seed) as model:
        u = nengo.Node(output=lambda t: np.sin(freq * 2 * np.pi * t))
        x = nengo.Ensemble(n_neurons, 1, neuron_type=neuron_type)
        y = nengo.Node(size_in=dims)
        y_on = nengo.Node(size_in=dims)
        y_off = nengo.Node(size_in=dims)

        e = nengo.Node(
            size_in=dims, output=lambda t, e: e if t < t_train else np.zeros_like(e)
        )

        nengo.Connection(u, y, synapse=None, transform=transform)
        nengo.Connection(u, x, synapse=None)
        conn_on = nengo.Connection(
            x,
            y_on,
            synapse=None,
            learning_rule_type=lr,
            function=lambda _: np.zeros(dims),
        )
        nengo.Connection(y, e, synapse=None, transform=-1)
        nengo.Connection(y_on, e, synapse=None)
        nengo.Connection(e, conn_on.learning_rule, synapse=tau)

        nengo.Connection(x, y_off, synapse=None, transform=transform)

        p_y = nengo.Probe(y, synapse=tau)
        p_y_on = nengo.Probe(y_on, synapse=tau)
        p_y_off = nengo.Probe(y_off, synapse=tau)
        p_inv_gamma = nengo.Probe(conn_on.learning_rule, "inv_gamma")

    with Simulator(model) as sim:
        sim.run(t_train + t_test)

    plt.plot(sim.trange(), sim.data[p_y_off], "k")
    plt.plot(sim.trange(), sim.data[p_y_on])

    # Check _descstr
    ops = [op for op in sim.model.operators if isinstance(op, SimRLS)]
    assert len(ops) == 1
    assert str(ops[0]).startswith("SimRLS")

    test = sim.trange() >= t_train

    on_versus_off = nrmse(sim.data[p_y_on][test], sim.data[p_y_off][test])
    on_versus_ideal = nrmse(sim.data[p_y_on][test], sim.data[p_y][test])
    off_versus_ideal = nrmse(sim.data[p_y_off][test], sim.data[p_y][test])

    A = get_activities(sim.data[x], x, np.linspace(-1, 1, 1000)[:, None])
    gamma_off = A.T.dot(A) + np.eye(n_neurons) / lr.learning_rate
    gamma_on = np.linalg.inv(sim.data[p_inv_gamma][-1])

    gamma_off /= np.linalg.norm(gamma_off)
    gamma_on /= np.linalg.norm(gamma_on)
    gamma_diff = nrmse(gamma_on, gamma_off)

    print()
    print(on_versus_off, on_versus_ideal, off_versus_ideal, gamma_diff)
    print()
    assert on_versus_off < tols[0]
    assert on_versus_ideal < tols[1]
    assert off_versus_ideal < tols[2]
    assert gamma_diff < tols[3]


def test_rls_scalar_rate(Simulator, seed, plt):
    # use artificially high rate to ensure we can make error arbitrarily small
    _test_rls_network(
        Simulator,
        seed,
        plt,
        lrate=100,
        tols=[0.02, 3e-3, 0.02, 0.3],
    )


def test_rls_multidim(Simulator, seed, plt):
    # use artificially high rate to ensure we can make error arbitrarily small
    _test_rls_network(
        Simulator,
        seed,
        plt,
        dims=11,
        lrate=100,
        tols=[0.02, 3e-3, 0.02, 0.3],
    )


def test_rls_scalar_spiking(Simulator, seed, plt):
    _test_rls_network(
        Simulator,
        seed,
        plt,
        neuron_type=nengo.LIF(),
        lrate=0.05,
        tau=0.01,
        tols=[0.04, 0.05, 0.04, 0.3],
    )


def test_frozen():
    """Test attributes inherited from FrozenObject"""
    a = PES(learning_rate=2e-3, pre_synapse=4e-3)
    b = PES(learning_rate=2e-3, pre_synapse=4e-3)
    c = PES(learning_rate=2e-3, pre_synapse=5e-3)

    assert hash(a) == hash(a)
    assert hash(b) == hash(b)
    assert hash(c) == hash(c)

    assert a == b
    assert hash(a) == hash(b)
    assert a != c
    assert hash(a) != hash(c)  # not guaranteed, but highly likely
    assert b != c
    assert hash(b) != hash(c)  # not guaranteed, but highly likely

    with pytest.raises((ValueError, RuntimeError)):
        a.learning_rate = 1e-1


def test_pes_direct_errors():
    """Test that applying a learning rule to a direct ensemble errors."""
    with nengo.Network():
        pre = nengo.Ensemble(10, 1, neuron_type=nengo.Direct())
        post = nengo.Ensemble(10, 1)
        conn = nengo.Connection(pre, post)
        with pytest.raises(ValidationError):
            conn.learning_rule_type = nengo.PES()


def test_custom_type(Simulator, allclose):
    """Test with custom learning rule type.

    A custom learning type may have ``size_in`` not equal to 0, 1, or None.
    """

    class TestRule(nengo.learning_rules.LearningRuleType):
        modifies = "decoders"

        def __init__(self):
            super().__init__(1.0, size_in=3)

    def build_test_rule(model, _, rule):
        error = Signal(np.zeros(rule.connection.size_in))
        model.add_op(Reset(error))
        model.sig[rule]["in"] = error[: rule.size_in]

        model.add_op(Copy(error, model.sig[rule]["delta"]))

    Builder.register(TestRule)(build_test_rule)

    with nengo.Network() as net:
        a = nengo.Ensemble(10, 1)
        b = nengo.Ensemble(10, 1)
        conn = nengo.Connection(
            a.neurons, b, transform=np.zeros((1, 10)), learning_rule_type=TestRule()
        )

        err = nengo.Node([1, 2, 3])
        nengo.Connection(err, conn.learning_rule, synapse=None)

        p = nengo.Probe(conn, "weights")

    with Simulator(net) as sim:
        sim.run(sim.dt * 5)

    assert allclose(sim.data[p][:, 0, :3], np.outer(np.arange(1, 6), np.arange(1, 4)))
    assert allclose(sim.data[p][:, :, 3:], 0)


def test_slicing(Simulator, seed, allclose):
    with nengo.Network(seed=seed) as model:
        a = nengo.Ensemble(50, 1)
        b = nengo.Ensemble(30, 2)
        conn = nengo.Connection(
            a, b, learning_rule_type=PES(), function=lambda x: (0, 0)
        )
        nengo.Connection(nengo.Node(1.0), a)

        err1 = nengo.Node(lambda t, x: x - 0.75, size_in=1)
        nengo.Connection(b[0], err1)
        nengo.Connection(err1, conn.learning_rule[0])

        err2 = nengo.Node(lambda t, x: x + 0.5, size_in=1)
        nengo.Connection(b[1], err2)
        nengo.Connection(err2, conn.learning_rule[1])

        p = nengo.Probe(b, synapse=0.03)

    with Simulator(model) as sim:
        sim.run(1.0)

    t = sim.trange() > 0.8
    assert allclose(sim.data[p][t, 0], 0.75, atol=0.15)
    assert allclose(sim.data[p][t, 1], -0.5, atol=0.15)


def test_null_error():
    with nengo.Network():
        a = nengo.Ensemble(1, 1)
        b = nengo.Ensemble(1, 1)

        # works with a decoded connection (since we'll be generating weights as
        # part of the decoding process)
        nengo.Connection(a, b, learning_rule_type=nengo.PES(), transform=None)

        # error on neuron connection for decoder learning rule
        with pytest.raises(ValidationError, match="does not have weights"):
            nengo.Connection(
                a.neurons, b, learning_rule_type=nengo.PES(), transform=None
            )

        # works for decoded connection with solver.weights=True
        nengo.Connection(
            a,
            b,
            solver=nengo.solvers.LstsqL2(weights=True),
            learning_rule_type=nengo.BCM(),
            transform=None,
        )

        # error on neuron connection for weights learning rule
        with pytest.raises(ValidationError, match="does not have weights"):
            nengo.Connection(
                a.neurons, b.neurons, learning_rule_type=nengo.BCM(), transform=None
            )

        # works with encoder learning rules (since they don't require a transform)
        nengo.Connection(a.neurons, b, learning_rule_type=Voja(), transform=None)


def test_high_learning_rate_warning():
    with pytest.warns(UserWarning, match="learning rate is very high"):
        nengo.PES(learning_rate=1e32)


def test_learning_post_error():
    with nengo.Network():
        ens = nengo.Ensemble(10, 1)
        conn = nengo.Connection(ens, ens, learning_rule_type=nengo.PES())

        with pytest.raises(ValidationError, match="'post' must.*'Ensemble', 'Neurons"):
            nengo.Connection(ens, conn.learning_rule, learning_rule_type=nengo.PES())


def test_encoder_learning_post_errors():
    with nengo.Network():
        ens = nengo.Ensemble(2, 2)

        with pytest.raises(ValidationError, match="'post' must be of type 'Ensemble'"):
            nengo.Connection(
                ens,
                nengo.Node(size_in=2),
                learning_rule_type=nengo.Voja(),
            )

        with pytest.raises(ValidationError, match="encoders are not used"):
            nengo.Connection(
                ens,
                ens,
                solver=nengo.solvers.LstsqL2(weights=True),
                learning_rule_type=nengo.Voja(),
            )


def test_bad_learning_rule_modifies(Simulator):
    class BadCustomRule(nengo.learning_rules.LearningRuleType):
        modifies = "badval"  # start with a bad value to hit API check

    class TrickCustomRule(nengo.learning_rules.LearningRuleType):
        # start with a valid value, then switch once we pass API check
        modifies = "encoders"

    with nengo.Network() as net:
        ens = nengo.Ensemble(2, 2)

        with pytest.raises(ValidationError, match="Unrecognized target 'badval'"):
            nengo.Connection(ens, ens, learning_rule_type=BadCustomRule())

        nengo.Connection(ens, ens, learning_rule_type=TrickCustomRule())

    TrickCustomRule.modifies = "badvalue"  # switch to invalid value
    with pytest.raises(BuildError, match="Unknown target 'badvalue'"):
        with Simulator(net):
            pass


def test_learning_rule_size_in_strings():
    class CustomLR(LearningRuleType):
        modifies = "decoders"

    with nengo.Network():
        a = nengo.Ensemble(10, 5)
        b = nengo.Ensemble(10, 2)

        ref = {
            "pre": 5,
            "mid": 3,
            "post": 2,
            "pre_state": 5,
            "post_state": 2,
        }
        for size_in, correct_size_in in ref.items():
            conn = nengo.Connection(
                a,
                b,
                function=lambda x: x[:3],
                transform=np.ones((2, 3)),
                learning_rule_type=CustomLR(size_in=size_in),
            )
            assert conn.learning_rule.size_in == correct_size_in

    with pytest.raises(ValidationError, match="is not a valid string value"):
        LearningRuleType(size_in="badval")


def test_bad_weight_learning_rule_transform_shape():
    with nengo.Network():
        ens = nengo.Ensemble(5, 1)

        with pytest.raises(ValidationError, match="Transform.*post_neurons x pre_neur"):
            nengo.Connection(
                ens,
                ens.neurons,
                transform=np.ones((5, 1)),
                learning_rule_type=nengo.BCM(),
            )


def test_probeable():
    net = nengo.Network()

    def check_learning_rule(learning_rule_type, expected, net=net):
        assert learning_rule_type.probeable == expected
        post = net.e if isinstance(learning_rule_type, Voja) else net.n
        transform = np.ones((1, 10)) if isinstance(learning_rule_type, Voja) else 1.0
        conn = nengo.Connection(
            net.n, post, transform=transform, learning_rule_type=learning_rule_type
        )
        assert conn.learning_rule.probeable == expected

    with net:
        net.e = nengo.Ensemble(10, 1)
        net.n = net.e.neurons
        check_learning_rule(nengo.PES(), ("error", "activities", "delta"))
        check_learning_rule(
            nengo.RLS(), ("pre_filtered", "error", "delta", "inv_gamma")
        )
        check_learning_rule(
            nengo.BCM(), ("theta", "pre_filtered", "post_filtered", "delta")
        )
        check_learning_rule(nengo.Oja(), ("pre_filtered", "post_filtered", "delta"))
        check_learning_rule(nengo.Voja(), ("post_filtered", "scaled_encoders", "delta"))
