From dc276b651bb5079415d10e6def9d7eb6fb4b17ba Mon Sep 17 00:00:00 2001
From: Jade <spacey-sooty@proton.me>
Date: Tue, 9 Jul 2024 09:25:42 +0800
Subject: [PATCH] [wpimath] Add reset methods to Odometry (#6702)

---
 .../wpi/first/math/kinematics/Odometry.java   | 32 ++++++++++++++++++
 .../native/include/frc/kinematics/Odometry.h  | 33 +++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/wpimath/src/main/java/edu/wpi/first/math/kinematics/Odometry.java b/wpimath/src/main/java/edu/wpi/first/math/kinematics/Odometry.java
index 0b4dab7a950..66ce592957d 100644
--- a/wpimath/src/main/java/edu/wpi/first/math/kinematics/Odometry.java
+++ b/wpimath/src/main/java/edu/wpi/first/math/kinematics/Odometry.java
@@ -6,6 +6,7 @@
 
 import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.math.geometry.Translation2d;
 
 /**
  * Class for odometry. Robot code should not use this directly- Instead, use the particular type for
@@ -63,6 +64,37 @@ public void resetPosition(Rotation2d gyroAngle, T wheelPositions, Pose2d poseMet
     m_previousWheelPositions = m_kinematics.copy(wheelPositions);
   }
 
+  /**
+   * Resets the pose.
+   *
+   * @param poseMeters The pose to reset to.
+   */
+  public void resetPose(Pose2d poseMeters) {
+    m_gyroOffset = m_gyroOffset.plus(poseMeters.getRotation().minus(m_poseMeters.getRotation()));
+    m_poseMeters = poseMeters;
+    m_previousAngle = m_poseMeters.getRotation();
+  }
+
+  /**
+   * Resets the translation of the pose.
+   *
+   * @param translation The translation to reset to.
+   */
+  public void resetTranslation(Translation2d translation) {
+    m_poseMeters = new Pose2d(translation, m_poseMeters.getRotation());
+  }
+
+  /**
+   * Resets the rotation of the pose.
+   *
+   * @param rotation The rotation to reset to.
+   */
+  public void resetRotation(Rotation2d rotation) {
+    m_gyroOffset = m_gyroOffset.plus(rotation.minus(m_poseMeters.getRotation()));
+    m_poseMeters = new Pose2d(m_poseMeters.getTranslation(), rotation);
+    m_previousAngle = m_poseMeters.getRotation();
+  }
+
   /**
    * Returns the position of the robot on the field.
    *
diff --git a/wpimath/src/main/native/include/frc/kinematics/Odometry.h b/wpimath/src/main/native/include/frc/kinematics/Odometry.h
index 53066c7041a..c0daea0e3e0 100644
--- a/wpimath/src/main/native/include/frc/kinematics/Odometry.h
+++ b/wpimath/src/main/native/include/frc/kinematics/Odometry.h
@@ -7,6 +7,8 @@
 #include <wpi/SymbolExports.h>
 
 #include "frc/geometry/Pose2d.h"
+#include "frc/geometry/Rotation2d.h"
+#include "frc/geometry/Translation2d.h"
 #include "frc/kinematics/Kinematics.h"
 #include "frc/kinematics/WheelPositions.h"
 
@@ -58,6 +60,37 @@ class WPILIB_DLLEXPORT Odometry {
     m_previousWheelPositions = wheelPositions;
   }
 
+  /**
+   * Resets the pose.
+   *
+   * @param pose The pose to reset to.
+   */
+  void ResetPose(const Pose2d& pose) {
+    m_gyroOffset = m_gyroOffset + (pose.Rotation() - m_pose.Rotation());
+    m_pose = pose;
+    m_previousAngle = pose.Rotation();
+  }
+
+  /**
+   * Resets the translation of the pose.
+   *
+   * @param translation The translation to reset to.
+   */
+  void ResetTranslation(const Translation2d& translation) {
+    m_pose = Pose2d{translation, m_pose.Rotation()};
+  }
+
+  /**
+   * Resets the rotation of the pose.
+   *
+   * @param rotation The rotation to reset to.
+   */
+  void ResetRotation(const Rotation2d& rotation) {
+    m_gyroOffset = m_gyroOffset + (rotation - m_pose.Rotation());
+    m_pose = Pose2d{m_pose.Translation(), rotation};
+    m_previousAngle = rotation;
+  }
+
   /**
    * Returns the position of the robot on the field.
    * @return The pose of the robot.