Source code for local_motion.velocity_profile

"""Module to calculate a velocity profile for a local planned path."""

import numpy as np
import numpy.typing as npt
from scipy.signal import find_peaks


class VelocityProfileConfig():
    """
    Configuration Container for VelocityProfile.

    Attributes
    ----------
    g : float
        Gravitational constant.
    a_y_factor : float
        Factor for lateral acceleration.
    a_x_factor : float
        Factor for longitudinal acceleration.
    stretch_time_for_mpc : bool
        Whether the velocity profile should be stretched for the mpc.
    mpc_min_time_horizont : float
        Time the velocity profile should be stretched to for the mpc. Deprecated.
    v_max : float
        Maximal allowed velocity.
    v_end : float
        Velocity for the last point of the trajectory.
    a_x_factor_completed : float
        Factor for longitudinal velocity for braking when mission is completed to come to a standstill.
    """

    def __init__(self, g: float = 9.81, a_y_factor: float = 0.8, a_x_factor: float = 0.9,
                 stretch_time_for_mpc: bool = False, mpc_min_time_horizont: float = 1.1,
                 v_max: float = 15.5, v_end: float = 0, a_x_factor_completed: float = 0.8) -> None:
        """
        Initialize Configuration for VelocityProfile.

        Parameters
        ----------
        g : float, optional
            Gravitational constant, by default 9.81
        a_y_factor : float, optional
            Factor for lateral acceleration, by default 0.8
        a_x_factor : float, optional
            Factor for longitudinal acceleration, by default 0.9
        stretch_time_for_mpc : bool, optional
            Whether the velocity profile should be stretched for the mpc, by default False
        mpc_min_time_horizont : float, optional
            Time the velocity profile should be stretched to for the mpc. Deprecated, by default 1.1
        v_max : float, optional
            Maximal allowed velocity, by default 15.5
        v_end : float, optional
            Velocity for the last point of the trajectory, by default 0
        a_x_factor_completed : float, optional
            Factor for longitudinal velocity for braking when
            mission is completed to come to a standstill, by default 0.8
        """
        self.g = g
        self.a_y_factor = a_y_factor
        self.a_x_factor = a_x_factor
        self.v_max = v_max
        self.v_end = v_end
        self.a_x_factor_completed = a_x_factor_completed
        self.stretch_time_for_mpc = stretch_time_for_mpc
        self.mpc_min_time_horizont = mpc_min_time_horizont


class VelocityProfile():
    """
    Class to calculate a velocity profile for a given path.

    Attributes
    ----------
    g : float
        Gravitational constant.
    a_y_max : float
        Maximal allowed lateral acceleration.
    a_x_max : float
        Maximal allowed longitudinal acceleration.
    stretch_time_for_mpc : bool
        Whether the velocity profile should be stretched for the mpc.
    mpc_min_time_horizont : float
        Time the velocity profile should be stretched to for the mpc. Deprecated.
    v_max : float
        Maximal allowed velocity.
    v_end : float
        Velocity for the last point of the trajectory.
    a_x_max_completed : float
        Maximal allowed longitudinal velocity for braking when mission is completed to come to a standstill.
    v_initial : Optional[float]
        Velocity at first path point.
    mission_completed : bool
        Indicates whether the mission is completed and the vehicle should come to a standstill.
    path : npt.NDArray
        Path for which the velocity profile should be calculated for, shape: (n, ) with n equal number of path points.
    curvature : npt.NDArray
        Curvature of the path,  shape: (n, ) with n equal number of path points.
    radius : npt.NDArray
        Radius of the path,  shape: (n, ) with n equal number of path points.
    delta_distance : npt.NDArray
        Distance between each subsequent pair of path points, shape: (n, ) with n equal number of path points.
    cum_distance : npt.NDArray
        Cumulative distance of path,  shape: (n, ) with n equal number of path points.
    apex : npt.NDArray
        List of apex points, shape: (m, 2) with m equal number of apex points.
    velocity_max : npt.NDArray
        Maximal possible velocity based on the path curvature and
        maximal longitudinal acceleration, shape: (n, ) with n equal number of path points.
    brake_velocity : npt.NDArray
        Maximal possible velocity based on accelerating reverse from one support point
        to the previous one, shape: (n, ) with n equal number of path points.
    acceleration_velocity : npt.NDArray
        Maximal possible velocity based on accelerating forward from one support point
        to the next one, shape: (n, ) with n equal number of path points.
    velocity : npt.NDArray
        Resulting velocity profile by combining maximal possible velocity based on curvature,
        braking and accelerating, shape: (n, ) with n equal number of path points.
    """

    def __init__(self, path: npt.NDArray, mission_completed: bool = False,
                 velocity_profile_config: VelocityProfileConfig = VelocityProfileConfig()):
        """
        Initialize VelocityProfile instance.

        Parameters
        ----------
        path : npt.NDArray
            Path for which the velocity profile should be calculated, shaoe: (n, 2) with n equal number of path points.
        mission_completed : bool, optional
            Indicates whether the mission is completed and the vehicle should come to a standstill, by default False
        velocity_profile_config : VelocityProfileConfig, optional
            Configuration for the velocity profile, by default VelocityProfileConfig()
        """
        self.g = velocity_profile_config.g
        self.a_y_max = round(velocity_profile_config.a_y_factor * self.g, 3)
        self.a_x_max = round(velocity_profile_config.a_x_factor * self.g, 3)
        self.a_x_max_completed = round(velocity_profile_config.a_x_factor_completed * self.g, 3)
        self.v_max = velocity_profile_config.v_max
        self.v_end = velocity_profile_config.v_end
        self.stretch_time_for_mpc = velocity_profile_config.stretch_time_for_mpc
        self.mpc_min_time_horizont = velocity_profile_config.mpc_min_time_horizont

        self.path = path
        derivation1 = np.gradient(self.path, axis=0, edge_order=1)
        derivation2 = np.gradient(np.gradient(self.path, axis=0, edge_order=1), axis=0, edge_order=1)
        self.curvature = np.divide(
            (np.multiply(derivation1[:, 0], derivation2[:, 1]) - np.multiply(derivation1[:, 1], derivation2[:, 0])),
            np.power(np.sum(np.square(derivation1), axis=1), 3 / 2))

        self.radius = 1 / self.curvature

        self.velocity_max = np.array([])
        self.brake_velocity = np.array([])
        self.acceleration_velocity = np.array([])
        self.velocity = np.array([])
        self.apex = np.array([])
        self.delta_distance: npt.NDArray
        self.cum_distance: npt.NDArray
        self.v_initial = None

        self.mission_completed = mission_completed

    def limit_v_max_for_skidpad(self):
        """
        Limit the maximum allowed velocity for skidpad.

        Uses maximal allowed lateral acceleration for calculating the maximal allowed velocity.
        """
        skidpad_radius = 9.125
        self.v_max = min(np.sqrt(self.a_y_max * skidpad_radius), self.v_max)

    def calculate(self, v_initial: float = None, v_end: float = None) -> npt.NDArray:
        """
        Calculate a velocity profile for the previously set path.

        Calculates apex points with their respective maximal allowed velocities
        according to the maximal lateral velocity and their radius.
        Calculates velocities between support points (first, last and apex points)
        constraint by the maximal longitudinal velocity.

        Parameters
        ----------
        v_initial : float, optional
            Velocity for the first path point., by default None
        v_end : float, optional
            Velocity for the last path point., by default None

        Returns
        -------
        npt.NDArray
            Calculated velocity profile, shape: (n, ) with n equal number of path points.
        """
        if v_end is not None:
            self.v_end = v_end
        self.v_initial = v_initial
        # calculate real Distance of the track
        self.distance()

        if self.mission_completed:
            if self.v_initial is None:
                self.v_initial = self.v_max
            return self.brake()

        # calculate max longitudinal speed for each point
        self.find_longitudinal_speed()

        # find apex (local minimum of curvature)
        self.find_apex()

        # calculate forward path of acceleration (between to apexes)
        self.acceleration_forward()

        # calculate reverse path of acceleration
        self.acceleration_reverse()

        # combine both velocity-profiles
        velocity = self.real_speed()

        if self.stretch_time_for_mpc:
            total_time = np.sum(self.delta_distance / velocity)
            if total_time < self.mpc_min_time_horizont:
                velocity /= (self.mpc_min_time_horizont / total_time)

        return velocity

    def distance(self):
        """Calculate distance between each path points and cumulative distance between path points."""
        self.delta_distance = np.sqrt(np.sum(np.power(np.diff(self.path, axis=0, prepend=self.path[:1]), 2), axis=1))
        self.cum_distance = np.cumsum(self.delta_distance)

    def find_longitudinal_speed(self):
        """Calculate maximal possible longitudinal velocity based on radius and the maximal allowed lateral velocity."""
        self.velocity_max = np.sqrt(self.a_y_max * np.absolute(self.radius))
        self.velocity_max[np.isposinf(self.velocity_max)] = self.v_max
        if (self.v_initial is not None):
            self.velocity_max[0] = self.v_initial
        if (self.v_end is not None):
            self.velocity_max[-1] = self.v_end
        self.velocity_max = np.minimum(self.velocity_max, self.v_max)

    def find_apex(self):
        """Find apexes in path by finding peaks in the curvature."""
        # calculate max. curvature at which the car can drive with fullspeed
        curve_fullspeed = self.a_y_max / (self.v_max**2)

        # min. turning-radius of the vehicle is 3m (0.33), max. turning-radius of 25m (0.04)

        peaks, _ = find_peaks(np.absolute(self.curvature), height=[curve_fullspeed, 1])  # , distance=i_min)
        self.apex = peaks

    def acceleration_forward(self):
        """
        Accelerate forward from one support point to the next.

        .. todo:: This Code should be reprogrammed and the mathematic approaches documented.
            When doing so, you could also implement a GGV-Map
        """
        self.acceleration_velocity = np.copy(self.velocity_max)

        # itterate trough every Apex
        if (self.v_initial is not None):
            start = np.sort(np.append(self.apex, 0))
        else:
            start = self.apex
        for index in start:
            itteration = True
            while itteration:
                # find velocity of the  "last" point
                v_i = self.acceleration_velocity[index]

                # calculate lateral accel. with the velocity an the curvature at the "new" point
                a_y_i_plus = round(v_i**2 * np.absolute(self.curvature[index + 1]), 3)
                # a_y_i = round(v_i**2 * np.absolute(self.Curvature[index]), 3)  # noqa: E800

                # check if new lat. accel. is less then maximum accel. else stop
                if (a_y_i_plus <= self.a_y_max):

                    # calculate new long. accel. with the new lat. accel.
                    a_x_i_plus = np.sqrt(self.a_y_max**2 - a_y_i_plus**2)

                    # calculate the distance between both points
                    delta_s = self.cum_distance[index + 1] - self.cum_distance[index]

                    # caculate new velocity
                    v_i_plus = np.sqrt(2 * a_x_i_plus * delta_s + v_i**2)

                    # append new velocity to the result array
                    self.acceleration_velocity[index + 1] = v_i_plus

                    index += 1
                else:
                    itteration = False
                    break

                # check if itteration is still True
                if ((index >= self.path.shape[0] - 1) or (1.1 * self.velocity_max[index] < v_i_plus)):
                    itteration = False

                # solve Problem with velocitys
                elif ((1.1 * self.velocity_max[index] >= v_i_plus) and (self.velocity_max[index] < v_i_plus)):
                    self.acceleration_velocity[index] = self.velocity_max[index]

    def brake(self) -> npt.NDArray:
        """
        Calculate velocity profile by delecerating forward from first point.

        Uses maximal allowed longitudinal acceleration for braking.

        Returns
        -------
        npt.NDArray
            Velocity profile, shape: (n, ) with n equal number of path points.
        """
        brake_velocity = np.zeros_like(self.cum_distance)
        brake_velocity[0] = self.v_initial

        for index in range(brake_velocity.shape[0] - 1):
            velocity = brake_velocity[index]
            if velocity == 0:
                break

            a_y = round(velocity**2 * np.absolute(self.curvature[index]), 3)

            if a_y < self.a_x_max_completed:
                a_x = np.sqrt(self.a_x_max_completed**2 - a_y**2)
            else:
                a_x = 0

            brake_velocity[index + 1] = max(velocity - a_x * (self.delta_distance[index] / velocity), 0)

        return brake_velocity

    def acceleration_reverse(self):
        """
        Accelerate reverse from one support point to the previous.

        .. todo:: This Code should be reprogrammed and the mathematic approaches documented.
            When doing so, you could also implement a GGV-Map
        """
        self.brake_velocity = np.copy(self.velocity_max)

        # check for initial and end condition and append them to Apex
        if (self.v_end is not None):
            start = np.append(self.apex, len(self.brake_velocity) - 1)
        else:
            start = self.apex

        # itterate trough every Apex, backwards
        for index in np.flip(start, 0):
            itteration = True
            while itteration:
                # find velocity of the  "last" point
                v_i = self.brake_velocity[index]

                # calculate lateral accel. with the velocity an the curvature at the "new" point
                a_y_i_minus = round(v_i**2 * np.absolute(self.curvature[index - 1]), 3)
                # a_y_i = round(v_i**2 * np.absolute(self.Curvature[index]), 3)  # noqa: E800

                # check if new lat. accel. is less then maximum accel. else stop
                if (a_y_i_minus <= self.a_y_max):

                    # calculate new long. accel. with the new lat. accel.
                    a_x_i_minus = np.sqrt(self.a_y_max**2 - a_y_i_minus**2)

                    # calculate the distance between both points
                    delta_s = self.cum_distance[index] - self.cum_distance[index - 1]

                    # caculate new velocity
                    v_i_minus = np.sqrt(2 * a_x_i_minus * delta_s + v_i**2)

                    # append new velocity to the result array
                    self.brake_velocity[index - 1] = v_i_minus

                    index -= 1
                else:
                    itteration = False
                    break

                # check if itteration is still True
                if ((index < 1) or (1.1 * self.velocity_max[index] < v_i_minus)):
                    itteration = False

                # solve problem with velocitys
                elif ((1.1 * self.velocity_max[index] >= v_i_minus) and (self.velocity_max[index] < v_i_minus)):
                    self.brake_velocity[index] = self.velocity_max[index]

    def plot(self):
        """Plot velocity profile and their statistics."""
        import matplotlib.pyplot as plt
        fig, (ax1, ax3, ax4) = plt.subplots(3, 1)

        ax1.plot(self.cum_distance, self.velocity_max, label='Maximal possible velocity', color='orange', alpha=0.5)
        ax1.plot(self.cum_distance, self.acceleration_velocity, label='Acceleration velocity', color='g', alpha=0.5)
        ax1.plot(self.cum_distance, self.brake_velocity, label='Deceleration velocity', color='k', alpha=0.5)
        ax1.plot(self.cum_distance, self.curvature, label='Curvature', color='b', alpha=0.5)
        ax3.plot(self.cum_distance, self.radius, label='Radius', color='y', alpha=0.5)
        ax3.set_ylim(top=10, bottom=-10)
        ax3.grid()
        ax1.plot(self.cum_distance, self.velocity, label='Resulting velocity', color='r')
        ax1.scatter(self.cum_distance[self.apex], self.curvature[self.apex], label='Apexes', color='pink')
        ax4.plot(self.path[:, 0], self.path[:, 1])
        ax1.legend(loc=1)
        ax3.legend(loc=2)
        plt.show()

    def real_speed(self) -> npt.NDArray:
        """
        Calculate resulting speed by combining maximal velocity, brake velocity and acceleration velocity.

        Returns
        -------
        npt.NDArray
            Velocity profile, shape: (n, ) with n equal numbe rof path points.
        """
        self.velocity = np.minimum(self.velocity_max, self.brake_velocity)
        self.velocity = np.minimum(self.velocity, self.acceleration_velocity)

        return self.velocity