#FORMAT python

from pyrobot.brain.behaviors import *
from time import *

class GatherPucksBrain(FSMBrain):
   def setup(self):
      if not self.robot.hasA("camera"):
         self.startDevice('camera') # or BlobCamera, or PlayerCamera
      if not self.robot.hasA("gripper"):
         self.startDevice('gripper')
      self.camera = self.robot.camera[0]
      self.camera.clearFilters()
      self.camera.addFilter("match",255,0,0,) # red
      self.camera.addFilter("blobify",0,255,255,0,1,1,1,) # red

class locatePuck(State):
   """
   Rotate until a red blob is seen or a complete scan has been done.
   """
   def onActivate(self):
      self.searches = 0             # counter for number of rotations

   def step(self):
      redBlobs = self.robot.camera[0].filterResults[1]
      if len(redBlobs) != 0:
         self.move(0, 0)
         print "found a puck!"
         self.goto('approachPuck')
      elif self.searches > 275:
         print "found all pucks"
         self.goto('done')
      else:
         print "searching for a puck"
         self.searches += 1         # update counter
         self.move(0, 0.2)

class approachPuck(State):
   """
   Move towards closest red blob by centering it in blobfinder.
   Once centered, then consider distance to the closest red blob.
   Once close, move slowly until it is within gripper.
   """
   def step(self):
      redBlobs = self.robot.camera[0].filterResults[1]
      if len(redBlobs) == 0:
         print "no puck in sight"
         self.goto('locatePuck')
      elif redBlobs[0][2] < 50: # sizes need to be changed for pyrobot
         print "puck far left"
         self.move(0, 0.2)
      elif redBlobs[0][2] < 75:
         print "puck to left"
         self.move(0.2, 0.1)
      elif redBlobs[0][2] > 100:
         print "puck far right"
         self.move(0, -0.2)
      elif redBlobs[0][2] > 85:
         print "puck to right"
         self.move(0.2, -0.1)
      else:
         if self.robot.gripper[0].getBreakBeam('inner') > 1:
            print "grabbing puck"
            self.move(0, 0)
            self.goto('grabPuck')
         elif redBlobs[0][4] > 500: #area
            print "puck straight ahead"
            self.move(0.3, 0)
         else:
            print "puck very close"
            self.move(0.1, 0)

class grabPuck(State):
   """
   Pick up the current puck and then return to locating another puck.
   """
   def step(self):
      self.robot.gripper[0].store()
      sleep(2.0)
      self.robot.gripper[0].open()
      sleep(1.0)
      self.goto('locatePuck')

class done(State):
   """
   No more visible pucks, so stop moving.
   """
   def update(self):
      self.move(0, 0)

def INIT(engine):
    brain = GatherPucksBrain(engine=engine)
    brain.add(locatePuck(1))
    brain.add(approachPuck())
    brain.add(grabPuck())
    brain.add(done())
    return brain