import nengo

model = nengo.Network(label='critter')
with model:
    # velocity input
    velocity = nengo.Node([0, 0], label='velocity')
    
    # motor output
    motor = nengo.Ensemble(200, dimensions=2, label='motor')
    nengo.Probe(motor)

    # make the position memory
    position = nengo.Ensemble(500, dimensions=2, label='position',
                              radius=3)
    nengo.Connection(position, position, synapse=0.1)
    nengo.Connection(motor, position, transform=0.1)
    nengo.Probe(position)

    # figure out which way is home
    home_location = [0.5, 0.5]
    home_dir = nengo.Ensemble(200, dimensions=2, label='home_dir')
    def compute_home(x):
        return (home_location - x) * 10
    nengo.Connection(position, home_dir, function=compute_home)
    nengo.Probe(home_dir)

    # what mode am I in?
    mode = nengo.Node(1, label='mode')

    # should I go where I'm told?
    d_velocity = nengo.Ensemble(300, dimensions=3, label='d_velocity',
                                radius=2)
    nengo.Connection(velocity, d_velocity[[0,1]])
    nengo.Connection(mode, d_velocity[2])
    def velocity_func(x):
        a, b, mode = x
        if mode > 0.5:
            return a, b
        else:
            return 0, 0

    # should I go home?
    nengo.Connection(d_velocity, motor, function=velocity_func)
    d_home = nengo.Ensemble(300, dimensions=3, label='d_home',
                                radius=2)
    nengo.Connection(home_dir, d_home[[0,1]])
    nengo.Connection(mode, d_home[2])
    def home_func(x):
        a, b, mode = x
        if mode < -0.5:
            return a, b
        else:
            return 0, 0
    nengo.Connection(d_home, motor, function=home_func)


import nengo
import nengo.spa as spa

model = nengo.Network(label="Action1")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    Q_A = nengo.Ensemble(50, 1, label='Q_A')
    Q_B = nengo.Ensemble(50, 1, label='Q_B')
    Q_C = nengo.Ensemble(50, 1, label='Q_C')
    Q_D = nengo.Ensemble(50, 1, label='Q_D')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q_A, transform=[vocab.parse('DOG').v])
    nengo.Connection(state, Q_B, transform=[vocab.parse('CAT').v])
    nengo.Connection(state, Q_C, transform=[vocab.parse('RAT').v])
    nengo.Connection(state, Q_D, transform=[vocab.parse('COW').v])
    
    nengo.Probe(state)
    nengo.Probe(Q_A)
    nengo.Probe(Q_B)
    nengo.Probe(Q_C)
    nengo.Probe(Q_D)


import nengo
import nengo.spa as spa

model = nengo.Network(label="Action2")
with model:
    state = nengo.Ensemble(500, 16, label='state')    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])

    nengo.Probe(state)
    nengo.Probe(Q.output)

import nengo
import nengo.spa as spa

model = nengo.Network(label="Action3")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])
    nengo.Probe(state)
    nengo.Probe(Q.output)
    
    Q_together = nengo.Ensemble(200, 4, label='Q_together')
    nengo.Probe(Q_together)
    nengo.Connection(Q.output, Q_together)
    
    def maximum(x):
        result = [0, 0, 0, 0]
        result[x.argmax()] = 1
        return result
    
    R = nengo.networks.EnsembleArray(50, 4, label='R')
    nengo.Connection(Q_together, R.input, function=maximum)
    nengo.Probe(R.output)


import nengo
import nengo.spa as spa

model = nengo.Network(label="Action4")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])
    nengo.Probe(state)
    nengo.Probe(Q.output)

    e = 0.1
    i = -1

    transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
    nengo.Connection(Q.output, Q.input, transform=transform)


import nengo
import nengo.spa as spa

model = nengo.Network(label="Action5")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])
    nengo.Probe(state)
    nengo.Probe(Q.output)

    e = 0.1
    i = -1

    def positive(x):
        if x[0]<0: return [0]
        else: return x
    Q.add_output('positive', positive)
    transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
    nengo.Connection(Q.positive, Q.input, transform=transform)

import nengo
import nengo.spa as spa

model = nengo.Network(label="Action5")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])
    nengo.Probe(state)
    nengo.Probe(Q.output)

    e = 1
    i = -1

    def positive(x):
        if x[0]<0: return [0]
        else: return x
    Q.add_output('positive', positive)
    transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
    nengo.Connection(Q.positive, Q.input, transform=transform)

import nengo
import nengo.spa as spa

model = nengo.Network(label="Action6")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])
    nengo.Probe(state)
    nengo.Probe(Q.output)

    e = 0.5
    i = -1

    def positive(x):
        if x[0]<0: return [0]
        else: return x
    Q.add_output('positive', positive)
    transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
    nengo.Connection(Q.positive, Q.input, transform=transform)

    def threshold(x):
        if x[0]<0: return [0]
        else: return 1
    Q.add_output('threshold', threshold)

    R = nengo.networks.EnsembleArray(50, 4, label='R')
    nengo.Probe(R.output)
    nengo.Connection(Q.threshold, R.input)


import nengo
import nengo.spa as spa

model = nengo.Network(label="Action6")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])
    nengo.Probe(state)
    nengo.Probe(Q.output)

    e = 0.2
    i = -1

    def positive(x):
        if x[0]<0: return [0]
        else: return x
    Q.add_output('positive', positive)
    transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
    nengo.Connection(Q.positive, Q.input, transform=transform)

    def threshold(x):
        if x[0]<0: return [0]
        else: return 1
    Q.add_output('threshold', threshold)

    R = nengo.networks.EnsembleArray(50, 4, label='R')
    nengo.Probe(R.output)
    nengo.Connection(Q.threshold, R.input)


import nengo

model = nengo.Network('Basal Ganglia')
with model:
    mm = 1
    mp = 1
    me = 1
    mg = 1
    ws = 1
    wt = 1
    wm = 1
    wg = 1
    wp = 0.9
    we = 0.3
    e = 0.2
    ep = -0.25
    ee = -0.2
    eg = -0.2
    le = 0.2
    lg = 0.2
    tau_ampa = 0.002
    tau_gaba = 0.008
    N = 50
    D = 4
    model.config[nengo.Ensemble].radius = 1.5
    model.config[nengo.Ensemble].encoders = [[1]]*N
    
    strD1 = nengo.networks.EnsembleArray(N, D, label="StrD1", 
                intercepts=nengo.objects.Uniform(e, 1))
    strD2 = nengo.networks.EnsembleArray(N, D, label="StrD2", 
                intercepts=nengo.objects.Uniform(e, 1))
    stn = nengo.networks.EnsembleArray(N, D, label="STN", 
                intercepts=nengo.objects.Uniform(ep, 1))
    gpi = nengo.networks.EnsembleArray(N, D, label="GPi", 
                intercepts=nengo.objects.Uniform(eg, 1))
    gpe = nengo.networks.EnsembleArray(N, D, label="GPe", 
                intercepts=nengo.objects.Uniform(ee, 1))

    input = nengo.Node([0]*D, label="input")
    output = nengo.Node(label="output", size_in=D)

    # spread the input to StrD1, StrD2, and STN
    nengo.Connection(input, strD1.input, synapse=None,
                     transform=ws * (1 + lg))
    nengo.Connection(input, strD2.input, synapse=None,
                     transform=ws * (1 - le))
    nengo.Connection(input, stn.input, synapse=None,
                     transform=wt)    
                     
    # connect the striatum to the GPi and GPe (inhibitory)
    def func_str(x):
        if x < e:
            return 0
        return mm * (x - e)
    strD1.add_output('func', func_str)
    import numpy
    nengo.Connection(strD1.func,
                     gpi.input, synapse=tau_gaba,
                     transform=-numpy.eye(D) * wm)
    strD2.add_output('func', func_str)
    nengo.Connection(strD2.func,
                     gpe.input, synapse=tau_gaba,
                     transform=-numpy.eye(D) * wm)                     
    def func_stn(x):
        if x < ep:
            return 0
        return mp * (x - ep)
                     
    # connect the STN to GPi and GPe (broad and excitatory)
    tr = wp * numpy.ones((D, D))
    stn.add_output('func', func_stn)
    nengo.Connection(stn.func, gpi.input,
                     transform=tr, synapse=tau_ampa)
    nengo.Connection(stn.func, gpe.input,
                     transform=tr, synapse=tau_ampa)

    def func_gpe(x):
        if x < ee:
            return 0
        return me * (x - ee)

    # connect the GPe to GPi and STN (inhibitory)
    gpe.add_output('func', func_gpe)
    nengo.Connection(gpe.func, gpi.input, synapse=tau_gaba,
                     transform=-we)
    nengo.Connection(gpe.func, stn.input, synapse=tau_gaba,
                     transform=-wg)

    def func_gpi(x):
        if x < eg:
            return 0
        return mg * (x - eg)
    # connect GPi to output (inhibitory)
    gpi.add_output('func', func_gpi)
    nengo.Connection(gpi.func, output)
    
    nengo.Probe(output)

import nengo
import nengo.spa as spa

model = nengo.Network(label="BG1")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])

    nengo.Probe(state)
    nengo.Probe(Q.output)
    
    bg = nengo.networks.BasalGanglia(4)
    nengo.Connection(Q.output, bg.input)
    
    R = nengo.networks.EnsembleArray(50, 4, label='R',
        encoders=[[1]]*50, intercepts=nengo.objects.Uniform(0.2, 1))
    nengo.Connection(bg.output, R.input)
    nengo.Probe(R.output) 
    
    import numpy
    bias = nengo.Node([1], label='bias')
    nengo.Connection(bias, R.input, transform=numpy.ones((4, 1)), synapse=None)

    nengo.Connection(R.output, R.input, transform=(numpy.eye(4)-1), synapse=0.008)


import nengo
import nengo.spa as spa

model = nengo.Network(label="BG4")
with model:
    state = nengo.Ensemble(500, 16, label='state')
    
    Q = nengo.networks.EnsembleArray(50, 4, label='Q')

    vocab = spa.Vocabulary(16)   
    state.vocab = vocab
    nengo.Connection(state, Q.input, transform=[vocab.parse('DOG').v,
                                                vocab.parse('CAT').v,
                                                vocab.parse('RAT').v,
                                                vocab.parse('COW').v])

    nengo.Probe(state)
    nengo.Probe(Q.output)
    
    bg = nengo.networks.BasalGanglia(4)
    nengo.Connection(Q.output, bg.input)
    
    R = nengo.networks.EnsembleArray(50, 4, label='R',
        encoders=[[1]]*50, intercepts=nengo.objects.Uniform(0.2, 1))
    nengo.Connection(bg.output, R.input)
    nengo.Probe(R.output) 
    
    import numpy
    bias = nengo.Node([1], label='bias')
    nengo.Connection(bias, R.input, transform=numpy.ones((4, 1)), synapse=None)

    # mutual inhibition on the actions
    nengo.Connection(R.output, R.input, transform=(numpy.eye(4)-1), synapse=0.008)
    
    motor = nengo.Ensemble(500, 16, label='motor')
    
    nengo.Connection(R.output, motor, 
        transform=numpy.array([vocab.parse('BARK').v,
                    vocab.parse('MEOW').v,
                    vocab.parse('SQUEAK').v,
                    vocab.parse('MOO').v]).T)
    nengo.Probe(motor)


import nengo
import nengo.spa as spa

model = spa.SPA(label="SPA1")
with model:
    model.state = spa.Buffer(16)
    model.motor = spa.Buffer(16)
    actions = spa.Actions(
        'dot(state, DOG) --> motor=BARK',
        'dot(state, CAT) --> motor=MEOW',
        'dot(state, RAT) --> motor=SQUEAK',
        'dot(state, COW) --> motor=MOO',
    )    
    model.bg = spa.BasalGanglia(actions)
    model.thalamus = spa.Thalamus(model.bg)
    
    nengo.Probe(model.state.state.output)
    nengo.Probe(model.motor.state.output)

import nengo
import nengo.spa as spa

model = spa.SPA(label="SPA1")
with model:
    model.state = spa.Buffer(16)
    actions = spa.Actions(
        'dot(state, A) --> state=B',
        'dot(state, B) --> state=C',
        'dot(state, C) --> state=D',
        'dot(state, D) --> state=E',
        'dot(state, E) --> state=A',
    )    
    model.bg = spa.BasalGanglia(actions)
    model.thalamus = spa.Thalamus(model.bg)
    
    nengo.Probe(model.state.state.output)


    nengo.Probe(model.state.state.output)


import nengo
import nengo.spa as spa

model = spa.SPA(label="SPA1")
with model:
    model.state = spa.Buffer(16)
    model.vision = spa.Buffer(16)
    actions = spa.Actions(
        'dot(vision, LETTER) --> state=vision',
        'dot(state, A) --> state=B',
        'dot(state, B) --> state=C',
        'dot(state, C) --> state=D',
        'dot(state, D) --> state=E',
        'dot(state, E) --> state=A',
    )    
    model.bg = spa.BasalGanglia(actions)
    model.thalamus = spa.Thalamus(model.bg)
    
    def my_input(t):
        if t < 0.1:
            return 'LETTER+D'
        else:
            return '0'
    model.input = spa.Input(vision=my_input)
    
    nengo.Probe(model.state.state.output)