# A Subsumption Behavior-based Brain
# Should be used on a pioneer robot with blob vision
# Uses four behaviors: Wander, Avoid, Seek (a given color), Grasp
# SubsumptionGrasp.py use with FindPucks.world

from pyro.brain import Brain
import time
import random

class Behavior:
    def __init__(self):
        self.translate = 0
        self.rotate = 0
        self.flag = 0

    def setRobot(self, robot):
        self.robot = robot

    def getRobot(self):
        return self.robot

    def move(self, translate, rotate):
        self.translate = translate
        self.rotate = rotate
        self.flag = 1

    # Returns a list of all blobs on the given color channel.   
    def getBlobChannel(self, channel):   
        index = 0   
        data = self.getRobot().getServiceData("blob")  
        if channel == 'red':   
            index = 0   
        elif channel == 'green':   
            index = 1   
        elif channel == 'blue':   
            index = 2   
        else:   
            print "unrecognized color channel"   
        if len(data) == 0 or len(data[1]) == 0:   
            return []   
        else:   
            return data[1][index]   
   
    # Returns a list of the two key features of the blob information for the   
    # closest blob on the given channel: direction and range.
    #
    # The blob information consists of a list of nine features:   
    # Index 2: is an integer between 0 and 160 which represents the x position
    # of the blob's centeroid.  A value of 0 indicates that the blob is at the
    # farthest left location, a value of 80 indicates that the blob is centered
    # and a value of 160 indicates that the blob is at the farthest right
    # location.   
    # Index 8: is an integer representing the range to the blob.
    def getClosestBlob(self, channel):   
        all = self.getBlobChannel(channel)   
        if len(all) == 0:   
            return 'None'   
        else:   
            if all[0][2] < 75:   
                turnDirection = 1   
            elif all[0][2] > 85:   
                turnDirection = -1   
            else:   
                turnDirection = 0   
            return turnDirection, all[0][8]   

class SubsumptionBrain(Brain):
    def __init__(self, engine):
        Brain.__init__(self, 'SubsumptionBrain', engine)
        self.behaviors = []
        self.getRobot().startServices(['blob', 'gripper'])   

    def getRobot(self):
        return self.robot

    def add(self, behavior):
        behavior.setRobot( self.robot )
        self.behaviors.append( behavior )

    def step(self):
        b = self.updateAll()
        print "%s is in control" % self.behaviors[b].__class__.__name__
        self.getRobot().move(self.behaviors[b].translate,
                             self.behaviors[b].rotate)
        time.sleep(1.0)

    def updateAll(self):
        # for all except lowest:
        for b in range(len(self.behaviors) - 1, 0, -1):
            self.behaviors[b].flag = 0
            self.behaviors[b].update()
            # if it fired, return number:
            if self.behaviors[b].flag:
                return b
        # if none fired, return lowest:
        self.behaviors[0].update()
        return 0


# -----------------------------------------------------

class Wander(Behavior):
    def update(self):
        self.move( .2, random.random() * 2 - 1)

class Avoid(Behavior):
    def update(self):
        front = self.getRobot().get('range', 'value', 'front-all', 'min')
        if front[1] < 1:
            if front[0] >= 1 and front[0] <= 3:
                self.move( -.1, -.15)
            else:
                self.move( -.1, .15)

class Seek(Behavior):
    def __init__(self, color):
        self.color = color
        
    def update(self):
        result = self.getClosestBlob(self.color)
        left = self.getRobot().get('range', 'value', 'front-left', 'min')[1]
        right = self.getRobot().get('range', 'value', 'front-right', 'min')[1]
        if result != 'None' and left > 0.4 and right > 0.4:   
            blobDirection = result[0]   
            blobRange = result[1]
            if blobRange < 380:
                self.move(0.01, 0)
            elif blobDirection > 0:   
                self.move(0.15, 0.05)   
            elif blobDirection < 0:   
                self.move(0.15, -0.05)   
            else:   
                self.move(0.2, 0)

class Grasp(Behavior):
    def update(self):
        gripper = self.getRobot().getService("gripper")
        if gripper.isClosed():
            self.move(0,0)
            gripper.store()   
            time.sleep(1)   
            gripper.open()   
            time.sleep(1)
        elif gripper.getBreakBeamState():
            self.move(0,0)
            gripper.close()   
            
def INIT(engine):
    subsumption = SubsumptionBrain( engine )
    # add behaviors, lowest priorities first:
    subsumption.add( Wander() )
    subsumption.add( Avoid() )
    subsumption.add( Seek('red') )
    subsumption.add( Grasp() )
    return subsumption