001// Copyright (c) Choreo contributors
002
003package choreo.trajectory;
004
005import choreo.util.ChoreoAllianceFlipUtil;
006import choreo.util.ChoreoArrayUtil;
007import java.nio.ByteBuffer;
008import org.wpilib.math.geometry.Pose2d;
009import org.wpilib.math.geometry.Rotation2d;
010import org.wpilib.math.kinematics.ChassisVelocities;
011import org.wpilib.math.util.MathUtil;
012import org.wpilib.util.struct.Struct;
013
014/** A single swerve robot sample in a Trajectory. */
015public class SwerveSample implements TrajectorySample<SwerveSample> {
016  private static final double[] EMPTY_MODULE_FORCES = new double[] {0, 0, 0, 0};
017
018  /** The timestamp of this sample, relative to the beginning of the trajectory. */
019  public final double t;
020
021  /** The X position of the sample relative to the blue alliance wall origin in meters. */
022  public final double x;
023
024  /** The Y position of the sample relative to the blue alliance wall origin in meters. */
025  public final double y;
026
027  /** The heading of the sample in radians, with 0 being in the +X direction. */
028  public final double heading;
029
030  /** The velocity of the sample in the X direction in m/s. */
031  public final double vx;
032
033  /** The velocity of the sample in the Y direction in m/s. */
034  public final double vy;
035
036  /** The angular velocity of the sample in rad/s. */
037  public final double omega;
038
039  /** The acceleration of the sample in the X direction in m/s². */
040  public final double ax;
041
042  /** The acceleration of the sample in the Y direction in m/s². */
043  public final double ay;
044
045  /** The angular acceleration of the sample in rad/s². */
046  public final double alpha;
047
048  /**
049   * The force on each swerve module in the X direction in Newtons. Module forces appear in the
050   * following order: [FL, FR, BL, BR].
051   */
052  private final double[] fx;
053
054  /**
055   * The force on each swerve module in the Y direction in Newtons Module forces appear in the
056   * following order: [FL, FR, BL, BR].
057   */
058  private final double[] fy;
059
060  /**
061   * Constructs a SwerveSample with the specified parameters.
062   *
063   * @param t The timestamp of this sample, relative to the beginning of the trajectory.
064   * @param x The X position of the sample in meters.
065   * @param y The Y position of the sample in meters.
066   * @param heading The heading of the sample in radians, with 0 being in the +X direction.
067   * @param vx The velocity of the sample in the X direction in m/s.
068   * @param vy The velocity of the sample in the Y direction in m/s.
069   * @param omega The angular velocity of the sample in rad/s.
070   * @param ax The acceleration of the sample in the X direction in m/s².
071   * @param ay The acceleration of the sample in the Y direction in m/s².
072   * @param alpha The angular acceleration of the sample in rad/s².
073   * @param moduleForcesX The force on each swerve module in the X direction in Newtons. Module
074   *     forces appear in the following order: [FL, FR, BL, BR].
075   * @param moduleForcesY The force on each swerve module in the Y direction in Newtons. Module
076   *     forces appear in the following order: [FL, FR, BL, BR].
077   */
078  public SwerveSample(
079      double t,
080      double x,
081      double y,
082      double heading,
083      double vx,
084      double vy,
085      double omega,
086      double ax,
087      double ay,
088      double alpha,
089      double[] moduleForcesX,
090      double[] moduleForcesY) {
091    this.t = t;
092    this.x = x;
093    this.y = y;
094    this.heading = heading;
095    this.vx = vx;
096    this.vy = vy;
097    this.omega = omega;
098    this.ax = ax;
099    this.ay = ay;
100    this.alpha = alpha;
101    this.fx = moduleForcesX;
102    this.fy = moduleForcesY;
103  }
104
105  /**
106   * A null safe getter for the module forces in the X direction.
107   *
108   * @return The module forces in the X direction.
109   */
110  public double[] moduleForcesX() {
111    if (fx == null || fx.length != 4) {
112      return EMPTY_MODULE_FORCES;
113    }
114    return fx;
115  }
116
117  /**
118   * A null safe getter for the module forces in the Y direction.
119   *
120   * @return The module forces in the Y direction.
121   */
122  public double[] moduleForcesY() {
123    if (fy == null || fy.length != 4) {
124      return EMPTY_MODULE_FORCES;
125    }
126    return fy;
127  }
128
129  @Override
130  public double getTimestamp() {
131    return t;
132  }
133
134  @Override
135  public Pose2d getPose() {
136    return new Pose2d(x, y, Rotation2d.fromRadians(heading));
137  }
138
139  @Override
140  public ChassisVelocities getChassisVelocities() {
141    return new ChassisVelocities(vx, vy, omega);
142  }
143
144  @Override
145  public SwerveSample interpolate(SwerveSample endValue, double timestamp) {
146    double scale = (timestamp - this.t) / (endValue.t - this.t);
147
148    double[] interp_fx = new double[4];
149    double[] interp_fy = new double[4];
150    for (int i = 0; i < 4; ++i) {
151      interp_fx[i] = MathUtil.lerp(this.moduleForcesX()[i], endValue.moduleForcesX()[i], scale);
152      interp_fy[i] = MathUtil.lerp(this.moduleForcesY()[i], endValue.moduleForcesY()[i], scale);
153    }
154
155    // Integrate the acceleration to get the rest of the state, since linearly
156    // interpolating the state gives an inaccurate result if the accelerations are changing between
157    // states
158    //
159    //   τ = timestamp − tₖ
160    //
161    //   x(τ) = xₖ + vₖτ + 1/2 aₖτ²
162    //   v(τ) = vₖ + aₖτ
163    double τ = timestamp - this.t;
164    double τ2 = τ * τ;
165    return new SwerveSample(
166        timestamp,
167        this.x + this.vx * τ + 0.5 * this.ax * τ2,
168        this.y + this.vy * τ + 0.5 * this.ay * τ2,
169        this.heading + this.omega * τ + 0.5 * this.alpha * τ2,
170        this.vx + this.ax * τ,
171        this.vy + this.ay * τ,
172        this.omega + this.alpha * τ,
173        this.ax,
174        this.ay,
175        this.alpha,
176        interp_fx,
177        interp_fy);
178  }
179
180  @Override
181  public SwerveSample offsetBy(double timestampOffset) {
182    return new SwerveSample(
183        this.t + timestampOffset,
184        this.x,
185        this.y,
186        this.heading,
187        this.vx,
188        this.vy,
189        this.omega,
190        this.ax,
191        this.ay,
192        this.alpha,
193        this.moduleForcesX(),
194        this.moduleForcesY());
195  }
196
197  @Override
198  public SwerveSample flipped() {
199    return ChoreoAllianceFlipUtil.flip(this);
200  }
201
202  @Override
203  public SwerveSample mirrorX() {
204    return ChoreoAllianceFlipUtil.getMirrorX().flip(this);
205  }
206
207  @Override
208  public SwerveSample mirrorY() {
209    return ChoreoAllianceFlipUtil.getMirrorY().flip(this);
210  }
211
212  @Override
213  public SwerveSample rotateAround() {
214    return ChoreoAllianceFlipUtil.getRotateAround().flip(this);
215  }
216
217  /** The struct for the SwerveSample class. */
218  public static final Struct<SwerveSample> struct = new SwerveSampleStruct();
219
220  private static final class SwerveSampleStruct implements Struct<SwerveSample> {
221    @Override
222    public Class<SwerveSample> getTypeClass() {
223      return SwerveSample.class;
224    }
225
226    @Override
227    public String getTypeName() {
228      return "SwerveSample";
229    }
230
231    @Override
232    public int getSize() {
233      return Double.BYTES * 18;
234    }
235
236    @Override
237    public String getSchema() {
238      return "double timestamp;"
239          + "Pose2d pose;"
240          + "double vx;"
241          + "double vy;"
242          + "double omega;"
243          + "double ax;"
244          + "double ay;"
245          + "double alpha;"
246          + "double moduleForcesX[4];"
247          + "double moduleForcesY[4];";
248    }
249
250    @Override
251    public Struct<?>[] getNested() {
252      return new Struct<?>[] {Pose2d.struct};
253    }
254
255    @Override
256    public SwerveSample unpack(ByteBuffer bb) {
257      return new SwerveSample(
258          bb.getDouble(),
259          bb.getDouble(),
260          bb.getDouble(),
261          bb.getDouble(),
262          bb.getDouble(),
263          bb.getDouble(),
264          bb.getDouble(),
265          bb.getDouble(),
266          bb.getDouble(),
267          bb.getDouble(),
268          new double[] {bb.getDouble(), bb.getDouble(), bb.getDouble(), bb.getDouble()},
269          new double[] {bb.getDouble(), bb.getDouble(), bb.getDouble(), bb.getDouble()});
270    }
271
272    @Override
273    public void pack(ByteBuffer bb, SwerveSample value) {
274      bb.putDouble(value.t);
275      bb.putDouble(value.x);
276      bb.putDouble(value.y);
277      bb.putDouble(value.heading);
278      bb.putDouble(value.vx);
279      bb.putDouble(value.vy);
280      bb.putDouble(value.omega);
281      bb.putDouble(value.ax);
282      bb.putDouble(value.ay);
283      bb.putDouble(value.alpha);
284      for (int i = 0; i < 4; ++i) {
285        bb.putDouble(value.moduleForcesX()[i]);
286      }
287      for (int i = 0; i < 4; ++i) {
288        bb.putDouble(value.moduleForcesY()[i]);
289      }
290    }
291  }
292
293  @Override
294  public boolean equals(Object obj) {
295    if (!(obj instanceof SwerveSample)) {
296      return false;
297    }
298
299    var other = (SwerveSample) obj;
300    return MathUtil.isNear(this.t, other.t, 1E-6)
301        && MathUtil.isNear(this.x, other.x, 1E-6)
302        && MathUtil.isNear(this.y, other.y, 1E-6)
303        && MathUtil.isNear(this.heading, other.heading, 1E-6)
304        && MathUtil.isNear(this.vx, other.vx, 1E-6)
305        && MathUtil.isNear(this.vy, other.vy, 1E-6)
306        && MathUtil.isNear(this.omega, other.omega, 1E-6)
307        && MathUtil.isNear(this.ax, other.ax, 1E-6)
308        && MathUtil.isNear(this.ay, other.ay, 1E-6)
309        && MathUtil.isNear(this.alpha, other.alpha, 1E-6)
310        && ChoreoArrayUtil.zipEquals(
311            this.fx, other.fx, (a, b) -> MathUtil.isNear(a.doubleValue(), b.doubleValue(), 1E-6))
312        && ChoreoArrayUtil.zipEquals(
313            this.fy, other.fy, (a, b) -> MathUtil.isNear(a.doubleValue(), b.doubleValue(), 1E-6));
314  }
315}