#FORMAT python

from pyrobot.brain import Brain  
from pyrobot.brain.conx import *  
from time import * 
  
class NNBrain(Brain):  
   def setup(self):
      self.net = Network()
      self.net.addLayers(3,2,2)  
      self.net.setEpsilon(0.25)  
      self.net.setMomentum(.1)  
      self.maxvalue = self.robot.range.getMaxvalue()
      self.counter = 0  
      self.doneLearning = 0 
    
   # Scale the range readings to be between 0 and 1.  Also make close 
   # obstacles register high readings, and distance obstacles low readings. 
   def scale(self, val):  
      return (1 - (val / self.maxvalue))  
  
   # The robot will get translate and rotate values in the range [-0.5,0.5],  
   # but the neural network will generate outputs in the range [0,1].  
   def toNetworkUnits(self, val):  
      return (val + 0.5)  
  
   def toRobotUnits(self, val):  
      return (val - 0.5) 
 
   def determineTargets(self, left, front, right): 
      if front < 1.0: 
         print "front" 
         return([0.0, 0.5]) 
      elif left < 1.0: 
         print "left" 
         return([0.0, -0.5]) 
      elif right < 1.0: 
         print "right" 
         return([0.0, 0.5]) 
      else: 
         print "clear" 
         return([0.5, 0.0]) 
 
   def step(self): 
      if self.doneLearning:  
          self.net.setLearning(0)  
      else:  
          self.net.setLearning(1) 
      # Set inputs 
      left  = min([s.distance() for s in self.robot.range["front-left"]])
      front = min([s.distance() for s in self.robot.range["front"]])
      right = min([s.distance() for s in self.robot.range["front-right"]])
      inputs = map(self.scale, [left, front, right])
      trnTarget,rotTarget = self.determineTargets(left, front, right) 
      targets = [self.toNetworkUnits(trnTarget), self.toNetworkUnits(rotTarget)]
      # Learn 
      self.net.step(input=inputs,output=targets)  # input and output are layer names
      trnActual = self.toRobotUnits(self.net['output'].activation[0]) 
      rotActual = self.toRobotUnits(self.net['output'].activation[1]) 
      # Check if the robot is stuck and give it a kick to unjam it 
      if self.robot.stall: 
          print "stuck--reversing" 
          self.move(-0.5, 0.0) 
          sleep(0.5) 
      else:
          if self.doneLearning:
              print "move", trnActual, rotActual
              self.move(trnActual, rotActual)
          else:
              print "step", self.counter, "target", trnTarget, rotTarget, \
                    "network", trnActual, rotActual 
              self.move(trnTarget, rotTarget) 
      self.counter += 1 
  
def INIT(engine): 
   return NNBrain('NNBrain', engine)