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    //   Δt = tₖ₊₁ − tₖ
163    //   τ = timestamp − tₖ
164    //
165    //   x(τ) = xₖ + vₖτ + 1/2 aₖτ² + 1/6 jₖτ³
166    //   v(τ) = vₖ + aₖτ + 1/2 jₖτ²
167    //   a(τ) = aₖ + jₖτ
168    //
169    // where jₖ = (aₖ₊₁ − aₖ)/Δt
170    double dt = endValue.t - this.t;
171    double τ = timestamp - this.t;
172    double τ2 = τ * τ;
173    double τ3 = τ * τ * τ;
174    double jx = (endValue.ax - this.ax) / dt;
175    double jy = (endValue.ay - this.ay) / dt;
176    double η = (endValue.alpha - this.alpha) / dt;
177    return new SwerveSample(
178        timestamp,
179        this.x + this.vx * τ + 0.5 * this.ax * τ2 + 1.0 / 6.0 * jx * τ3,
180        this.y + this.vy * τ + 0.5 * this.ay * τ2 + 1.0 / 6.0 * jy * τ3,
181        this.heading + this.omega * τ + 0.5 * this.alpha * τ2 + 1.0 / 6.0 * η * τ3,
182        this.vx + this.ax * τ + 0.5 * jx * τ2,
183        this.vy + this.ay * τ + 0.5 * jy * τ2,
184        this.omega + this.alpha * τ + 0.5 * η * τ2,
185        this.ax + jx * τ,
186        this.ay + jy * τ,
187        this.alpha + η * τ,
188        interp_fx,
189        interp_fy);
190  }
191
192  @Override
193  public SwerveSample offsetBy(double timestampOffset) {
194    return new SwerveSample(
195        this.t + timestampOffset,
196        this.x,
197        this.y,
198        this.heading,
199        this.vx,
200        this.vy,
201        this.omega,
202        this.ax,
203        this.ay,
204        this.alpha,
205        this.moduleForcesX(),
206        this.moduleForcesY());
207  }
208
209  @Override
210  public SwerveSample flipped() {
211    return switch (ChoreoAllianceFlipUtil.getFlipper()) {
212      case MIRRORED ->
213          new SwerveSample(
214              this.t,
215              ChoreoAllianceFlipUtil.flipX(this.x),
216              ChoreoAllianceFlipUtil.flipY(this.y),
217              ChoreoAllianceFlipUtil.flipHeading(this.heading),
218              -this.vx,
219              this.vy,
220              -this.omega,
221              -this.ax,
222              this.ay,
223              -this.alpha,
224              // FL, FR, BL, BR
225              // Mirrored
226              // -FR, -FL, -BR, -BL
227              new double[] {
228                -this.moduleForcesX()[1],
229                -this.moduleForcesX()[0],
230                -this.moduleForcesX()[3],
231                -this.moduleForcesX()[2]
232              },
233              // FL, FR, BL, BR
234              // Mirrored
235              // FR, FL, BR, BL
236              new double[] {
237                this.moduleForcesY()[1],
238                this.moduleForcesY()[0],
239                this.moduleForcesY()[3],
240                this.moduleForcesY()[2]
241              });
242      case ROTATE_AROUND ->
243          new SwerveSample(
244              this.t,
245              ChoreoAllianceFlipUtil.flipX(this.x),
246              ChoreoAllianceFlipUtil.flipY(this.y),
247              ChoreoAllianceFlipUtil.flipHeading(this.heading),
248              -this.vx,
249              -this.vy,
250              this.omega,
251              -this.ax,
252              -this.ay,
253              this.alpha,
254              Arrays.stream(this.moduleForcesX()).map(x -> -x).toArray(),
255              Arrays.stream(this.moduleForcesY()).map(y -> -y).toArray());
256    };
257  }
258
259  /** The struct for the SwerveSample class. */
260  public static final Struct<SwerveSample> struct = new SwerveSampleStruct();
261
262  private static final class SwerveSampleStruct implements Struct<SwerveSample> {
263    @Override
264    public Class<SwerveSample> getTypeClass() {
265      return SwerveSample.class;
266    }
267
268    @Override
269    public String getTypeName() {
270      return "SwerveSample";
271    }
272
273    @Override
274    public int getSize() {
275      return Struct.kSizeDouble * 18;
276    }
277
278    @Override
279    public String getSchema() {
280      return "double timestamp;"
281          + "Pose2d pose;"
282          + "double vx;"
283          + "double vy;"
284          + "double omega;"
285          + "double ax;"
286          + "double ay;"
287          + "double alpha;"
288          + "double moduleForcesX[4];"
289          + "double moduleForcesY[4];";
290    }
291
292    @Override
293    public Struct<?>[] getNested() {
294      return new Struct<?>[] {Pose2d.struct};
295    }
296
297    @Override
298    public SwerveSample unpack(ByteBuffer bb) {
299      return new SwerveSample(
300          bb.getDouble(),
301          bb.getDouble(),
302          bb.getDouble(),
303          bb.getDouble(),
304          bb.getDouble(),
305          bb.getDouble(),
306          bb.getDouble(),
307          bb.getDouble(),
308          bb.getDouble(),
309          bb.getDouble(),
310          new double[] {bb.getDouble(), bb.getDouble(), bb.getDouble(), bb.getDouble()},
311          new double[] {bb.getDouble(), bb.getDouble(), bb.getDouble(), bb.getDouble()});
312    }
313
314    @Override
315    public void pack(ByteBuffer bb, SwerveSample value) {
316      bb.putDouble(value.t);
317      bb.putDouble(value.x);
318      bb.putDouble(value.y);
319      bb.putDouble(value.heading);
320      bb.putDouble(value.vx);
321      bb.putDouble(value.vy);
322      bb.putDouble(value.omega);
323      bb.putDouble(value.ax);
324      bb.putDouble(value.ay);
325      bb.putDouble(value.alpha);
326      for (int i = 0; i < 4; ++i) {
327        bb.putDouble(value.moduleForcesX()[i]);
328      }
329      for (int i = 0; i < 4; ++i) {
330        bb.putDouble(value.moduleForcesY()[i]);
331      }
332    }
333  }
334
335  @Override
336  public boolean equals(Object obj) {
337    if (!(obj instanceof SwerveSample)) {
338      return false;
339    }
340
341    var other = (SwerveSample) obj;
342    return this.t == other.t
343        && this.x == other.x
344        && this.y == other.y
345        && this.heading == other.heading
346        && this.vx == other.vx
347        && this.vy == other.vy
348        && this.omega == other.omega
349        && this.ax == other.ax
350        && this.ay == other.ay
351        && this.alpha == other.alpha
352        && ChoreoArrayUtil.zipEquals(
353            this.fx, other.fx, (a, b) -> a.doubleValue() == b.doubleValue())
354        && ChoreoArrayUtil.zipEquals(
355            this.fy, other.fy, (a, b) -> a.doubleValue() == b.doubleValue());
356  }
357}