#FORMAT python

# brain to collect data for offline learning
# robot goes forward and then slows to a stop when it detects something

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

def saveListToFile(ls, file):
   for i in range(len(ls)):
      file.write(str(ls[i]) + " ")
   file.write("\n")

class CollectDataBrain(Brain):

   def setup(self):
      self.counter = 0
      self.countstopping = 0
      self.datafile1 = open("frontsensors.dat", "w")
      self.datafile2 = open("translatetargets.dat", "w")
      self.maxvalue = self.robot.range.getMaxvalue()
      print "max sensor value is ", self.maxvalue

   def determineMove(self, front):
      if front[0] < 0.5 or front[1] < 0.5:
         print "collision imminent"
         self.countstopping = self.countstopping+1
         return(0)
      elif front[0] < 0.8 or front[1] < 0.8:
         print "object detected"
         return(0.08)
      else:
         print "clear"
         return(0.15)
   def scale(self, val):
      x = val / self.maxvalue
      if x > 1:
         print "scaled > 1"
         return 1
      else:
         return x
   def step(self):
      front = [s.distance() for s in self.robot.range["front"]]
      translation = self.determineMove(front)
      print "front sensors", front[0], front[1]
      if self.counter > 1000:
         self.datafile1.close()
         self.datafile2.close()
         print "done collecting data"
         self.stop()
      elif self.countstopping > 5:
         self.move(-0.3, 0.2)
         sleep(0.5)
         self.countstopping = 0
      else:
         print "move", self.counter, translation
         saveListToFile([self.scale(front[0]), self.scale(front[1])],
                        self.datafile1)
         saveListToFile([translation] , self.datafile2)
         self.move(translation, 0.0)
         self.counter += 1

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