from pyrobot.brain import Brain
from pyrobot.brain.conx import *
from time import *

class NNBrain(Brain):
    """
    Test the trained network and measure its performance.  NOTE: Must
    translate the motor commands generated by the network from 0-1
    range back to -0.5-+0.5 range.
    """
    def setup(self):
        self.n = Network()
        self.n.addLayers(8,2,1)
        #? self.maxvalue = self.robot.range.getMaxvalue()
        self.n.loadWeightsFromFile("balance-net.wts")
        #? self.n.setLearning(0)
        self.counter = 1
        self.totalDist = 0
        
    def scaleSensor(self, val):
        x = val / self.maxvalue
        if x > 1:
            return 1.0
        else:
            return x
        
    def scaleMotor(self, val):
        return(val - 0.5)
    
    def step(self):
        """
        if self.counter % 200 == 0:
            print "average distance from wall", self.totalDist/200.0
            self.totalDist = 0
        self.totalDist += max([s.distance() for s in self.robot.range["left"]])
       # frontLeft = self.robot.range[0].distance()
       # backLeft = self.robot.range[15].distance()
       # front = min([s.distance() for s in self.robot.range["front"]])
        leftSensor = self.robot.light[0][0].value
        rightSensor = self.robot.light[0][1].value

        self.n['input'].copyActivations([self.scaleSensor(leftSensor),
                                         self.scaleSensor(rightSensor)])
                                         #self.scaleSensor(backLeft)])
        self.n.propagate()
        translate = self.scaleMotor(float(self.n['output'].activation[0]))
        rotate = self.scaleMotor(float(self.n['output'].activation[1]))
        self.move(translate, rotate)
        self.counter += 1
        """
        self.n['input'].copyActivations([self.scaleSensor(leftSensor),
                                         self.scaleSensor(rightSensor)])
                                         #self.scaleSensor(backLeft)])
        self.n.propagate()
        
        
    

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