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