import nef

net = nef.Network('Creature')

net.make_input('command_input', [0,0])

net.make('command', neurons=100, dimensions=2)

net.make('motor', neurons=100, dimensions=2)

net.make('position', neurons=1000, dimensions=2, radius=5)

net.make('scared_direction', neurons=100, dimensions=2)

def negative(x):
    return -x[0], -x[1]

net.connect('position', 'scared_direction', func=negative)

net.connect('position', 'position')

net.make('plan', neurons=500, dimensions=5)
net.connect('command', 'plan', index_post=[0,1])
net.connect('scared_direction', 'plan', index_post=[2,3])
net.make('scared', neurons=50, dimensions=1)
net.make_input('scared_input', [0])
net.connect('scared_input', 'scared')
net.connect('scared', 'plan', index_post=[4])

def plan_function(x):
    c_x, c_y, s_x, s_y, s = x
    return s*(s_x)+(1-s)*c_x, s*(s_y)+(1-s)*c_y
    
net.connect('plan', 'motor', func=plan_function)    


def rescale(x):
    return x[0]*0.1, x[1]*0.1

net.connect('motor', 'position', func=rescale)


net.connect('command_input', 'command')

net.add_to_nengo()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make('Q_A', neurons=50, dimensions=1)
net.make('Q_B', neurons=50, dimensions=1)
net.make('Q_C', neurons=50, dimensions=1)
net.make('Q_D', neurons=50, dimensions=1)
net.connect('s', 'Q_A', transform=[1,0])
net.connect('s', 'Q_B', transform=[-1,0])
net.connect('s', 'Q_C', transform=[0,1])
net.connect('s', 'Q_D', transform=[0,-1])

net.make_input('input', [0,0])
net.connect('input', 's')

net.add_to_nengo()
net.view()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

net.make_input('input', [0,0])
net.connect('input', 's')

net.add_to_nengo()
net.view()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.make('Qall', neurons=200, dimensions=4)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

def maximum(x):
    max_x = max(x)
    result = [0,0,0,0]
    result[x.index(max_x)] = 1
    return result
net.make('Action', neurons=200, dimensions=4)
net.connect('Q', 'Qall')
net.connect('Qall', 'Action', func=maximum)

net.make_input('input', [0,0])
net.connect('input', 's')

net.add_to_nengo()
net.view()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

net.make_input('input', [0,0])
net.connect('input', 's')

e = 0.1
i = -1

transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
net.connect('Q', 'Q', transform=transform)


net.add_to_nengo()
net.view()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

net.make_array('Action', neurons=50, length=4, dimensions=1)
net.connect('Q', 'Action')

net.make_input('input', [0,0])
net.connect('input', 's')

e = 0.5
i = -1

transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]

# Let's force the feedback connection to only consider positive values
def positive(x):
    if x[0]<0: return [0]
    else: return x
net.connect('Action', 'Action', func=positive, transform=transform)


net.add_to_nengo()
net.view()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

net.make_array('Action', neurons=50, length=4, dimensions=1)
net.connect('Q', 'Action')

net.make_input('input', [0,0])
net.connect('input', 's')

e = 1
i = -1

transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
def positive(x):
    if x[0]<0: return [0]
    else: return x
net.connect('Action', 'Action', func=positive, transform=transform)


net.add_to_nengo()
net.view()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

net.make_array('Action', neurons=50, length=4, dimensions=1)
net.connect('Q', 'Action')

net.make_input('input', [0,0])
net.connect('input', 's')

e = 0.5
i = -1

transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
def positive(x):
    if x[0]<0: return [0]
    else: return x
net.connect('Action', 'Action', func=positive, transform=transform)

# Apply this function on the output
def select(x):
    if x[0]<=0: return [0]
    else: return [1]
net.make_array('ActionValue', neurons=50, length=4, dimensions=1)
net.connect('Action', 'ActionValue', func=select)


net.add_to_nengo()
net.view()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

net.make_array('Action', neurons=50, length=4, dimensions=1)
net.connect('Q', 'Action')

net.make_input('input', [0,0])
net.connect('input', 's')

e = 0.2
i = -1

transform = [[e, i, i, i], [i, e, i, i], [i, i, e, i], [i, i, i, e]]
def positive(x):
    if x[0]<0: return [0]
    else: return x
net.connect('Action', 'Action', func=positive, transform=transform)

# Apply this function on the output
def select(x):
    if x[0]<=0: return [0]
    else: return [1]
net.make_array('ActionValue', neurons=50, length=4, dimensions=1)
net.connect('Action', 'ActionValue', func=select)


net.add_to_nengo()
net.view()

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

D = 5
tau_ampa=0.002
tau_gaba=0.008
N = 50
radius = 1.5

import nef
net = nef.Network('Basal Ganglia')

net.make_input('input', [0]*D)

net.make_array('StrD1',N, D,intercept=(e,1),encoders=[[1]],radius=radius)
net.make_array('StrD2',N, D,intercept=(e,1),encoders=[[1]],radius=radius)
net.make_array('STN',N, D,intercept=(ep,1),encoders=[[1]],radius=radius)
net.make_array('GPi',N, D,intercept=(eg,1),encoders=[[1]],radius=radius)
net.make_array('GPe',N, D,intercept=(ee,1),encoders=[[1]],radius=radius)

net.connect('input','StrD1',weight=ws*(1+lg),pstc=tau_ampa)
net.connect('input','StrD2',weight=ws*(1-le),pstc=tau_ampa)
net.connect('input','STN',weight=wt,pstc=tau_ampa)

def func_str(x):
    if x[0]<e: return 0
    return mm*(x[0]-e)
net.connect('StrD1','GPi',func=func_str,weight=-wm,pstc=tau_gaba)
net.connect('StrD2','GPe',func=func_str,weight=-wm,pstc=tau_gaba)

def func_stn(x):
    if x[0]<ep: return 0
    return mp*(x[0]-ep)
tr=[[wp]*D for i in range(D)]    
net.connect('STN','GPi',func=func_stn,transform=tr,pstc=tau_ampa)
net.connect('STN','GPe',func=func_stn,transform=tr,pstc=tau_ampa)        

def func_gpe(x):
    if x[0]<ee: return 0
    return me*(x[0]-ee)
net.connect('GPe','GPi',func=func_gpe,weight=-we,pstc=tau_gaba)
net.connect('GPe','STN',func=func_gpe,weight=-wg,pstc=tau_gaba)

net.make_array('Action',N, D,intercept=(0.2,1),encoders=[[1]])
net.make_input('bias', [1]*D)
net.connect('bias', 'Action')
import numeric as np
net.connect('Action', 'Action', (np.eye(D)-1), pstc=tau_gaba)       


def func_gpi(x):
    if x[0]<eg: return 0
    return mg*(x[0]-eg)
net.connect('GPi','Action',func=func_gpi,pstc=tau_gaba,weight=-3)

net.add_to_nengo()
net.view()


import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

net.make_input('input', [0,0])
net.connect('input', 's')

D = 4
net.make_array('Action', neurons=50, length=D, dimensions=1, encoders=[[1]], intercept=(0.2,1))
net.make_input('bias', [1]*D)
net.connect('bias', 'Action')
import numeric as np
net.connect('Action', 'Action', (np.eye(D)-1), pstc=0.008)

import nps
nps.basalganglia.make_basal_ganglia(net, 'Q', 'Action', D, same_neurons=False, output_weight=-3)

net.add_to_nengo()
net.view()

import nef
net = nef.Network('Selection')
net.make('s', neurons=200, dimensions=2)
net.make_array('Q', neurons=50, length=4, dimensions=1)
net.connect('s', 'Q', transform=[[1,0],[-1,0],[0,1],[0,-1]])

net.make_input('input', [0,0])
net.connect('input', 's')

D = 4
net.make_array('Action', neurons=50, length=D, dimensions=1, encoders=[[1]], intercept=(0.2,1))
net.make_input('bias', [1]*D)
net.connect('bias', 'Action')
import numeric as np
net.connect('Action', 'Action', (np.eye(D)-1), pstc=0.008)

import nps
nps.basalganglia.make_basal_ganglia(net, 'Q', 'Action', D, same_neurons=False, output_weight=-3)

net.make('motor', neurons=100, dimensions=2)
net.connect('Action.0', 'motor', transform=[1,0]) 
net.connect('Action.1', 'motor', transform=[-1,0]) 
net.connect('Action.2', 'motor', transform=[0,1]) 
net.connect('Action.3', 'motor', transform=[0,-1]) 


net.add_to_nengo()
net.view()

import nef

net = nef.Network('Creature')

net.make_input('command_input', [0,0])

net.make('command', neurons=100, dimensions=2)

net.make('motor', neurons=100, dimensions=2)

net.make('position', neurons=1000, dimensions=2, radius=5)

net.make('scared_direction', neurons=100, dimensions=2)

def negative(x):
    return -x[0], -x[1]

net.connect('position', 'scared_direction', func=negative)

net.connect('position', 'position')


def rescale(x):
    return x[0]*0.1, x[1]*0.1

net.connect('motor', 'position', func=rescale)

net.connect('command_input', 'command')


D = 4

net.make_input('Q_input', [0]*D)
net.make_array('Q', neurons=50, length=D)
net.connect('Q_input', 'Q')
net.make_array('Action', neurons=50, length=D, dimensions=1, encoders=[[1]], intercept=(0.2,1))
net.make_input('bias', [1]*D)
net.connect('bias', 'Action')
import numeric as np
net.connect('Action', 'Action', (np.eye(D)-1), pstc=0.008)
import nps
nps.basalganglia.make_basal_ganglia(net, 'Q', 'Action', D, same_neurons=False, output_weight=-3)


net.make('do_command', 300, 3)
net.connect('command', 'do_command', index_post=[0,1])
net.connect('Action.0', 'do_command', index_post=[2])
def command(x):
    return x[2]*x[0], x[2]*x[1]
net.connect('do_command', 'motor', func=command)

net.make('do_scared', 300, 3)
net.connect('scared_direction', 'do_scared', index_post=[0,1])
net.connect('Action.1', 'do_scared', index_post=[2])
def command(x):
    return x[2]*x[0], x[2]*x[1]
net.connect('do_scared', 'motor', func=command)

net.add_to_nengo()

import nef

net = nef.Network('Creature')

net.make_input('command_input', [0,0])

net.make('command', neurons=100, dimensions=2)

net.make('motor', neurons=100, dimensions=2)

net.make('position', neurons=1000, dimensions=2, radius=5)

net.make('scared_direction', neurons=100, dimensions=2)

def negative(x):
    return -x[0], -x[1]

net.connect('position', 'scared_direction', func=negative)

net.connect('position', 'position')


def rescale(x):
    return x[0]*0.1, x[1]*0.1

net.connect('motor', 'position', func=rescale)

net.connect('command_input', 'command')


D = 4

net.make_input('Q_input', [0]*D)
net.make_array('Q', neurons=50, length=D)
net.connect('Q_input', 'Q')
net.make_array('Action', neurons=50, length=D, dimensions=1, encoders=[[1]], intercept=(0.2,1))
net.make_input('bias', [1]*D)
net.connect('bias', 'Action')
import numeric as np
net.connect('Action', 'Action', (np.eye(D)-1), pstc=0.008)
import nps
nps.basalganglia.make_basal_ganglia(net, 'Q', 'Action', D, same_neurons=False, output_weight=-3)


net.make('do_command', 200, 2)
net.connect('command', 'do_command')
net.connect('do_command', 'motor')
net.connect('GPi.0', 'do_command', encoders=-10)

net.make('do_scared', 300, 2)
net.connect('scared_direction', 'do_scared')
net.connect('do_scared', 'motor')
net.connect('GPi.1', 'do_scared', encoders=-10)

net.add_to_nengo()

from spa import *

D=16

class Rules: #Define the rules by specifying the start state and the 
             #desired next state
    def A(state='A'): #e.g. If in state A
        set(state='B') # then go to state B
    def B(state='B'):
        set(state='C')
    def C(state='C'):
        set(state='D')
    def D(state='D'):
        set(state='E')
    def E(state='E'):
        set(state='A')
    


class Sequence(SPA): #Define an SPA model (cortex, basal ganglia, thalamus)
    dimensions=16
    
    state=Buffer() #Create a working memory (recurrent network) object: 
                   #i.e. a Buffer
    BG=BasalGanglia(Rules()) #Create a basal ganglia with the prespecified 
                             #set of rules
    thal=Thalamus(BG) # Create a thalamus for that basal ganglia (so it 
                      # uses the same rules)
    
seq=Sequence()


from spa import *

D=16

class Rules: #Define the rules by specifying the start state and the 
             #desired next state
    def A(state='A'): #e.g. If in state A
        set(state='B') # then go to state B
    def B(state='B'):
        set(state='C')
    def C(state='C'):
        set(state='D')
    def D(state='D'):
        set(state='E')
    def E(state='E'):
        set(state='A')
    


class Sequence(SPA): #Define an SPA model (cortex, basal ganglia, thalamus)
    dimensions=16
    
    state=Buffer() #Create a working memory (recurrent network) object: 
                   #i.e. a Buffer
    BG=BasalGanglia(Rules()) #Create a basal ganglia with the prespecified 
                             #set of rules
    thal=Thalamus(BG) # Create a thalamus for that basal ganglia (so it 
                      # uses the same rules)
    
    input=Input(0.1,state='D') #Define an input; set the input to 
                               #state D for 100 ms

seq=Sequence()


from spa import *

D=16

class Rules: #Define the rules by specifying the start state and the 
             #desired next state
    def start(vision='(A+B+C+D+E)*2'):
        set(state=vision)
    def A(state='A'): #e.g. If in state A
        set(state='B') # then go to state B
    def B(state='B'):
        set(state='C')
    def C(state='C'):
        set(state='D')
    def D(state='D'):
        set(state='E')
    def E(state='E'):
        set(state='A')
    


class Routing(SPA): #Define an SPA model (cortex, basal ganglia, thalamus)
    dimensions=16

    state=Buffer() #Create a working memory (recurrent network) 
                   #object: i.e. a Buffer
    vision=Buffer(feedback=0) #Create a cortical network object with no 
                              #recurrence (so no memory properties, just 
                              #transient states)
    BG=BasalGanglia(Rules) #Create a basal ganglia with the prespecified 
                           #set of rules
    thal=Thalamus(BG) # Create a thalamus for that basal ganglia (so it 
                      # uses the same rules)

    input=Input(0.1,vision='D') 
    
model=Routing()