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