def test_inputgatedmemory(Simulator, plt, seed):
    to_memorize = 0.5
    start_memorizing = 0.4
    with nengo.Network(seed=seed) as net:
        test_input = nengo.Node(piecewise(
            {0.0: 0, 0.1: to_memorize, start_memorizing + 0.1: 0}))
        gate_input = nengo.Node(piecewise({0.0: 0, start_memorizing: 1}))
        reset_input = nengo.Node(0)

        mem = nengo.networks.InputGatedMemory(100, 1, difference_gain=5.0)
        nengo.Connection(test_input, mem.input)
        nengo.Connection(gate_input, mem.gate)
        nengo.Connection(reset_input, mem.reset)

        mem_p = nengo.Probe(mem.output, synapse=0.01)

    sim = Simulator(net)
    sim.run(0.5)

    data = sim.data[mem_p]
    t = sim.trange()

    plt.title("gating at %.1f s" % start_memorizing)
    plt.plot(t, data, label="value in memory")
    plt.axhline(to_memorize, c='k', lw=2, label="value to remember")
    plt.axvline(start_memorizing, c='k', ls=':', label="start gating")
    plt.legend(loc='best')

    assert abs(np.mean(data[t < 0.1])) < 0.01
    assert abs(np.mean(data[(t > 0.2) & (t <= 0.4)]) - 0.5) < 0.02
    assert abs(np.mean(data[t > 0.4]) - 0.5) < 0.02

def test_inputgatedmemory(Simulator):
    with nengo.Network(seed=123) as net:
        test_input = nengo.Node(piecewise({0.0: 0, 0.3: 0.5, 1.0: 0}))

        gate_input = nengo.Node(piecewise({0.0: 0, 0.8: 1}))

        reset_input = nengo.Node(0)

        mem = nengo.networks.InputGatedMemory(100, 1, difference_gain=5.0)
        nengo.Connection(test_input, mem.input)

        nengo.Connection(gate_input, mem.gate)

        nengo.Connection(reset_input, mem.reset_node)

        mem_p = nengo.Probe(mem.output, synapse=0.01)

    sim = Simulator(net)
    sim.run(1.2)

    data = sim.data[mem_p]
    trange = sim.trange()

    with Plotter(Simulator) as plt:
        plt.plot(trange, data)

        plt.savefig('test_workingmemory.test_inputgatedmemory.pdf')
        plt.close()

    assert abs(np.mean(data[trange < 0.3])) < 0.01
    assert abs(np.mean(data[(trange > 0.8) & (trange < 1.0)]) - 0.5) < 0.02
    assert abs(np.mean(data[trange > 1.0]) - 0.5) < 0.02

def test_node_to_neurons(Simulator, nl_nodirect):
    name = 'node_to_neurons'
    N = 30

    m = nengo.Model(name, seed=123)
    a = nengo.Ensemble(nl_nodirect(N), dimensions=1)
    inn = nengo.Node(output=np.sin)
    inh = nengo.Node(piecewise({0: 0, 2.5: 1}))
    nengo.Connection(inn, a)
    nengo.Connection(inh, a.neurons, transform=[[-2.5]]*N)

    inn_p = nengo.Probe(inn, 'output')
    a_p = nengo.Probe(a, 'decoded_output', filter=0.1)
    inh_p = nengo.Probe(inh, 'output')

    sim = Simulator(m)
    sim.run(5.0)
    t = sim.trange()
    ideal = np.sin(t)
    ideal[t >= 2.5] = 0

    with Plotter(Simulator, nl_nodirect) as plt:
        plt.plot(t, sim.data(inn_p), label='Input')
        plt.plot(t, sim.data(a_p), label='Neuron approx, filter=0.1')
        plt.plot(t, sim.data(inh_p), label='Inhib signal')
        plt.plot(t, ideal, label='Ideal output')
        plt.legend(loc=0, prop={'size': 10})
        plt.savefig('test_connection.test_' + name + '.pdf')
        plt.close()

    assert np.allclose(sim.data(a_p)[-10:], 0, atol=.1, rtol=.01)

def test_ensemble_to_neurons(Simulator, nl_nodirect, plt, seed):
    N = 30

    m = nengo.Network(seed=seed)
    with m:
        m.config[nengo.Ensemble].neuron_type = nl_nodirect()
        a = nengo.Ensemble(N, dimensions=1)
        b = nengo.Ensemble(N, dimensions=1)
        inn = nengo.Node(output=np.sin)
        inh = nengo.Node(piecewise({0: 0, 0.5: 1}))
        nengo.Connection(inn, a)
        nengo.Connection(inh, b)
        nengo.Connection(b, a.neurons, transform=[[-2.5]] * N)

        inn_p = nengo.Probe(inn, 'output')
        a_p = nengo.Probe(a, 'decoded_output', synapse=0.1)
        b_p = nengo.Probe(b, 'decoded_output', synapse=0.1)
        inh_p = nengo.Probe(inh, 'output')

    sim = Simulator(m)
    sim.run(1.0)
    t = sim.trange()
    ideal = np.sin(t)
    ideal[t >= 0.5] = 0

    plt.plot(t, sim.data[inn_p], label='Input')
    plt.plot(t, sim.data[a_p], label='Neuron approx, pstc=0.1')
    plt.plot(t, sim.data[b_p], label='Neuron approx of inhib sig, pstc=0.1')
    plt.plot(t, sim.data[inh_p], label='Inhib signal')
    plt.plot(t, ideal, label='Ideal output')
    plt.legend(loc=0, prop={'size': 10})

    assert np.allclose(sim.data[a_p][-10:], 0, atol=.1, rtol=.01)
    assert np.allclose(sim.data[b_p][-10:], 1, atol=.1, rtol=.01)

def test_integrator(Simulator, plt, seed):
    model = nengo.Network(seed=seed)
    with model:
        inputs = {0: 0, 0.2: 1, 1: 0, 2: -2, 3: 0, 4: 1, 5: 0}
        input = nengo.Node(piecewise(inputs))

        tau = 0.1
        T = nengo.networks.Integrator(tau, n_neurons=100, dimensions=1)
        nengo.Connection(input, T.input, synapse=tau)

        A = nengo.Ensemble(100, dimensions=1)
        nengo.Connection(A, A, synapse=tau)
        nengo.Connection(input, A, transform=tau, synapse=tau)

        input_p = nengo.Probe(input, sample_every=0.01)
        A_p = nengo.Probe(A, synapse=0.01, sample_every=0.01)
        T_p = nengo.Probe(T.ensemble, synapse=0.01, sample_every=0.01)

    sim = Simulator(model)
    sim.run(6.0)

    t = sim.trange(dt=0.01)
    plt.plot(t, sim.data[A_p], label='Manual')
    plt.plot(t, sim.data[T_p], label='Template')
    plt.plot(t, sim.data[input_p], 'k', label='Input')
    plt.legend(loc='best')

    assert rmse(sim.data[A_p], sim.data[T_p]) < 0.1

def test_oscillator(Simulator, plt, seed):
    model = nengo.Network(seed=seed)
    with model:
        inputs = {0: [1, 0], 0.5: [0, 0]}
        input = nengo.Node(piecewise(inputs), label='Input')

        tau = 0.1
        freq = 5
        T = nengo.networks.Oscillator(tau, freq,  n_neurons=100)
        nengo.Connection(input, T.input)

        A = nengo.Ensemble(100, dimensions=2)
        nengo.Connection(A, A, synapse=tau,
                         transform=[[1, -freq*tau], [freq*tau, 1]])
        nengo.Connection(input, A)

        in_probe = nengo.Probe(input, sample_every=0.01)
        A_probe = nengo.Probe(A, synapse=0.01, sample_every=0.01)
        T_probe = nengo.Probe(T.ensemble, synapse=0.01, sample_every=0.01)

    sim = Simulator(model)
    sim.run(3.0)

    t = sim.trange(dt=0.01)
    plt.plot(t, sim.data[A_probe], label='Manual')
    plt.plot(t, sim.data[T_probe], label='Template')
    plt.plot(t, sim.data[in_probe], 'k', label='Input')
    plt.legend(loc='best')

    assert rmse(sim.data[A_probe], sim.data[T_probe]) < 0.2

def test_integrator(Simulator, nl, plt):
    model = nengo.Network(label='Integrator', seed=892)
    with model:
        model.config[nengo.Ensemble].neuron_type = nl()
        inputs = {0: 0, 0.2: 1, 1: 0, 2: -2, 3: 0, 4: 1, 5: 0}
        input = nengo.Node(piecewise(inputs))

        tau = 0.1
        T = nengo.networks.Integrator(tau, n_neurons=100, dimensions=1)
        nengo.Connection(input, T.input, synapse=tau)

        A = nengo.Ensemble(100, dimensions=1)
        nengo.Connection(A, A, synapse=tau)
        nengo.Connection(input, A, transform=tau, synapse=tau)

        input_p = nengo.Probe(input)
        A_p = nengo.Probe(A, synapse=0.01)
        T_p = nengo.Probe(T.ensemble, synapse=0.01)

    sim = Simulator(model, dt=0.001)
    sim.run(6.0)

    t = sim.trange()
    plt.plot(t, sim.data[A_p], label='Manual')
    plt.plot(t, sim.data[T_p], label='Template')
    plt.plot(t, sim.data[input_p], 'k', label='Input')
    plt.legend(loc=0)

    assert rmse(sim.data[A_p], sim.data[T_p]) < 0.2

def test_node_to_neurons(Simulator, nl_nodirect):
    name = "node_to_neurons"
    N = 30

    m = nengo.Network(label=name, seed=123)
    with m:
        m.config[nengo.Ensemble].neuron_type = nl_nodirect()
        a = nengo.Ensemble(N, dimensions=1)
        inn = nengo.Node(output=np.sin)
        inh = nengo.Node(piecewise({0: 0, 2.5: 1}))
        nengo.Connection(inn, a)
        nengo.Connection(inh, a.neurons, transform=[[-2.5]] * N)

        inn_p = nengo.Probe(inn, "output")
        a_p = nengo.Probe(a, "decoded_output", synapse=0.1)
        inh_p = nengo.Probe(inh, "output")

    sim = Simulator(m)
    sim.run(5.0)
    t = sim.trange()
    ideal = np.sin(t)
    ideal[t >= 2.5] = 0

    with Plotter(Simulator, nl_nodirect) as plt:
        plt.plot(t, sim.data[inn_p], label="Input")
        plt.plot(t, sim.data[a_p], label="Neuron approx, synapse=0.1")
        plt.plot(t, sim.data[inh_p], label="Inhib signal")
        plt.plot(t, ideal, label="Ideal output")
        plt.legend(loc=0, prop={"size": 10})
        plt.savefig("test_connection.test_" + name + ".pdf")
        plt.close()

    assert np.allclose(sim.data[a_p][-10:], 0, atol=0.1, rtol=0.01)

def test_oscillator(Simulator, nl):
    model = nengo.Network(label='Oscillator', seed=789)
    with model:
        inputs = {0: [1, 0], 0.5: [0, 0]}
        input = nengo.Node(piecewise(inputs), label='Input')

        tau = 0.1
        freq = 5
        T = nengo.networks.Oscillator(
            tau, freq, label="Oscillator", neurons=nl(100))
        nengo.Connection(input, T.input)

        A = nengo.Ensemble(nl(100), label='A', dimensions=2)
        nengo.Connection(A, A, synapse=tau,
                         transform=[[1, -freq*tau], [freq*tau, 1]])
        nengo.Connection(input, A)

        in_probe = nengo.Probe(input, "output")
        A_probe = nengo.Probe(A, "decoded_output", synapse=0.01)
        T_probe = nengo.Probe(T.ensemble, "decoded_output", synapse=0.01)

    sim = Simulator(model)
    sim.run(3.0)

    with Plotter(Simulator, nl) as plt:
        t = sim.trange()
        plt.plot(t, sim.data[A_probe], label='Manual')
        plt.plot(t, sim.data[T_probe], label='Template')
        plt.plot(t, sim.data[in_probe], 'k', label='Input')
        plt.legend(loc=0)
        plt.savefig('test_oscillator.test_oscillator.pdf')
        plt.close()

    assert rmse(sim.data[A_probe], sim.data[T_probe]) < 0.3

def test_integrator(Simulator, nl):
    model = nengo.Model('Integrator')
    inputs = {0: 0, 0.2: 1, 1: 0, 2: -2, 3: 0, 4: 1, 5: 0}
    input = nengo.Node(piecewise(inputs))

    tau = 0.1
    T = nengo.networks.Integrator(tau, neurons=nl(100), dimensions=1)
    nengo.Connection(input, T.input, filter=tau)

    A = nengo.Ensemble(nl(100), dimensions=1)
    nengo.Connection(A, A, filter=tau)
    nengo.Connection(input, A, transform=tau, filter=tau)

    input_p = nengo.Probe(input, 'output')
    A_p = nengo.Probe(A, 'decoded_output', filter=0.01)
    T_p = nengo.Probe(T.ensemble, 'decoded_output', filter=0.01)

    sim = Simulator(model, dt=0.001)
    sim.run(6.0)

    with Plotter(Simulator, nl) as plt:
        t = sim.trange()
        plt.plot(t, sim.data(A_p), label='Manual')
        plt.plot(t, sim.data(T_p), label='Template')
        plt.plot(t, sim.data(input_p), 'k', label='Input')
        plt.legend(loc=0)
        plt.savefig('test_integrator.test_integrator.pdf')
        plt.close()

    assert rmse(sim.data(A_p), sim.data(T_p)) < 0.2

def test_function():
    f = piecewise({0: np.sin, 0.5: np.cos})
    assert np.allclose(f(0), [np.sin(0)])
    assert np.allclose(f(0.25), [np.sin(0.25)])
    assert np.allclose(f(0.4999), [np.sin(0.4999)])
    assert np.allclose(f(0.5), [np.cos(0.5)])
    assert np.allclose(f(0.75), [np.cos(0.75)])
    assert np.allclose(f(1.0), [np.cos(1.0)])

def test_basic():
    f = piecewise({0.5: 1, 1.0: 0})
    assert np.allclose(f(-10), [0])
    assert np.allclose(f(0), [0])
    assert np.allclose(f(0.25), [0])
    assert np.allclose(f(0.5), [1])
    assert np.allclose(f(0.75), [1])
    assert np.allclose(f(1.0), [0])
    assert np.allclose(f(1.5), [0])
    assert np.allclose(f(100), [0])

def test_lists():
    f = piecewise({0.5: [1, 0], 1.0: [0, 1]})
    assert np.allclose(f(-10), [0, 0])
    assert np.allclose(f(0), [0, 0])
    assert np.allclose(f(0.25), [0, 0])
    assert np.allclose(f(0.5), [1, 0])
    assert np.allclose(f(0.75), [1, 0])
    assert np.allclose(f(1.0), [0, 1])
    assert np.allclose(f(1.5), [0, 1])
    assert np.allclose(f(100), [0, 1])

def test_function_list():
    def func1(t):
        return t, t ** 2, t ** 3

    def func2(t):
        return t ** 4, t ** 5, t ** 6

    f = piecewise({0: func1, 0.5: func2})
    assert np.allclose(f(0), func1(0))
    assert np.allclose(f(0.25), func1(0.25))
    assert np.allclose(f(0.4999), func1(0.4999))
    assert np.allclose(f(0.5), func2(0.5))
    assert np.allclose(f(0.75), func2(0.75))
    assert np.allclose(f(1.0), func2(1.0))

from nengo.utils.functions import piecewise

__author__ = 'bogdanp'

import nengo
import nengo.spa as spa

model = nengo.Network("Oscillator network")

with model:
    input = nengo.Node(piecewise({0: [1, 0], 0.1: [0, 0]}))
    osc = nengo.networks.Oscillator(1, 1, 100)
    nengo.Connection(input, osc.input)

import nengo

model = nengo.Model('Integrator')
A = nengo.Ensemble(nengo.LIF(100), dimensions=1, label='Integrator')

from nengo.utils.functions import piecewise
input = nengo.Node(piecewise({0: 0, 0.2: 1, 1: 0, 2: -2, 3: 0, 4: 1, 5: 0}), label='Piecewise input')

tau = 0.1
nengo.Connection(A, A, transform=[[1]], filter=tau) # Using a long time constant for stability\n"
nengo.Connection(input, A, transform=[[tau]], filter=tau) # The same time constant as recurrent to make it more 'ideal'"

def test_invalid_function_length():
    with pytest.raises(Exception):
        f = piecewise({0.5: 0, 1.0: lambda t: [t, t ** 2]})
        assert f

def test_invalid_length():
    with pytest.raises(Exception):
        f = piecewise({0.5: [1, 0], 1.0: [1, 0, 0]})
        assert f

def test_invalid_key():
    with pytest.raises(TypeError):
        f = piecewise({0.5: 1, 1: 0, 'a': 0.2})
        assert f

from __future__ import division
from nengo.utils.functions import piecewise

import numpy as np
import nengo 
import matplotlib.pyplot as plt
import raphan_matrices as rmat

# MODEL PARAMETERS
# =====
TAU_RECUR = 0.1
input_func = input_func = piecewise({0: 0, 0.2: 5, 0.3: 0, 0.44: -10, 0.54: 0, 0.8: 5, 0.9: 0})
control_func = piecewise({0: 1, 0.6: 0.5}) # function to define integrator behaviour

model = nengo.Network(label='1dinteg')

with model:
    A = nengo.Ensemble(200, dimensions=2, radius=1.5)

    inp = nengo.Node(input_func)
    control = nengo.Node(output=control_func)

    # wire the stimulus input
    nengo.Connection(inp, A, transform=[[TAU_RECUR], [0]], synapse=TAU_RECUR)
    # wire the control quality input  
    nengo.Connection(control, A[1], synapse=0.005)

    # set up the recurrent input
    nengo.Connection(A, A[0],
                     function = lambda x: x[0] * x[1],
                     synapse=TAU_RECUR)