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    public RotTrlTransform3d() {
066        this(new Rotation3d(), new Translation3d());
067    }
068
069    /**
070     * The rotation-translation transformation that makes poses in the world consider this pose as the
071     * new origin, or change the basis to this pose.
072     *
073     * @param pose The new origin
074     */
075    public static RotTrlTransform3d makeRelativeTo(Pose3d pose) {
076        return new RotTrlTransform3d(pose.getRotation(), pose.getTranslation()).inverse();
077    }
078
079    /** The inverse of this transformation. Applying the inverse will "undo" this transformation. */
080    public RotTrlTransform3d inverse() {
081        var inverseRot = rot.unaryMinus();
082        var inverseTrl = trl.rotateBy(inverseRot).unaryMinus();
083        return new RotTrlTransform3d(inverseRot, inverseTrl);
084    }
085
086    /** This transformation as a Transform3d (as if of the origin) */
087    public Transform3d getTransform() {
088        return new Transform3d(trl, rot);
089    }
090
091    /** The translation component of this transformation */
092    public Translation3d getTranslation() {
093        return trl;
094    }
095
096    /** The rotation component of this transformation */
097    public Rotation3d getRotation() {
098        return rot;
099    }
100
101    public Translation3d apply(Translation3d trl) {
102        return trl.rotateBy(rot).plus(this.trl);
103    }
104
105    public List<Translation3d> applyTrls(List<Translation3d> trls) {
106        return trls.stream().map(this::apply).collect(Collectors.toList());
107    }
108
109    public Rotation3d apply(Rotation3d rot) {
110        return rot.plus(this.rot);
111    }
112
113    public List<Rotation3d> applyRots(List<Rotation3d> rots) {
114        return rots.stream().map(this::apply).collect(Collectors.toList());
115    }
116
117    public Pose3d apply(Pose3d pose) {
118        return new Pose3d(apply(pose.getTranslation()), apply(pose.getRotation()));
119    }
120
121    public List<Pose3d> applyPoses(List<Pose3d> poses) {
122        return poses.stream().map(this::apply).collect(Collectors.toList());
123    }
124}