001package com.stuypulse.robot.util.simulation;
002
003import static edu.wpi.first.units.Units.*;
004
005import java.nio.ByteBuffer;
006
007import edu.wpi.first.util.struct.Struct;
008import edu.wpi.first.math.geometry.Pose2d;
009import edu.wpi.first.math.geometry.Pose3d;
010import edu.wpi.first.math.geometry.Rotation2d;
011import edu.wpi.first.math.geometry.Rotation3d;
012import edu.wpi.first.math.geometry.Translation2d;
013import edu.wpi.first.math.geometry.Translation3d;
014import edu.wpi.first.units.measure.*;
015import edu.wpi.first.util.struct.StructSerializable;
016
017/**
018 *
019 *
020 * <h2>Record that holds CAD Offsets</h2>
021 *
022 * <p>
023 * Sourced from CAD exports to hold a component's positional offsets from their
024 * position in CAD
025 * to their position in sim
026 *
027 * <pre>{@code
028 * Offsets SHOOTER = new Offsets(-0.1016, 0.2032, 0.3255,
029 *         Degrees.of(90), Degrees.of(0), Degrees.of(-90));
030 *
031 * Pose3d shooterPose = SHOOTER.withRotation(
032 *         new Rotation3d(0, 0, turretSim.getAngle().getRadians()));
033 * }</pre>
034 *
035 * @param x     X translation from the robot origin (meters)
036 * @param y     Y translation from the robot origin (meters)
037 * @param z     Z translation from the robot origin (meters)
038 * @param roll  Rotation about the X axis
039 * @param pitch Rotation about the Y axis
040 * @param yaw   Rotation about the Z axis
041 */
042public record Offsets (
043        Distance x, Distance y, Distance z, Angle roll, Angle pitch, Angle yaw)
044        implements StructSerializable {
045    public static final OffsetsStruct struct = new OffsetsStruct();
046
047    /**
048     *
049     *
050     * <h4>Constructs an Offsets instance with no rotation</h4>
051     *
052     * @param x X translation from the robot origin (meters)
053     * @param y Y translation from the robot origin (meters)
054     * @param z Z translation from the robot origin (meters)
055     */
056    public Offsets(double x, double y, double z) {
057        this(Meters.of(x), Meters.of(y), Meters.of(z), Radians.of(0), Radians.of(0), Radians.of(0));
058    }
059
060    /**
061     *
062     *
063     * <h4>Constructs an Offsets instance with no rotation, using your unit of
064     * choice</h4>
065     *
066     * @param x X translation from the robot origin
067     * @param y Y translation from the robot origin
068     * @param z Z translation from the robot origin
069     */
070    public Offsets(Distance x, Distance y, Distance z) {
071        this(x.in(Meters), y.in(Meters), z.in(Meters));
072    }
073
074    /**
075     *
076     *
077     * <h4>Constructs an Offsets instance with no translation</h4>
078     *
079     * @param roll  Rotation about the X axis
080     * @param pitch Rotation about the Y axis
081     * @param yaw   Rotation about the Z axis
082     */
083    public Offsets(Angle roll, Angle pitch, Angle yaw) {
084        this(Meters.of(0), Meters.of(0), Meters.of(0), roll, pitch, yaw);
085    }
086
087    /**
088     *
089     *
090     * <h4>Constructs an Offsets instance with no rotation</h4>
091     *
092     * @param x     X translation from the robot origin (meters)
093     * @param y     Y translation from the robot origin (meters)
094     * @param z     Z translation from the robot origin (meters)
095     * @param roll  Rotation about the X axis
096     * @param pitch Rotation about the Y axis
097     * @param yaw   Rotation about the Z axis
098     */
099    public Offsets(double x, double y, double z, Angle roll, Angle pitch, Angle yaw) {
100        this(Meters.of(x), Meters.of(y), Meters.of(z), roll, pitch, yaw);
101    }
102
103    /**
104     *
105     *
106     * <h4>Translational component as a {@link Translation3d}</h4>
107     *
108     * @return translation from the robot origin
109     */
110    public Translation3d toTranslation3d() {
111        return new Translation3d(x.in(Meters), y.in(Meters), z.in(Meters));
112    }
113
114    /**
115     *
116     *
117     * <h4>Rotational component as a {@link Rotation3d}</h4>
118     *
119     * @return Rotation3d of the roll, pitch, and yaw components of the offset
120     */
121    public Rotation3d toRotation3d() {
122        return new Rotation3d(roll, pitch, yaw);
123    }
124
125    /**
126     *
127     *
128     * <h4>This offset as a {@link Pose3d}.</h4>
129     *
130     * <p>
131     * Useful for components whose pose is fully static and requires no adjustments
132     *
133     * @return pose at this offset's position and orientation
134     */
135    public Pose3d toPose3d() {
136        return new Pose3d(toTranslation3d(), toRotation3d());
137    }
138
139    /**
140     *
141     *
142     * <h4>Applies this offset's translation and rotation onto an existing
143     * {@link Pose3d}</h4>
144     *
145     * @param pose the base pose to offset
146     * @return a new pose with this offset applied to both translation and rotation
147     */
148    public Pose3d applyToPose3d(Pose3d pose) {
149        return new Pose3d(
150                pose.getMeasureX().plus(x),
151                pose.getMeasureY().plus(y),
152                pose.getMeasureZ().plus(z),
153                applyToRotation3d(pose.getRotation()));
154    }
155
156    /**
157     *
158     *
159     * <h4>Applies this offset's translation and rotation onto an existing
160     * {@link Pose3d}</h4>
161     *
162     * <p>
163     * Translation is applied, but it is <b>robot robot relative</b>
164     *
165     * @param pose the base pose to offset
166     * @return a new pose with this offset applied to both translation and rotation
167     */
168    public Pose3d applyToPose3dRobotRelative(Pose3d pose) {
169        Translation3d rotatedOffset = toTranslation3d().rotateBy(pose.getRotation());
170        return new Pose3d(
171                pose.getTranslation().plus(rotatedOffset), applyToRotation3d(pose.getRotation()));
172    }
173
174    /**
175     *
176     *
177     * <h4>Composes this offset's rotation onto an existing {@link Rotation3d}</h4>
178     *
179     * @param rotation the base rotation to offset
180     * @return a new rotation with this offset's rotation applied on top
181     */
182    public Rotation3d applyToRotation3d(Rotation3d rotation) {
183        return new Rotation3d(
184                rotation.getX() + roll.in(Radians),
185                rotation.getY() + pitch.in(Radians),
186                rotation.getZ() + yaw.in(Radians));
187    }
188
189    /**
190     *
191     *
192     * <h4>Pose at this offset's translation with an additional rotation</h4>
193     *
194     * <p>
195     * Applies a rotation to this offsets rotation
196     *
197     * <pre>{@code
198     * SHOOTER_OFFSETS.withRotation(
199     *         new Rotation3d(0, 0, turretSim.getAngle().getRadians()));
200     * }</pre>
201     *
202     * @param rotation the additional rotation to add to this offset's rotation
203     * @return a new pose at this offset's translation with the combined rotation
204     */
205    public Pose3d withRotation(Rotation3d rotation) {
206        return new Pose3d(toTranslation3d(), toRotation3d().plus(rotation));
207    }
208
209    /**
210     *
211     *
212     * <h4>Applies this offset's X/Y translation and yaw onto an existing
213     * {@link Pose2d}</h4>
214     *
215     * <p>
216     * Roll and pitch are ignored because they have no 2D equivalent
217     *
218     * @param pose the base 2D pose to offset
219     * @return a new pose of this offset's X, Y, and yaw
220     */
221    public Pose2d applyToPose2d(Pose2d pose) {
222        return new Pose2d(
223                pose.getMeasureX().plus(x),
224                pose.getMeasureY().plus(y),
225                applyToRotation2d(pose.getRotation()));
226    }
227
228    /**
229     *
230     *
231     * <h4>Applies this offset's X/Y translation and yaw onto an existing
232     * {@link Pose2d}</h4>
233     *
234     * <p>
235     * Translation is applied, but it is <b>robot robot relative</b>
236     *
237     * @param pose the base 2D pose to offset
238     * @return a new pose with this offset applied
239     */
240    public Pose2d applyToPose2dRobotRelative(Pose2d pose) {
241        Translation2d rotatedOffset = new Translation2d(x, y).rotateBy(pose.getRotation());
242        return new Pose2d(
243                pose.getTranslation().plus(rotatedOffset), applyToRotation2d(pose.getRotation()));
244    }
245
246    /**
247     *
248     *
249     * <h4>Applies this offset's yaw onto an existing {@link Rotation2d}</h4>
250     *
251     * <p>
252     * Roll and pitch are ignored because they have no 2D equivalent
253     *
254     * @param rotation the base 2D rotation to offset
255     * @return a new rotation with this offset's yaw applied on top
256     */
257    public Rotation2d applyToRotation2d(Rotation2d rotation) {
258        return new Rotation2d(rotation.getRadians() + yaw.in(Radians));
259    }
260}
261
262final class OffsetsStruct implements Struct<Offsets> {
263    @Override
264    public Class<Offsets> getTypeClass() {
265        return Offsets.class;
266    }
267
268    @Override
269    public String getTypeName() {
270        return "Offsets";
271    }
272
273    @Override
274    public int getSize() {
275        return kSizeDouble * 6;
276    }
277
278    @Override
279    public String getSchema() {
280        return "double x;double y;double z;double roll;double pitch;double yaw";
281    }
282
283    @Override
284    public Offsets unpack(ByteBuffer bb) {
285        return new Offsets(Meters.of(bb.getDouble()), Meters.of(bb.getDouble()), Meters.of(bb.getDouble()), 
286            Radians.of(bb.getDouble()), Radians.of(bb.getDouble()), Radians.of(bb.getDouble()));
287    }
288
289    @Override
290    public void pack(ByteBuffer bb, Offsets value) {
291        bb.putDouble(value.x().in(Meters));
292        bb.putDouble(value.y().in(Meters));
293        bb.putDouble(value.z().in(Meters));
294        bb.putDouble(value.roll().in(Radians));
295        bb.putDouble(value.pitch().in(Radians));
296        bb.putDouble(value.yaw().in(Radians));
297    }
298}