from py_trees.behaviour import Behaviour
from py_trees.common import Status
from setup import robot, motors, sensors
from constants import JOINT_TOLERANCE, TIMEOUT

# Move the provided joints towards the target values.
# Mainly used to control the arm.
class MoveArm(Behaviour):
    def __init__(self, name, targets):
        super().__init__(name)
        self.targets = targets

    def initialise(self):
        self.initial_time = robot.getTime()
        for joint, target in self.targets.items():
            motors[joint].setPosition(target)
        print("%s.initialise()" % (self.name))

    def update(self):
        if robot.getTime() - self.initial_time > TIMEOUT:
            print(
                "%s.update()[%s->%s], Timed out"
                % (self.name, self.status, Status.FAILURE)
            )
            return Status.FAILURE

        # All joints are within JOINT_TOLERANCE from their respective targets.
        if all([abs(sensors[joint].getValue() - target) < JOINT_TOLERANCE
                for joint, target in self.targets.items()]):
            print(
                "%s.update()[%s->%s]"
                % (self.name, self.status, Status.SUCCESS)
            )
            return Status.SUCCESS
        return Status.RUNNING