#!/usr/bin/python3

import numpy


class TrapezoidProfile(object):
    """Computes a trapezoidal motion profile

  Attributes:
    _acceleration_time: the amount of time the robot will travel at the
        specified acceleration (s)
    _acceleration: the acceleration the robot will use to get to the target
        (unit/s^2)
    _constant_time: amount of time to travel at a constant velocity to reach
        target (s)
    _deceleration_time: amount of time to decelerate (at specified
        deceleration) to target (s)
    _deceleration: decceleration the robot needs to get to goal velocity
        (units/s^2)
    _maximum_acceleration: the maximum acceleration (units/s^2)
    _maximum_velocity: the maximum velocity (unit/s)
    _timestep: time between calls to Update (delta_time)
    _output: output array containing distance to goal and velocity
  """

    def __init__(self, delta_time):
        """Constructs a TrapezoidProfile.

    Args:
      delta_time: time between calls to Update (seconds)
    """
        self._acceleration_time = 0
        self._acceleration = 0
        self._constant_time = 0
        self._deceleration_time = 0
        self._deceleration = 0

        self._maximum_acceleration = 0
        self._maximum_velocity = 0
        self._timestep = delta_time

        self._output = numpy.array(numpy.zeros((2, 1)))

    # Updates the state
    def Update(self, goal_position, goal_velocity):
        self._CalculateTimes(goal_position - self._output[0], goal_velocity)

        next_timestep = self._timestep

        # We now have the amount of time we need to accelerate to follow the
        # profile, the amount of time we need to move at constant velocity
        # to follow the profile, and the amount of time we need to decelerate to
        # follow the profile.  Do as much of that as we have time left in dt.
        if self._acceleration_time > next_timestep:
            self._UpdateVals(self._acceleration, next_timestep)
        else:
            self._UpdateVals(self._acceleration, self._acceleration_time)

            next_timestep -= self._acceleration_time
            if self._constant_time > next_timestep:
                self._UpdateVals(0, next_timestep)
            else:
                self._UpdateVals(0, self._constant_time)
                next_timestep -= self._constant_time
                if self._deceleration_time > next_timestep:
                    self._UpdateVals(self._deceleration, next_timestep)
                else:
                    self._UpdateVals(self._deceleration,
                                     self._deceleration_time)
                    next_timestep -= self._deceleration_time
                    self._UpdateVals(0, next_timestep)

        return self._output

    # Useful for preventing windup etc.
    def MoveCurrentState(self, current):
        self._output = current

    # Useful for preventing windup etc.
    def MoveGoal(self, dx):
        self._output[0] += dx

    def SetGoal(self, x):
        self._output[0] = x

    def set_maximum_acceleration(self, maximum_acceleration):
        self._maximum_acceleration = maximum_acceleration

    def set_maximum_velocity(self, maximum_velocity):
        self._maximum_velocity = maximum_velocity

    def _UpdateVals(self, acceleration, delta_time):
        self._output[0, 0] += (self._output[1, 0] * delta_time +
                               0.5 * acceleration * delta_time * delta_time)
        self._output[1, 0] += acceleration * delta_time

    def _CalculateTimes(self, distance_to_target, goal_velocity):
        if distance_to_target == 0:
            self._acceleration_time = 0
            self._acceleration = 0
            self._constant_time = 0
            self._deceleration_time = 0
            self._deceleration = 0
            return
        elif distance_to_target < 0:
            # Recurse with everything inverted.
            self._output[1] *= -1
            self._CalculateTimes(-distance_to_target, -goal_velocity)
            self._output[1] *= -1
            self._acceleration *= -1
            self._deceleration *= -1
            return

        self._constant_time = 0
        self._acceleration = self._maximum_acceleration
        maximum_acceleration_velocity = (
            distance_to_target * 2 * numpy.abs(self._acceleration) +
            self._output[1] * self._output[1])
        if maximum_acceleration_velocity > 0:
            maximum_acceleration_velocity = numpy.sqrt(
                maximum_acceleration_velocity)
        else:
            maximum_acceleration_velocity = -numpy.sqrt(
                -maximum_acceleration_velocity)

        # Since we know what we'd have to do if we kept after it to decelerate, we
        # know the sign of the acceleration.
        if maximum_acceleration_velocity > goal_velocity:
            self._deceleration = -self._maximum_acceleration
        else:
            self._deceleration = self._maximum_acceleration

        # We now know the top velocity we can get to.
        top_velocity = numpy.sqrt(
            (distance_to_target + (self._output[1] * self._output[1]) /
             (2.0 * self._acceleration) + (goal_velocity * goal_velocity) /
             (2.0 * self._deceleration)) / (-1.0 /
                                            (2.0 * self._deceleration) + 1.0 /
                                            (2.0 * self._acceleration)))

        # If it can go too fast, we now know how long we get to accelerate for and
        # how long to go at constant velocity.
        if top_velocity > self._maximum_velocity:
            self._acceleration_time = (
                (self._maximum_velocity - self._output[1]) /
                self._maximum_acceleration)
            self._constant_time = (
                distance_to_target +
                (goal_velocity * goal_velocity -
                 self._maximum_velocity * self._maximum_velocity) /
                (2.0 * self._maximum_acceleration)) / self._maximum_velocity
        else:
            self._acceleration_time = ((top_velocity - self._output[1]) /
                                       self._acceleration)

        if self._output[1] > self._maximum_velocity:
            self._constant_time = 0
            self._acceleration_time = 0

        self._deceleration_time = ((goal_velocity - top_velocity) /
                                   self._deceleration)