from py_trees.behaviour import Behaviour
from py_trees.common import Status
from constants import WALK_SPEED, TIMEOUT
from setup import robot, gps, compass, left_motor, right_motor
from utilities import euclidean, bound_by, deg, fix_arctan
from numpy import arctan2

# Move back a little to give some space for arm movement, while trying to
# maintain the original heading.
class GoBack(Behaviour):
    def __init__(self, name, distance):
        super().__init__(name)
        self.target = distance

    def initialise(self):
        self.initial_time = robot.getTime()
        self.initial_pos = gps.getValues()[:2]
        self.initial_theta = arctan2(*compass.getValues()[:2])
        print(
            "%s.initialise(), from (%.3f, %.3f), at %.1f°"
            % (self.name, *self.initial_pos, deg(self.initial_theta))
        )

    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(theta - self.initial_theta)

        if euclidean(self.initial_pos, gps.getValues()) >= self.target:
            print(
                "%s.update()[%s->%s], angle error = %.1f°"
                % (self.name, self.status, Status.SUCCESS, deg(alpha))
            )
            return Status.SUCCESS

        left_motor.setVelocity(-WALK_SPEED + alpha)
        right_motor.setVelocity(-WALK_SPEED - alpha)
        return Status.RUNNING

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