SwerveSample.h

// Copyright (c) Choreo contributors
 
#pragma once
 
#include <algorithm>
#include <array>
#include <type_traits>
 
#include <frc/kinematics/ChassisSpeeds.h>
#include <units/acceleration.h>
#include <units/angle.h>
#include <units/angular_acceleration.h>
#include <units/angular_velocity.h>
#include <units/force.h>
#include <units/length.h>
#include <units/time.h>
#include <units/velocity.h>
#include <wpi/MathExtras.h>
#include <wpi/json_fwd.h>
 
#include "choreo/util/AllianceFlipperUtil.h"
 
namespace choreo {
 
class SwerveSample {
 public:
  constexpr SwerveSample() = default;
 
  constexpr SwerveSample(units::second_t timestamp, units::meter_t x,
                         units::meter_t y, units::radian_t heading,
                         units::meters_per_second_t vx,
                         units::meters_per_second_t vy,
                         units::radians_per_second_t omega,
                         units::meters_per_second_squared_t ax,
                         units::meters_per_second_squared_t ay,
                         units::radians_per_second_squared_t alpha,
                         std::array<units::newton_t, 4> moduleForcesX,
                         std::array<units::newton_t, 4> moduleForcesY)
      : timestamp{timestamp},
        x{x},
        y{y},
        heading{heading},
        vx{vx},
        vy{vy},
        omega{omega},
        ax{ax},
        ay{ay},
        alpha{alpha},
        moduleForcesX{moduleForcesX},
        moduleForcesY{moduleForcesY} {}
 
  constexpr units::second_t GetTimestamp() const { return timestamp; }
 
  constexpr frc::Pose2d GetPose() const {
    return frc::Pose2d{x, y, frc::Rotation2d{heading}};
  }
 
  constexpr frc::ChassisSpeeds GetChassisSpeeds() const {
    return frc::ChassisSpeeds{vx, vy, omega};
  }
 
  template <int Year = util::kDefaultYear>
  constexpr SwerveSample Flipped() const {
    constexpr auto flipper = choreo::util::GetFlipperForYear<Year>();
    if constexpr (flipper.isMirrored) {
      return SwerveSample{timestamp,
                          flipper.FlipX(x),
                          flipper.FlipY(y),
                          flipper.FlipHeading(heading),
                          -vx,
                          vy,
                          -omega,
                          -ax,
                          ay,
                          -alpha,
                          // FL, FR, BL, BR
                          // Mirrored
                          // -FR, -FL, -BR, -BL
                          {-moduleForcesX[1], -moduleForcesX[0],
                           -moduleForcesX[3], -moduleForcesX[2]},
                          // FL, FR, BL, BR
                          // Mirrored
                          // -FR, -FL, -BR, -BL
                          {moduleForcesY[1], moduleForcesY[0], moduleForcesY[3],
                           moduleForcesY[2]}};
    } else {
      return SwerveSample{timestamp,
                          flipper.FlipX(x),
                          flipper.FlipY(y),
                          flipper.FlipHeading(heading),
                          -vx,
                          -vy,
                          omega,
                          -ax,
                          -ay,
                          alpha,
                          {-moduleForcesX[0], -moduleForcesX[1],
                           -moduleForcesX[2], -moduleForcesX[3]},
                          {-moduleForcesY[0], -moduleForcesY[1],
                           -moduleForcesY[2], -moduleForcesY[3]}};
    }
  }
 
  constexpr SwerveSample OffsetBy(units::second_t timeStampOffset) const {
    return SwerveSample{timestamp + timeStampOffset,
                        x,
                        y,
                        heading,
                        vx,
                        vy,
                        omega,
                        ax,
                        ay,
                        alpha,
                        moduleForcesX,
                        moduleForcesY};
  }
 
  constexpr SwerveSample Interpolate(const SwerveSample& endValue,
                                     units::second_t t) const {
    units::scalar_t scale = (t - timestamp) / (endValue.timestamp - timestamp);
 
    std::array<units::newton_t, 4> interpolatedForcesX;
    std::array<units::newton_t, 4> interpolatedForcesY;
    for (int i = 0; i < 4; i++) {
      interpolatedForcesX[i] =
          wpi::Lerp(moduleForcesX[i], endValue.moduleForcesX[i], scale.value());
      interpolatedForcesY[i] =
          wpi::Lerp(moduleForcesY[i], endValue.moduleForcesY[i], scale.value());
    }
 
    // Integrate the acceleration to get the rest of the state, since linearly
    // interpolating the state gives an inaccurate result if the accelerations
    // are changing between states
    //
    //   Δt = tₖ₊₁ − tₖ
    //   τ = timestamp − tₖ
    //
    //   x(τ) = xₖ + vₖτ + 1/2 aₖτ² + 1/6 jₖτ³
    //   v(τ) = vₖ + aₖτ + 1/2 jₖτ²
    //   a(τ) = aₖ + jₖτ
    //
    // where jₖ = (aₖ₊₁ − aₖ)/Δt
    auto dt = endValue.timestamp - timestamp;
    auto τ = t - timestamp;
    auto τ2 = τ * τ;
    auto τ3 = τ * τ * τ;
    auto jx = (endValue.ax - ax) / dt;
    auto jy = (endValue.ay - ay) / dt;
    auto η = (endValue.alpha - alpha) / dt;
    return SwerveSample{
        wpi::Lerp(timestamp, endValue.timestamp, scale),
        x + vx * τ + 0.5 * ax * τ2 + 1.0 / 6.0 * jx * τ3,
        y + vy * τ + 0.5 * ay * τ2 + 1.0 / 6.0 * jy * τ3,
        heading + omega * τ + 0.5 * alpha * τ2 + 1.0 / 6.0 * η * τ3,
        vx + ax * τ + 0.5 * jx * τ2,
        vy + ay * τ + 0.5 * jy * τ2,
        omega + alpha * τ + 0.5 * η * τ2,
        ax + jx * τ,
        ay + jy * τ,
        alpha + η * τ,
        interpolatedForcesX,
        interpolatedForcesY};
  }
 
  constexpr bool operator==(const SwerveSample& other) const {
    constexpr double epsilon = 1e-6;
 
    auto compare_units = [epsilon](const auto& a, const auto& b) {
      using UnitType =
          std::remove_const_t<std::remove_reference_t<decltype(a)>>;
      return units::math::abs(a - b) < UnitType(epsilon);
    };
 
    auto compare_arrays = [&compare_units](const auto& arr1, const auto& arr2) {
      return std::equal(arr1.begin(), arr1.end(), arr2.begin(), compare_units);
    };
 
    return compare_units(timestamp, other.timestamp) &&
           compare_units(x, other.x) && compare_units(y, other.y) &&
           compare_units(heading, other.heading) &&
           compare_units(vx, other.vx) && compare_units(vy, other.vy) &&
           compare_units(omega, other.omega) && compare_units(ax, other.ax) &&
           compare_units(ay, other.ay) && compare_units(alpha, other.alpha) &&
           compare_arrays(moduleForcesX, other.moduleForcesX) &&
           compare_arrays(moduleForcesY, other.moduleForcesY);
  }
 
  units::second_t timestamp = 0_s;
 
  units::meter_t x = 0_m;
 
  units::meter_t y = 0_m;
 
  units::radian_t heading = 0_rad;
 
  units::meters_per_second_t vx = 0_mps;
 
  units::meters_per_second_t vy = 0_mps;
 
  units::radians_per_second_t omega = 0_rad_per_s;
 
  units::meters_per_second_squared_t ax = 0_mps_sq;
 
  units::meters_per_second_squared_t ay = 0_mps_sq;
 
  units::radians_per_second_squared_t alpha = 0_rad_per_s_sq;
 
  std::array<units::newton_t, 4> moduleForcesX{0_N, 0_N, 0_N, 0_N};
 
  std::array<units::newton_t, 4> moduleForcesY{0_N, 0_N, 0_N, 0_N};
};
 
void to_json(wpi::json& json, const SwerveSample& trajectorySample);
void from_json(const wpi::json& json, SwerveSample& trajectorySample);
 
}  // namespace choreo
 
#include "choreo/trajectory/struct/SwerveSampleStruct.h"