From cc02a948a0f1a23a758e7437c07bb7682a724cce Mon Sep 17 00:00:00 2001
From: Jade <spacey-sooty@proton.me>
Date: Tue, 9 Jul 2024 09:26:08 +0800
Subject: [PATCH] [wpimath] Add Reset methods to PoseEstimator (#6751)

---
 .../first/math/estimator/PoseEstimator.java   | 31 +++++++++++++++++++
 .../include/frc/estimator/PoseEstimator.h     | 22 +++++++++++++
 .../include/frc/estimator/PoseEstimator.inc   | 23 ++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/wpimath/src/main/java/edu/wpi/first/math/estimator/PoseEstimator.java b/wpimath/src/main/java/edu/wpi/first/math/estimator/PoseEstimator.java
index 5f17c83fd7f..2674e2ceaed 100644
--- a/wpimath/src/main/java/edu/wpi/first/math/estimator/PoseEstimator.java
+++ b/wpimath/src/main/java/edu/wpi/first/math/estimator/PoseEstimator.java
@@ -11,6 +11,7 @@
 import edu.wpi.first.math.VecBuilder;
 import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.math.geometry.Translation2d;
 import edu.wpi.first.math.geometry.Twist2d;
 import edu.wpi.first.math.interpolation.TimeInterpolatableBuffer;
 import edu.wpi.first.math.kinematics.Kinematics;
@@ -125,6 +126,36 @@ public void resetPosition(Rotation2d gyroAngle, T wheelPositions, Pose2d poseMet
     m_poseEstimate = m_odometry.getPoseMeters();
   }
 
+  /**
+   * Resets the robot's pose.
+   *
+   * @param pose The pose to reset to.
+   */
+  public void resetPose(Pose2d pose) {
+    m_odometry.resetPose(pose);
+    m_odometryPoseBuffer.clear();
+  }
+
+  /**
+   * Resets the robot's translation.
+   *
+   * @param translation The pose to translation to.
+   */
+  public void resetTranslation(Translation2d translation) {
+    m_odometry.resetTranslation(translation);
+    m_odometryPoseBuffer.clear();
+  }
+
+  /**
+   * Resets the robot's rotation.
+   *
+   * @param rotation The rotation to reset to.
+   */
+  public void resetRotation(Rotation2d rotation) {
+    m_odometry.resetRotation(rotation);
+    m_odometryPoseBuffer.clear();
+  }
+
   /**
    * Gets the estimated robot pose.
    *
diff --git a/wpimath/src/main/native/include/frc/estimator/PoseEstimator.h b/wpimath/src/main/native/include/frc/estimator/PoseEstimator.h
index cd0600f698b..da9f4f3ee68 100644
--- a/wpimath/src/main/native/include/frc/estimator/PoseEstimator.h
+++ b/wpimath/src/main/native/include/frc/estimator/PoseEstimator.h
@@ -15,6 +15,7 @@
 
 #include "frc/geometry/Pose2d.h"
 #include "frc/geometry/Rotation2d.h"
+#include "frc/geometry/Translation2d.h"
 #include "frc/interpolation/TimeInterpolatableBuffer.h"
 #include "frc/kinematics/Kinematics.h"
 #include "frc/kinematics/Odometry.h"
@@ -87,6 +88,27 @@ class WPILIB_DLLEXPORT PoseEstimator {
   void ResetPosition(const Rotation2d& gyroAngle,
                      const WheelPositions& wheelPositions, const Pose2d& pose);
 
+  /**
+   * Resets the robot's pose.
+   *
+   * @param pose The pose to reset to.
+   */
+  void ResetPose(const Pose2d& pose);
+
+  /**
+   * Resets the robot's translation.
+   *
+   * @param translation The pose to translation to.
+   */
+  void ResetTranslation(const Translation2d& translation);
+
+  /**
+   * Resets the robot's rotation.
+   *
+   * @param rotation The rotation to reset to.
+   */
+  void ResetRotation(const Rotation2d& rotation);
+
   /**
    * Gets the estimated robot pose.
    *
diff --git a/wpimath/src/main/native/include/frc/estimator/PoseEstimator.inc b/wpimath/src/main/native/include/frc/estimator/PoseEstimator.inc
index f8bddca52d0..cb39dd26053 100644
--- a/wpimath/src/main/native/include/frc/estimator/PoseEstimator.inc
+++ b/wpimath/src/main/native/include/frc/estimator/PoseEstimator.inc
@@ -5,6 +5,9 @@
 #pragma once
 
 #include "frc/estimator/PoseEstimator.h"
+#include "frc/geometry/Pose2d.h"
+#include "frc/geometry/Translation2d.h"
+#include "frc/kinematics/WheelPositions.h"
 
 namespace frc {
 
@@ -53,6 +56,26 @@ void PoseEstimator<WheelSpeeds, WheelPositions>::ResetPosition(
   m_poseEstimate = m_odometry.GetPose();
 }
 
+template <typename WheelSpeeds, typename WheelPositions>
+void PoseEstimator<WheelSpeeds, WheelPositions>::ResetPose(const Pose2d& pose) {
+  m_odometry.ResetPose(pose);
+  m_odometryPoseBuffer.Clear();
+}
+
+template <typename WheelSpeeds, typename WheelPositions>
+void PoseEstimator<WheelSpeeds, WheelPositions>::ResetTranslation(
+    const Translation2d& translation) {
+  m_odometry.ResetTranslation(translation);
+  m_odometryPoseBuffer.Clear();
+}
+
+template <typename WheelSpeeds, typename WheelPositions>
+void PoseEstimator<WheelSpeeds, WheelPositions>::ResetRotation(
+    const Rotation2d& rotation) {
+  m_odometry.ResetTranslation(rotation);
+  m_odometryPoseBuffer.Clear();
+}
+
 template <typename WheelSpeeds, typename WheelPositions>
 Pose2d PoseEstimator<WheelSpeeds, WheelPositions>::GetEstimatedPosition()
     const {