From f298a37c9c79723f5cc610a616b0b73b5513f4e1 Mon Sep 17 00:00:00 2001
From: Laurens Valk <laurens@pybricks.com>
Date: Thu, 19 Sep 2024 10:15:27 +0200
Subject: [PATCH] pbio/imu: Prevent calibration while motors in use.

---
 CHANGELOG.md                    |  3 +++
 lib/pbio/include/pbio/dcmotor.h |  5 +++++
 lib/pbio/src/dcmotor.c          | 14 ++++++++++++++
 lib/pbio/src/imu.c              | 10 ++++++++--
 4 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e586c02c0..ac18ec8f6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,6 +35,8 @@
 - On SPIKE Prime Hub and Robot Inventor Hub, moved Bluetooth indications to
   the Bluetooth light. Only warning lights will be shown on the main button
   light. See ([support#1716]) and ([pybricks-micropython#261]).
+- Allow gyro calibration only while all motors are coasting ([support#1840]) to
+  prevent recalibration during very steady moves.
 
 ### Fixed
 - Fixed not able to connect to new Technic Move hub with `LWP3Device()`.
@@ -56,6 +58,7 @@
 [support#1622]: https://github.com/pybricks/support/issues/1622
 [support#1678]: https://github.com/pybricks/support/issues/1678
 [support#1716]: https://github.com/pybricks/support/issues/1716
+[support#1840]: https://github.com/pybricks/support/issues/1840
 
 ## [3.5.0] - 2024-04-11
 
diff --git a/lib/pbio/include/pbio/dcmotor.h b/lib/pbio/include/pbio/dcmotor.h
index 5bc1585eb..182d53ea9 100644
--- a/lib/pbio/include/pbio/dcmotor.h
+++ b/lib/pbio/include/pbio/dcmotor.h
@@ -71,6 +71,7 @@ typedef struct _pbio_dcmotor_t {
 
 /** @cond INTERNAL */
 void pbio_dcmotor_stop_all(bool clear_parents);
+bool pbio_dcmotor_all_coasting(void);
 pbio_error_t pbio_dcmotor_coast(pbio_dcmotor_t *dcmotor);
 pbio_error_t pbio_dcmotor_set_voltage(pbio_dcmotor_t *dcmotor, int32_t voltage);
 int32_t pbio_dcmotor_get_max_voltage(pbdrv_legodev_type_id_t id);
@@ -100,6 +101,10 @@ pbio_error_t pbio_dcmotor_user_command(pbio_dcmotor_t *dcmotor, bool coast, int3
 static inline void pbio_dcmotor_stop_all(bool clear_parents) {
 }
 
+static inline bool pbio_dcmotor_any_active(void) {
+    return false;
+}
+
 static inline pbio_error_t pbio_dcmotor_get_dcmotor(pbdrv_legodev_dev_t *legodev, pbio_dcmotor_t **dcmotor) {
     *dcmotor = NULL;
     return PBIO_ERROR_NOT_SUPPORTED;
diff --git a/lib/pbio/src/dcmotor.c b/lib/pbio/src/dcmotor.c
index c12a92364..aab1a8668 100644
--- a/lib/pbio/src/dcmotor.c
+++ b/lib/pbio/src/dcmotor.c
@@ -59,6 +59,20 @@ void pbio_dcmotor_stop_all(bool clear_parents) {
     }
 }
 
+/**
+ * Tests if all dc motors are coasting.
+ *
+ * @return                  @c true if all motors are coasting, @c false otherwise.
+ */
+bool pbio_dcmotor_all_coasting(void) {
+    for (uint8_t i = 0; i < PBIO_CONFIG_DCMOTOR_NUM_DEV; i++) {
+        if (dcmotors[i].actuation_now != PBIO_DCMOTOR_ACTUATION_COAST) {
+            return false;
+        }
+    }
+    return true;
+}
+
 /**
  * Stops and closes DC motor instance so it can be used in another application.
  *
diff --git a/lib/pbio/src/imu.c b/lib/pbio/src/imu.c
index 065ff814d..62d7c1f09 100644
--- a/lib/pbio/src/imu.c
+++ b/lib/pbio/src/imu.c
@@ -10,6 +10,7 @@
 
 #include <pbio/angle.h>
 #include <pbio/config.h>
+#include <pbio/dcmotor.h>
 #include <pbio/error.h>
 #include <pbio/geometry.h>
 #include <pbio/imu.h>
@@ -60,6 +61,11 @@ bool pbio_imu_is_ready(void) {
 // Called by driver to process unfiltered gyro and accelerometer data recorded while stationary.
 static void pbio_imu_handle_stationary_data_func(const int32_t *gyro_data_sum, const int32_t *accel_data_sum, uint32_t num_samples) {
 
+    // Don't update if not stationary
+    if (!pbio_imu_is_stationary()) {
+        return;
+    }
+
     // If the IMU calibration hasn't been updated in a long time, reset the
     // stationary counter so that the calibration values get a large weight.
     if (!pbio_imu_is_ready()) {
@@ -126,12 +132,12 @@ pbio_error_t pbio_imu_set_base_orientation(pbio_geometry_xyz_t *front_side_axis,
 }
 
 /**
- * Checks if the IMU is currently stationary.
+ * Checks if the IMU is currently stationary and motors are not moving.
  *
  * @return    True if it has been stationary for about a second, false if moving.
  */
 bool pbio_imu_is_stationary(void) {
-    return pbdrv_imu_is_stationary(imu_dev);
+    return pbdrv_imu_is_stationary(imu_dev) && pbio_dcmotor_all_coasting();
 }
 
 // Measured rotation of the Z axis in the user frame for exactly one rotation