import nef
net = nef.Network('Example')
net.make_input('input', [0])

net.make('A', neurons=75, dimensions=1)
net.make('B', neurons=50, dimensions=1)

def square(x):
    return x[0]*x[0]

net.connect('A', 'B', func=square)

net.connect('input', 'A')

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')

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()