tau = 0.1

import nef

net=nef.Network('Creature', seed=4)

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

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


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

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

net.connect('motor', 'position', weight=tau)

net.make('v to home', neurons=100, dimensions=2)

def direction_home(position):
    x = -position[0]
    y = -position[1]
    length = math.sqrt(x*x+y*y)
    if length>0:
        x = x/length
        y = y/length
    return [x,y]

#net.connect('position', 'v to home', func=direction_home)
net.connect('position', 'v to home', weight=-10)


net.make('scared', neurons=50, dimensions=1)
net.make_input('scared input', [0])
net.connect('scared input', 'scared')


net.make('s and vh', neurons=500, dimensions=3)
net.make('s and vc', neurons=500, dimensions=3)

net.connect('scared', 's and vh', index_post=[0])
net.connect('v to home', 's and vh', index_post=[1,2])

net.connect('scared', 's and vc', index_post=[0])
net.connect('command', 's and vc', index_post=[1,2])

def s_and_vc(state):
    s, x, y=state
    return [(1-s)*x, (1-s)*y]
    
def s_and_vh(x):
    return [x[0]*x[1], x[0]*x[2]]
    

net.connect('s and vc', 'motor', func=s_and_vc)
net.connect('s and vh', 'motor', func=s_and_vh)

net.add_to_nengo()