#FORMAT python

#VisRobotBrain.py
from pyrobot.brain import Brain
from pyrobot.brain.vconx import *
from pyrobot.gui.plot.scatter import *
from pyrobot.gui.plot.hinton import *
from os import system, unlink

def a2s(a):
   """ array to string """
   retval = ''
   for i in range(len(a)):
      retval += "%.5f " % a[i]
   return retval

class NNBrain(Brain):
   """
   This is an example brain controlled by a neural network.
   This simple example loads the range sensor data into the
   input layer, and trains the network to stay away from
   things.
   """
   def setup(self):
      """ Create the network. """
      self.net = VINetwork()
      self.hiddenLayerSize = 10
      self.net.addLayers(self.robot.range.count,
                              self.hiddenLayerSize, 2)
      self.net.setBatch(0)
      self.net.initialize()
      self.net.setEpsilon(0.5)
      self.net.setMomentum(.1)
      self.net.setLearning(1)
      self.counter = 0
      self.maxvalue = self.robot.range.getMaxvalue()
      
   def destroy(self):
      self.net.destroy()
      
   def scale(self, val):
      return (val / self.maxvalue)
   
   def step(self):         
      # First, set inputs and targets:
      ins = [self.scale(s.value) for s in self.robot.range["all"]]
      # Compute targets:
      target_rotate = 0.5
      if min([s.value for s in self.robot.range["front"]]) < 1:
         target_trans = 0.0
         target_rotate = 0.0
      elif min([s.value for s in self.robot.range["back"]]) < 1:
         target_trans = 1.0
      else:
         target_trans = 1.0
      if min([s.value for s in self.robot.range["left"]])  < 1:
         target_rotate = 0.0
      elif min([s.value for s in self.robot.range["right"]])  < 1:
         target_rotate = 1.0
      target = [target_trans, target_rotate]

      # set inputs and targets with step
      self.net.step( input = ins, output = target)

      # get the output, and move:
      trans = (self.net['output'].activation[0] - .5) / 2.0
      rotate = (self.net['output'].activation[1] - .5) / 2.0
      self.robot.move(trans, rotate)
      self.counter += 1

def INIT(engine):
   return NNBrain('NNBrain', engine)