Wait for at least one index pulse while zeroing.

Change-Id: Ic5648d533bb9a9bab05c7683c7874b2c1de8fde4
diff --git a/frc971/zeroing/zeroing.cc b/frc971/zeroing/zeroing.cc
index c2150d5..d279318 100644
--- a/frc971/zeroing/zeroing.cc
+++ b/frc971/zeroing/zeroing.cc
@@ -10,6 +10,7 @@
     const constants::Values::ZeroingConstants& constants) {
   index_diff_ = constants.index_difference;
   max_sample_count_ = constants.average_filter_size;
+  index_pulse_count_after_reset_ = 0;
 
   start_pos_samples_.reserve(max_sample_count_);
 
@@ -21,9 +22,19 @@
   start_pos_ = 0;
   start_pos_samples_.clear();
   zeroed_ = false;
+  wait_for_index_pulse_ = true;
 }
 
 void ZeroingEstimator::UpdateEstimate(const PotAndIndexPosition& info) {
+  // We want to make sure that we encounter at least one index pulse while
+  // zeroing. So we take the index pulse count from the first sample after
+  // reset and wait for that count to change before we consider ourselves
+  // zeroed.
+  if (wait_for_index_pulse_) {
+    index_pulse_count_after_reset_ = info.index_pulses;
+    wait_for_index_pulse_ = false;
+  }
+
   if (start_pos_samples_.size() < max_sample_count_) {
     start_pos_samples_.push_back(info.pot - info.encoder);
   } else {
@@ -45,7 +56,8 @@
   // If there are no index pulses to use or we don't have enough samples yet to
   // have a well-filtered starting position then we use the filtered value as
   // our best guess.
-  if (info.index_pulses == 0 || offset_ratio_ready() < 1.0) {
+  if (info.index_pulses == index_pulse_count_after_reset_ ||
+      offset_ratio_ready() < 1.0) {
     start_pos_ = start_average;
   } else {
     // We calculate an aproximation of the value of the last index position.
diff --git a/frc971/zeroing/zeroing.h b/frc971/zeroing/zeroing.h
index d85c8b8..13be062 100644
--- a/frc971/zeroing/zeroing.h
+++ b/frc971/zeroing/zeroing.h
@@ -1,17 +1,15 @@
 #ifndef FRC971_ZEROING_ZEROING_H_
 #define FRC971_ZEROING_ZEROING_H_
 
+#include <cstdint>
 #include <vector>
+
 #include "frc971/control_loops/control_loops.q.h"
 #include "frc971/constants.h"
 
 // TODO(pschrader): Support the ZeroingConstants::measured_index_position
 // parameter.
 //
-// TODO(pschrader): Wait for an index pulse during zeroing. If we start up with
-// a non-zero number of index pulses then the logic will use that pulse to
-// compute the starting position (a/k/a the offset).
-//
 // TODO(pschrader): Create an error API to flag faults/errors etc..
 //
 // TODO(pschrader): Flag an error if encoder index pulse is not n revolutions
@@ -78,6 +76,13 @@
   // The estimated starting position of the mechanism. We also call this the
   // 'offset' in some contexts.
   double start_pos_;
+  // Flag for triggering logic that takes note of the current index pulse count
+  // after a reset. See `index_pulse_count_after_reset_'.
+  bool wait_for_index_pulse_;
+  // After a reset we keep track of the index pulse count with this. Only after
+  // the index pulse count changes (i.e. increments at least once or wraps
+  // around) will we consider the mechanism zeroed.
+  uint32_t index_pulse_count_after_reset_;
   // Marker to track whether we're fully zeroed yet or not.
   bool zeroed_;
 };
diff --git a/frc971/zeroing/zeroing_test.cc b/frc971/zeroing/zeroing_test.cc
index 51d1274..5f1d263 100644
--- a/frc971/zeroing/zeroing_test.cc
+++ b/frc971/zeroing/zeroing_test.cc
@@ -166,11 +166,48 @@
   ZeroingEstimator estimator(
       Values::ZeroingConstants{kSampleSize, index_diff, 0.0});
 
+  MoveTo(&sim, &estimator, 3.1 * index_diff);
+
   for (unsigned int i = 0; i < kSampleSize; i++) {
     MoveTo(&sim, &estimator, 5.0 * index_diff);
   }
+
   ASSERT_NEAR(3.1 * index_diff, estimator.offset(), 0.001);
 }
 
+TEST_F(ZeroingTest, WaitForIndexPulseAfterReset) {
+  double index_diff = 0.6;
+  PositionSensorSimulator sim(index_diff);
+  sim.Initialize(3.1 * index_diff, index_diff / 3.0);
+  ZeroingEstimator estimator(
+      Values::ZeroingConstants{kSampleSize, index_diff, 0.0});
+
+  // Make sure to fill up the averaging filter with samples.
+  for (unsigned int i = 0; i < kSampleSize; i++) {
+    MoveTo(&sim, &estimator, 3.1 * index_diff);
+  }
+
+  // Make sure we're not zeroed until we hit an index pulse.
+  ASSERT_FALSE(estimator.zeroed());
+
+  // Trigger an index pulse; we should now be zeroed.
+  MoveTo(&sim, &estimator, 4.5 * index_diff);
+  ASSERT_TRUE(estimator.zeroed());
+
+  // Reset the zeroing logic and supply a bunch of samples within the current
+  // index segment.
+  estimator.Reset();
+  for (unsigned int i = 0; i < kSampleSize; i++) {
+    MoveTo(&sim, &estimator, 4.2 * index_diff);
+  }
+
+  // Make sure we're not zeroed until we hit an index pulse.
+  ASSERT_FALSE(estimator.zeroed());
+
+  // Trigger another index pulse; we should be zeroed again.
+  MoveTo(&sim, &estimator, 3.1 * index_diff);
+  ASSERT_TRUE(estimator.zeroed());
+}
+
 }  // namespace zeroing
 }  // namespace frc971