001// Copyright (c) Choreo contributors
002
003package choreo.trajectory;
004
005import choreo.util.ChoreoAllianceFlipUtil;
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;
012
013/** A single differential drive robot sample in a Trajectory. */
014public class DifferentialSample implements TrajectorySample<DifferentialSample> {
015  /** The timestamp of this sample relative to the beginning of the trajectory. */
016  public final double t;
017
018  /** The X position of the sample relative to the blue alliance wall origin in meters. */
019  public final double x;
020
021  /** The Y position of the sample relative to the blue alliance wall origin in meters. */
022  public final double y;
023
024  /** The heading of the sample in radians, with 0 being in the +X direction. */
025  public final double heading;
026
027  /** The velocity of the left side in m/s. */
028  public final double vl;
029
030  /** The velocity of the right side in m/s. */
031  public final double vr;
032
033  /** The chassis angular velocity in rad/s. */
034  public final double omega;
035
036  /** The acceleration of the left side in m/s². */
037  public final double al;
038
039  /** The acceleration of the right side in m/s². */
040  public final double ar;
041
042  /** The force of the left side in Newtons. */
043  public final double fl;
044
045  /** The force of the right side in Newtons. */
046  public final double fr;
047
048  /**
049   * Constructs a DifferentialSample with the specified parameters.
050   *
051   * @param timestamp The timestamp of this sample.
052   * @param x The X position of the sample in meters.
053   * @param y The Y position of the sample in meters.
054   * @param heading The heading of the sample in radians, with 0 being in the +X direction.
055   * @param vl The velocity of the left side in m/s.
056   * @param vr The velocity of the right side in m/s.
057   * @param omega The chassis angular velocity in rad/s.
058   * @param al The acceleration of the left side in m/s².
059   * @param ar The acceleration of the right side in m/s².
060   * @param fl The force of the left side in Newtons.
061   * @param fr The force of the right side in Newtons.
062   */
063  public DifferentialSample(
064      double timestamp,
065      double x,
066      double y,
067      double heading,
068      double vl,
069      double vr,
070      double omega,
071      double al,
072      double ar,
073      double fl,
074      double fr) {
075    this.t = timestamp;
076    this.x = x;
077    this.y = y;
078    this.heading = heading;
079    this.vl = vl;
080    this.vr = vr;
081    this.omega = omega;
082    this.al = al;
083    this.ar = ar;
084    this.fl = fl;
085    this.fr = fr;
086  }
087
088  @Override
089  public double getTimestamp() {
090    return t;
091  }
092
093  @Override
094  public Pose2d getPose() {
095    return new Pose2d(x, y, Rotation2d.fromRadians(heading));
096  }
097
098  /**
099   * Returns the field-relative chassis speeds of this sample.
100   *
101   * @return the field-relative chassis speeds of this sample.
102   * @see edu.wpi.first.math.kinematics.DifferentialDriveKinematics#toChassisSpeeds
103   */
104  @Override
105  public ChassisSpeeds getChassisSpeeds() {
106    return new ChassisSpeeds((vl + vr) / 2, 0, omega);
107  }
108
109  @Override
110  public DifferentialSample interpolate(DifferentialSample endValue, double timestamp) {
111    double scale = (timestamp - this.t) / (endValue.t - this.t);
112    var interp_pose = getPose().interpolate(endValue.getPose(), scale);
113
114    return new DifferentialSample(
115        MathUtil.interpolate(this.t, endValue.t, scale),
116        interp_pose.getX(),
117        interp_pose.getY(),
118        interp_pose.getRotation().getRadians(),
119        MathUtil.interpolate(this.vl, endValue.vl, scale),
120        MathUtil.interpolate(this.vr, endValue.vr, scale),
121        MathUtil.interpolate(this.omega, endValue.omega, scale),
122        MathUtil.interpolate(this.al, endValue.al, scale),
123        MathUtil.interpolate(this.ar, endValue.ar, scale),
124        MathUtil.interpolate(this.fl, endValue.fl, scale),
125        MathUtil.interpolate(this.fr, endValue.fr, scale));
126  }
127
128  public DifferentialSample flipped() {
129    return switch (ChoreoAllianceFlipUtil.getFlipper()) {
130      case MIRRORED ->
131          new DifferentialSample(
132              t,
133              ChoreoAllianceFlipUtil.flipX(x),
134              ChoreoAllianceFlipUtil.flipY(y), // No-op for mirroring
135              ChoreoAllianceFlipUtil.flipHeading(heading),
136              vr,
137              vl,
138              -omega,
139              ar,
140              al,
141              fr,
142              fl);
143      case ROTATE_AROUND ->
144          new DifferentialSample(
145              t,
146              ChoreoAllianceFlipUtil.flipX(x),
147              ChoreoAllianceFlipUtil.flipY(y),
148              ChoreoAllianceFlipUtil.flipHeading(heading),
149              vl,
150              vr,
151              omega,
152              al,
153              ar,
154              fl,
155              fr);
156    };
157  }
158
159  public DifferentialSample offsetBy(double timestampOffset) {
160    return new DifferentialSample(
161        t + timestampOffset, x, y, heading, vl, vr, omega, al, ar, fl, fr);
162  }
163
164  /** The struct for the DifferentialSample class. */
165  public static final Struct<DifferentialSample> struct = new DifferentialSampleStruct();
166
167  private static final class DifferentialSampleStruct implements Struct<DifferentialSample> {
168    @Override
169    public Class<DifferentialSample> getTypeClass() {
170      return DifferentialSample.class;
171    }
172
173    @Override
174    public String getTypeName() {
175      return "DifferentialSample";
176    }
177
178    @Override
179    public int getSize() {
180      return Struct.kSizeDouble * 10;
181    }
182
183    @Override
184    public String getSchema() {
185      return "double timestamp;"
186          + "Pose2d pose;"
187          + "double vl;"
188          + "double vr;"
189          + "double omega;"
190          + "double al;"
191          + "double ar;"
192          + "double fl;"
193          + "double fr;";
194    }
195
196    @Override
197    public Struct<?>[] getNested() {
198      return new Struct<?>[] {Pose2d.struct};
199    }
200
201    @Override
202    public DifferentialSample unpack(ByteBuffer bb) {
203      return new DifferentialSample(
204          bb.getDouble(),
205          bb.getDouble(),
206          bb.getDouble(),
207          bb.getDouble(),
208          bb.getDouble(),
209          bb.getDouble(),
210          bb.getDouble(),
211          bb.getDouble(),
212          bb.getDouble(),
213          bb.getDouble(),
214          bb.getDouble());
215    }
216
217    @Override
218    public void pack(ByteBuffer bb, DifferentialSample value) {
219      bb.putDouble(value.t);
220      bb.putDouble(value.x);
221      bb.putDouble(value.y);
222      bb.putDouble(value.heading);
223      bb.putDouble(value.vl);
224      bb.putDouble(value.vr);
225      bb.putDouble(value.omega);
226      bb.putDouble(value.al);
227      bb.putDouble(value.ar);
228      bb.putDouble(value.fl);
229      bb.putDouble(value.fr);
230    }
231  }
232
233  @Override
234  public boolean equals(Object obj) {
235    if (!(obj instanceof DifferentialSample)) {
236      return false;
237    }
238
239    var other = (DifferentialSample) obj;
240    return this.t == other.t
241        && this.x == other.x
242        && this.y == other.y
243        && this.heading == other.heading
244        && this.vl == other.vl
245        && this.vr == other.vr
246        && this.omega == other.omega
247        && this.al == other.al
248        && this.ar == other.ar
249        && this.fl == other.fl
250        && this.fr == other.fr;
251  }
252}