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

joints = ["gripper_left_finger_joint", "gripper_right_finger_joint", ]
# For this module, ensure motors and sensors only contain the relevant entries.
motors = [motors[joint] for joint in joints]
sensors = [sensors[joint] for joint in joints]

class CloseGripper(Behaviour):
    def __init__(self, name):
        super().__init__(name)
        self.position = .045

    def initialise(self):
        for motor in motors:
            motor.setPosition(self.position)
        self.initial_time = robot.getTime()
        print("%s.initialise()" % (self.name, ))

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

        # Close the grip little by little, as setting the position to 0 all at
        # once creates strain on the gripper.
        self.position -= .003
        for motor in motors:
            motor.setPosition(max(0, self.position))

        # Terminate with SUCCESS when both grippers are exerting enough force.
        if all([motor.getForceFeedback() < GRASP_THRESHOLD for motor in motors]):
            print(
                "%s.update()[%s->%s], both ends grasped with force > %.2f"
                % (self.name, self.status, Status.SUCCESS, abs(GRASP_THRESHOLD))
            )
            return Status.SUCCESS

        # Both grippers reached 0 without enough force being exerted, meaning
        # the jar was not grabbed.
        if all([sensor.getValue() < JOINT_TOLERANCE for sensor in sensors]):
            print(
                "%s.update()[%s->%s], gripper completely closed (empty)"
                % (self.name, self.status, Status.FAILURE)
            )
            return Status.FAILURE

        return Status.RUNNING

    def terminate(self, new_status):
        print(
            "%s.terminate()[%s->%s], gripper positions: %.2f, %.2f"
            % (self.name, self.status, new_status, *[s.getValue() for s in sensors])
        )

class OpenGripper(Behaviour):
    def __init__(self, name):
        super().__init__(name)

    def initialise(self):
        print("%s.initialise()" % (self.name, ))
        for motor in motors:
            motor.setPosition(0.045)
        self.initial_time = robot.getTime()

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

        # Both joints are within JOINT_TOLERANCE from their respective targets.
        if all([abs(sensor.getValue() - .045) < JOINT_TOLERANCE for sensor in sensors]):
            print(
                "%s.update()[%s->%s], gripper completely opened"
                % (self.name, self.status, Status.SUCCESS)
            )
            return Status.SUCCESS

        return Status.RUNNING