001/*
002 * Copyright (C) Photon Vision.
003 *
004 * This program is free software: you can redistribute it and/or modify
005 * it under the terms of the GNU General Public License as published by
006 * the Free Software Foundation, either version 3 of the License, or
007 * (at your option) any later version.
008 *
009 * This program is distributed in the hope that it will be useful,
010 * but WITHOUT ANY WARRANTY; without even the implied warranty of
011 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
012 * GNU General Public License for more details.
013 *
014 * You should have received a copy of the GNU General Public License
015 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
016 */
017
018package org.photonvision.estimation;
019
020import edu.wpi.first.math.geometry.Pose3d;
021import edu.wpi.first.math.geometry.Rotation3d;
022import edu.wpi.first.math.geometry.Transform3d;
023import edu.wpi.first.math.geometry.Translation3d;
024import java.util.List;
025import java.util.stream.Collectors;
026
027/**
028 * Represents a transformation that first rotates a pose around the origin, and then translates it.
029 */
030public class RotTrlTransform3d {
031    private final Translation3d trl;
032    private final Rotation3d rot;
033
034    /**
035     * A rotation-translation transformation.
036     *
037     * <p>Applying this RotTrlTransform3d to poses will preserve their current origin-to-pose
038     * transform as if the origin was transformed by these components instead.
039     *
040     * @param rot The rotation component
041     * @param trl The translation component
042     */
043    public RotTrlTransform3d(Rotation3d rot, Translation3d trl) {
044        this.rot = rot;
045        this.trl = trl;
046    }
047
048    public RotTrlTransform3d(Pose3d initial, Pose3d last) {
049        this.rot = last.getRotation().minus(initial.getRotation());
050        this.trl = last.getTranslation().minus(initial.getTranslation().rotateBy(rot));
051    }
052
053    /**
054     * Creates a rotation-translation transformation from a Transform3d.
055     *
056     * <p>Applying this RotTrlTransform3d to poses will preserve their current origin-to-pose
057     * transform as if the origin was transformed by trf instead.
058     *
059     * @param trf The origin transformation
060     */
061    public RotTrlTransform3d(Transform3d trf) {
062        this(trf.getRotation(), trf.getTranslation());
063    }
064
065    /** Constructs the identity transform -- maps an initial pose to itself. */
066    public RotTrlTransform3d() {
067        this(new Rotation3d(), new Translation3d());
068    }
069
070    /**
071     * The rotation-translation transformation that makes poses in the world consider this pose as the
072     * new origin, or change the basis to this pose.
073     *
074     * @param pose The new origin
075     */
076    public static RotTrlTransform3d makeRelativeTo(Pose3d pose) {
077        return new RotTrlTransform3d(pose.getRotation(), pose.getTranslation()).inverse();
078    }
079
080    /**
081     * The inverse of this transformation. Applying the inverse will "undo" this transformation.
082     *
083     * @return The inverse transformation
084     */
085    public RotTrlTransform3d inverse() {
086        var inverseRot = rot.unaryMinus();
087        var inverseTrl = trl.rotateBy(inverseRot).unaryMinus();
088        return new RotTrlTransform3d(inverseRot, inverseTrl);
089    }
090
091    /**
092     * This transformation as a Transform3d (as if of the origin)
093     *
094     * @return The Transform2d
095     */
096    public Transform3d getTransform() {
097        return new Transform3d(trl, rot);
098    }
099
100    /**
101     * The translation component of this transformation
102     *
103     * @return The Translation3d
104     */
105    public Translation3d getTranslation() {
106        return trl;
107    }
108
109    /**
110     * The rotation component of this transformation
111     *
112     * @return The Rotation3d
113     */
114    public Rotation3d getRotation() {
115        return rot;
116    }
117
118    public Translation3d apply(Translation3d trl) {
119        return trl.rotateBy(rot).plus(this.trl);
120    }
121
122    public List<Translation3d> applyTrls(List<Translation3d> trls) {
123        return trls.stream().map(this::apply).collect(Collectors.toList());
124    }
125
126    public Rotation3d apply(Rotation3d rot) {
127        return rot.plus(this.rot);
128    }
129
130    public List<Rotation3d> applyRots(List<Rotation3d> rots) {
131        return rots.stream().map(this::apply).collect(Collectors.toList());
132    }
133
134    public Pose3d apply(Pose3d pose) {
135        return new Pose3d(apply(pose.getTranslation()), apply(pose.getRotation()));
136    }
137
138    public List<Pose3d> applyPoses(List<Pose3d> poses) {
139        return poses.stream().map(this::apply).collect(Collectors.toList());
140    }
141}