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