from py_trees.behaviour import Behaviour
from py_trees.common import Status
from constants import MAX_SPEED, ROTATION_SPEED, P1, P2, APPROACH_STOP, ANGLE_TOLERANCE, TIMEOUT
from setup import blackboard, robot, gps, compass, left_motor, right_motor
from utilities import deg, bound_by, euclidean, fix_arctan
from numpy import arctan2, sign

# Move to the target_pos set by JarFinder.
class ApproachJar(Behaviour):
    def __init__(self, name):
        super().__init__(name)

    def initialise(self):
        self.initial_time = robot.getTime()
        print(
            "%s.initialise(), going towards (%.3f, %.3f)"
            % (self.name, *blackboard["target_pos"])
        )

    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

        robot_pos = gps.getValues()[:2]
        theta = arctan2(*compass.getValues()[:2])
        d_dist = blackboard["target_pos"] - robot_pos

        rho = euclidean(robot_pos, blackboard["target_pos"])
        alpha = fix_arctan(arctan2(d_dist[1], d_dist[0]) - theta)

        left_speed = bound_by(-P1 * alpha + P2 * rho , MAX_SPEED)
        right_speed = bound_by(P1 * alpha + P2 * rho , MAX_SPEED)
        left_motor.setVelocity(left_speed)
        right_motor.setVelocity(right_speed)

        if rho < APPROACH_STOP: # The robot is now close enough.
            print(
                "%s.update()[%s->%s], angle error = %.1f°"
                % (self.name, self.status, Status.SUCCESS, deg(alpha))
            )
            return Status.SUCCESS
        return Status.RUNNING

    def terminate(self, new_status):
        left_motor.setVelocity(0.)
        right_motor.setVelocity(0.)
        print(
            "%s.terminate()[%s->%s], new position: (%.3f, %.3f)"
            % (self.name, self.status, new_status, *gps.getValues()[:2])
        )

class RotateRobot(Behaviour):
    def __init__(self, name, target_angle):
        super().__init__(name)
        self.target_angle = target_angle

    def initialise(self):
        self.initial_time = robot.getTime()
        print(
            "%s.initialise(), rotating towards %.1f°"
            % (self.name, deg(self.target_angle))
        )

    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

        theta = arctan2(*compass.getValues()[:2])
        alpha = fix_arctan(self.target_angle - theta)

        if abs(alpha) < ANGLE_TOLERANCE: # Close enough.
            print(
                "%s.update()[%s->%s], angle error = %.1f°"
                % (self.name, self.status, Status.SUCCESS, deg(alpha))
            )
            return Status.SUCCESS

        direction = sign(alpha)
        left_motor.setVelocity(-ROTATION_SPEED * direction)
        right_motor.setVelocity(ROTATION_SPEED * direction)
        return Status.RUNNING

    def terminate(self, new_status):
        left_motor.setVelocity(0.)
        right_motor.setVelocity(0.)
        print(
            "%s.terminate()[%s->%s], new position: (%.3f, %.3f)"
            % (self.name, self.status, new_status, *gps.getValues()[:2])
        )