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