allwpilib/wpimath/src/main/native/include/frc/MathUtil.h

// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.

#pragma once

#include <numbers>
#include <type_traits>

#include <gcem.hpp>
#include <wpi/SymbolExports.h>

#include "frc/geometry/Translation2d.h"
#include "frc/geometry/Translation3d.h"
#include "units/angle.h"
#include "units/base.h"
#include "units/length.h"
#include "units/math.h"
#include "units/time.h"
#include "units/velocity.h"
#include "wpimath/MathShared.h"

namespace frc {

/**
 * Returns 0.0 if the given value is within the specified range around zero. The
 * remaining range between the deadband and the maximum magnitude is scaled from
 * 0.0 to the maximum magnitude.
 *
 * @param value Value to clip.
 * @param deadband Range around zero.
 * @param maxMagnitude The maximum magnitude of the input (defaults to 1). Can
 * be infinite.
 * @return The value after the deadband is applied.
 */
template <typename T>
  requires std::is_arithmetic_v<T> || units::traits::is_unit_t_v<T>
constexpr T ApplyDeadband(T value, T deadband, T maxMagnitude = T{1.0}) {
  T magnitude;
  if constexpr (std::is_arithmetic_v<T>) {
    magnitude = gcem::abs(value);
  } else {
    magnitude = units::math::abs(value);
  }

  if (magnitude < deadband) {
    return T{0.0};
  }

  if (value > T{0.0}) {
    // Map deadband to 0 and map max to max with a linear relationship.
    //
    //   y - y₁ = m(x - x₁)
    //   y - y₁ = (y₂ - y₁)/(x₂ - x₁) (x - x₁)
    //   y = (y₂ - y₁)/(x₂ - x₁) (x - x₁) + y₁
    //
    // (x₁, y₁) = (deadband, 0) and (x₂, y₂) = (max, max).
    //
    //   x₁ = deadband
    //   y₁ = 0
    //   x₂ = max
    //   y₂ = max
    //   y = (max - 0)/(max - deadband) (x - deadband) + 0
    //   y = max/(max - deadband) (x - deadband)
    //
    // To handle high values of max, rewrite so that max only appears on the
    // denominator.
    //
    //   y = ((max - deadband) + deadband)/(max - deadband) (x - deadband)
    //   y = (1 + deadband/(max - deadband)) (x - deadband)
    return (1.0 + deadband / (maxMagnitude - deadband)) * (value - deadband);
  } else {
    // Map -deadband to 0 and map -max to -max with a linear relationship.
    //
    //   y - y₁ = m(x - x₁)
    //   y - y₁ = (y₂ - y₁)/(x₂ - x₁) (x - x₁)
    //   y = (y₂ - y₁)/(x₂ - x₁) (x - x₁) + y₁
    //
    // (x₁, y₁) = (-deadband, 0) and (x₂, y₂) = (-max, -max).
    //
    //   x₁ = -deadband
    //   y₁ = 0
    //   x₂ = -max
    //   y₂ = -max
    //   y = (-max - 0)/(-max + deadband) (x + deadband) + 0
    //   y = max/(max - deadband) (x + deadband)
    //
    // To handle high values of max, rewrite so that max only appears on the
    // denominator.
    //
    //   y = ((max - deadband) + deadband)/(max - deadband) (x + deadband)
    //   y = (1 + deadband/(max - deadband)) (x + deadband)
    return (1.0 + deadband / (maxMagnitude - deadband)) * (value + deadband);
  }
}

/**
 * Returns a zero vector if the given vector is within the specified
 * distance from the origin. The remaining distance between the deadband and the
 * maximum distance is scaled from the origin to the maximum distance.
 *
 * @param value Value to clip.
 * @param deadband Distance from origin.
 * @param maxMagnitude The maximum distance from the origin of the input
 * (defaults to 1). Can be infinite.
 * @return The value after the deadband is applied.
 */
template <typename T, int N>
  requires std::is_arithmetic_v<T> || units::traits::is_unit_t_v<T>
Eigen::Vector<T, N> ApplyDeadband(const Eigen::Vector<T, N>& value, T deadband,
                                  T maxMagnitude = T{1.0}) {
  if constexpr (std::is_arithmetic_v<T>) {
    if (value.norm() < T{1e-9}) {
      return Eigen::Vector<T, N>::Zero();
    }
    return value.normalized() *
           ApplyDeadband(value.norm(), deadband, maxMagnitude);
  } else {
    const Eigen::Vector<double, N> asDouble = value.template cast<double>();
    const Eigen::Vector<double, N> processed =
        ApplyDeadband(asDouble, deadband.value(), maxMagnitude.value());
    return processed.template cast<T>();
  }
}

/**
 * Raises the input to the power of the given exponent while preserving its
 * sign.
 *
 * The function normalizes the input value to the range [0, 1] based on the
 * maximum magnitude so that the output stays in the range.
 *
 * This is useful for applying smoother or more aggressive control response
 * curves (e.g. joystick input shaping).
 *
 * @param value The input value to transform.
 * @param exponent The exponent to apply (e.g. 1.0 = linear, 2.0 = squared
 * curve). Must be positive.
 * @param maxMagnitude The maximum expected absolute value of input (defaults to
 * 1). Must be positive.
 * @return The transformed value with the same sign and scaled to the input
 * range.
 */
template <typename T>
  requires std::is_arithmetic_v<T> || units::traits::is_unit_t_v<T>
constexpr T CopyDirectionPow(T value, double exponent,
                             T maxMagnitude = T{1.0}) {
  if constexpr (std::is_arithmetic_v<T>) {
    return gcem::copysign(
        gcem::pow(gcem::abs(value) / maxMagnitude, exponent) * maxMagnitude,
        value);
  } else {
    return units::math::copysign(
        gcem::pow((units::math::abs(value) / maxMagnitude).value(), exponent) *
            maxMagnitude,
        value);
  }
}

/**
 * Raises the norm of the input to the power of the given exponent while
 * preserving its direction.
 *
 * The function normalizes the input value to the range [0, 1] based on the
 * maximum magnitude so that the output stays in the range.
 *
 * This is useful for applying smoother or more aggressive control response
 * curves (e.g. joystick input shaping).
 *
 * @param value The input vector to transform.
 * @param exponent The exponent to apply (e.g. 1.0 = linear, 2.0 = squared
 * curve). Must be positive.
 * @param maxMagnitude The maximum expected distance from origin of input
 * (defaults to 1). Must be positive.
 * @return The transformed value with the same direction and norm scaled to
 * the input range.
 */
template <typename T, int N>
  requires std::is_arithmetic_v<T> || units::traits::is_unit_t_v<T>
Eigen::Vector<T, N> CopyDirectionPow(const Eigen::Vector<T, N>& value,
                                     double exponent, T maxMagnitude = T{1.0}) {
  if constexpr (std::is_arithmetic_v<T>) {
    if (value.norm() < T{1e-9}) {
      return Eigen::Vector<T, N>::Zero();
    }
    return value.normalized() *
           CopyDirectionPow(value.norm(), exponent, maxMagnitude);
  } else {
    const Eigen::Vector<double, N> asDouble = value.template cast<double>();
    const Eigen::Vector<double, N> processed =
        CopyDirectionPow(asDouble, exponent, maxMagnitude.value());
    return processed.template cast<T>();
  }
}

/**
 * Returns modulus of input.
 *
 * @param input        Input value to wrap.
 * @param minimumInput The minimum value expected from the input.
 * @param maximumInput The maximum value expected from the input.
 */
template <typename T>
constexpr T InputModulus(T input, T minimumInput, T maximumInput) {
  T modulus = maximumInput - minimumInput;

  // Wrap input if it's above the maximum input
  int numMax = (input - minimumInput) / modulus;
  input -= numMax * modulus;

  // Wrap input if it's below the minimum input
  int numMin = (input - maximumInput) / modulus;
  input -= numMin * modulus;

  return input;
}

/**
 * Checks if the given value matches an expected value within a certain
 * tolerance.
 *
 * @param expected The expected value
 * @param actual The actual value
 * @param tolerance The allowed difference between the actual and the expected
 * value
 * @return Whether or not the actual value is within the allowed tolerance
 */
template <typename T>
  requires std::is_arithmetic_v<T> || units::traits::is_unit_t_v<T>
constexpr bool IsNear(T expected, T actual, T tolerance) {
  if constexpr (std::is_arithmetic_v<T>) {
    return std::abs(expected - actual) < tolerance;
  } else {
    return units::math::abs(expected - actual) < tolerance;
  }
}

/**
 * Checks if the given value matches an expected value within a certain
 * tolerance. Supports continuous input for cases like absolute encoders.
 *
 * Continuous input means that the min and max value are considered to be the
 * same point, and tolerances can be checked across them. A common example
 * would be for absolute encoders: calling isNear(2, 359, 5, 0, 360) returns
 * true because 359 is 1 away from 360 (which is treated as the same as 0) and
 * 2 is 2 away from 0, adding up to an error of 3 degrees, which is within the
 * given tolerance of 5.
 *
 * @param expected The expected value
 * @param actual The actual value
 * @param tolerance The allowed difference between the actual and the expected
 * value
 * @param min Smallest value before wrapping around to the largest value
 * @param max Largest value before wrapping around to the smallest value
 * @return Whether or not the actual value is within the allowed tolerance
 */
template <typename T>
  requires std::is_arithmetic_v<T> || units::traits::is_unit_t_v<T>
constexpr bool IsNear(T expected, T actual, T tolerance, T min, T max) {
  T errorBound = (max - min) / 2.0;
  T error = frc::InputModulus<T>(expected - actual, -errorBound, errorBound);

  if constexpr (std::is_arithmetic_v<T>) {
    return std::abs(error) < tolerance;
  } else {
    return units::math::abs(error) < tolerance;
  }
}

/**
 * Wraps an angle to the range -π to π radians (-180 to 180 degrees).
 *
 * @param angle Angle to wrap.
 */
WPILIB_DLLEXPORT
constexpr units::radian_t AngleModulus(units::radian_t angle) {
  return InputModulus<units::radian_t>(angle,
                                       units::radian_t{-std::numbers::pi},
                                       units::radian_t{std::numbers::pi});
}

// floorDiv and floorMod algorithms taken from Java

/**
 * Returns the largest (closest to positive infinity)
 * {@code int} value that is less than or equal to the algebraic quotient.
 *
 * @param x the dividend
 * @param y the divisor
 * @return the largest (closest to positive infinity)
 *   {@code int} value that is less than or equal to the algebraic quotient.
 */
constexpr std::signed_integral auto FloorDiv(std::signed_integral auto x,
                                             std::signed_integral auto y) {
  auto quot = x / y;
  auto rem = x % y;
  // if the signs are different and modulo not zero, round down
  if ((x < 0) != (y < 0) && rem != 0) {
    --quot;
  }
  return quot;
}

/**
 * Returns the floor modulus of the {@code int} arguments.
 * <p>
 * The floor modulus is {@code r = x - (floorDiv(x, y) * y)},
 * has the same sign as the divisor {@code y} or is zero, and
 * is in the range of {@code -std::abs(y) < r < +std::abs(y)}.
 *
 * @param x the dividend
 * @param y the divisor
 * @return the floor modulus {@code x - (floorDiv(x, y) * y)}
 */
constexpr std::signed_integral auto FloorMod(std::signed_integral auto x,
                                             std::signed_integral auto y) {
  return x - FloorDiv(x, y) * y;
}

/**
 * Limits translation velocity.
 *
 * @param current Translation at current timestep.
 * @param next Translation at next timestep.
 * @param dt Timestep duration.
 * @param maxVelocity Maximum translation velocity.
 * @return Returns the next Translation2d limited to maxVelocity
 */
constexpr Translation2d SlewRateLimit(const Translation2d& current,
                                      const Translation2d& next,
                                      units::second_t dt,
                                      units::meters_per_second_t maxVelocity) {
  if (maxVelocity < 0_mps) {
    wpi::math::MathSharedStore::ReportError(
        "maxVelocity must be a non-negative number, got {}!", maxVelocity);
    return next;
  }
  Translation2d diff = next - current;
  units::meter_t dist = diff.Norm();
  if (dist < 1e-9_m) {
    return next;
  }
  if (dist > maxVelocity * dt) {
    // Move maximum allowed amount in direction of the difference
    // NOLINTNEXTLINE(bugprone-integer-division)
    return current + diff * (maxVelocity * dt / dist);
  }
  return next;
}

/**
 * Limits translation velocity.
 *
 * @param current Translation at current timestep.
 * @param next Translation at next timestep.
 * @param dt Timestep duration.
 * @param maxVelocity Maximum translation velocity.
 * @return Returns the next Translation3d limited to maxVelocity
 */
constexpr Translation3d SlewRateLimit(const Translation3d& current,
                                      const Translation3d& next,
                                      units::second_t dt,
                                      units::meters_per_second_t maxVelocity) {
  if (maxVelocity < 0_mps) {
    wpi::math::MathSharedStore::ReportError(
        "maxVelocity must be a non-negative number, got {}!", maxVelocity);
    return next;
  }
  Translation3d diff = next - current;
  units::meter_t dist = diff.Norm();
  if (dist < 1e-9_m) {
    return next;
  }
  if (dist > maxVelocity * dt) {
    // Move maximum allowed amount in direction of the difference
    // NOLINTNEXTLINE(bugprone-integer-division)
    return current + diff * (maxVelocity * dt / dist);
  }
  return next;
}

}  // namespace frc