Spring works.

Change-Id: I370012cc80e9019467100bec631c488c10a73141
diff --git a/frc971/zeroing/BUILD b/frc971/zeroing/BUILD
index 5271cc2..96315c7 100644
--- a/frc971/zeroing/BUILD
+++ b/frc971/zeroing/BUILD
@@ -1,73 +1,75 @@
-package(default_visibility = ['//visibility:public'])
+package(default_visibility = ["//visibility:public"])
 
-load('//aos/build:queues.bzl', 'queue_library')
+load("//aos/build:queues.bzl", "queue_library")
+load("//tools:environments.bzl", "mcu_cpus")
 
 cc_library(
-  name = 'averager',
-  hdrs = [
-    'averager.h',
-  ],
+    name = "averager",
+    hdrs = [
+        "averager.h",
+    ],
 )
 
 cc_test(
-  name = 'averager_test',
-  srcs = [
-    'averager_test.cc',
-  ],
-  deps = [
-    ':averager',
-    '//aos/testing:googletest',
-  ],
+    name = "averager_test",
+    srcs = [
+        "averager_test.cc",
+    ],
+    deps = [
+        ":averager",
+        "//aos/testing:googletest",
+    ],
 )
 
 cc_library(
-  name = 'zeroing',
-  srcs = [
-    'zeroing.cc',
-  ],
-  hdrs = [
-    'zeroing.h',
-  ],
-  deps = [
-    ':wrap',
-    '//frc971/control_loops:queues',
-    '//frc971:constants',
-  ],
+    name = "zeroing",
+    srcs = [
+        "zeroing.cc",
+    ],
+    hdrs = [
+        "zeroing.h",
+    ],
+    deps = [
+        ":wrap",
+        "//frc971:constants",
+        "//frc971/control_loops:queues",
+    ],
 )
 
 cc_test(
-  name = 'zeroing_test',
-  srcs = [
-    'zeroing_test.cc',
-  ],
-  deps = [
-    '//aos/testing:googletest',
-    '//aos/testing:test_shm',
-    ':zeroing',
-    '//aos/common/util:thread',
-    '//aos/common:die',
-    '//frc971/control_loops:position_sensor_sim',
-    '//frc971/control_loops:queues',
-  ],
+    name = "zeroing_test",
+    srcs = [
+        "zeroing_test.cc",
+    ],
+    deps = [
+        ":zeroing",
+        "//aos/common:die",
+        "//aos/common/util:thread",
+        "//aos/testing:googletest",
+        "//aos/testing:test_shm",
+        "//frc971/control_loops:position_sensor_sim",
+        "//frc971/control_loops:queues",
+    ],
 )
 
 cc_library(
-  name = 'wrap',
-  srcs = [
-    'wrap.cc',
-  ],
-  hdrs = [
-    'wrap.h',
-  ],
+    name = "wrap",
+    srcs = [
+        "wrap.cc",
+    ],
+    hdrs = [
+        "wrap.h",
+    ],
+    compatible_with = mcu_cpus,
 )
 
 cc_test(
-  name = 'wrap_test',
-  srcs = [
-    'wrap_test.cc',
-  ],
-  deps = [
-    '//aos/testing:googletest',
-    ':wrap',
-  ],
+    name = "wrap_test",
+    srcs = [
+        "wrap_test.cc",
+    ],
+    deps = [
+        ":wrap",
+        "//aos/testing:googletest",
+    ],
 )
diff --git a/frc971/zeroing/wrap.cc b/frc971/zeroing/wrap.cc
index 2698928..593ac97 100644
--- a/frc971/zeroing/wrap.cc
+++ b/frc971/zeroing/wrap.cc
@@ -5,8 +5,12 @@
 namespace frc971 {
 namespace zeroing {
 
+float Wrap(float nearest, float value, float period) {
+  return remainderf(value - nearest, period) + nearest;
+}
+
 double Wrap(double nearest, double value, double period) {
-  return ::std::remainder(value - nearest, period) + nearest;
+  return remainder(value - nearest, period) + nearest;
 }
 
 }  // namespace zeroing
diff --git a/frc971/zeroing/wrap.h b/frc971/zeroing/wrap.h
index 84b049d..c6ec085 100644
--- a/frc971/zeroing/wrap.h
+++ b/frc971/zeroing/wrap.h
@@ -7,6 +7,7 @@
 // Returns a modified value which has been wrapped such that it is +- period/2
 // away from nearest.
 double Wrap(double nearest, double value, double period);
+float Wrap(float nearest, float value, float period);
 
 }  // namespace zeroing
 }  // namespace frc971
diff --git a/frc971/zeroing/wrap_test.cc b/frc971/zeroing/wrap_test.cc
index e195385..b5f9100 100644
--- a/frc971/zeroing/wrap_test.cc
+++ b/frc971/zeroing/wrap_test.cc
@@ -1,6 +1,7 @@
-#include <random>
-
 #include "frc971/zeroing/wrap.h"
+
+#include <cmath>
+
 #include "gtest/gtest.h"
 
 namespace frc971 {
@@ -39,6 +40,18 @@
   }
 }
 
+// Tests some various positive and negative values for wrap (with floats).
+TEST(WrapTest, TestFloatWrap) {
+  EXPECT_NEAR(1.0f, Wrap(0.0f, 1.0f, 10.0f), 1e-6f);
+  EXPECT_NEAR(-1.0f, Wrap(0.0f, -1.0f, 10.0f), 1e-6f);
+
+  EXPECT_NEAR(1.0f, Wrap(5.0f, 1.0f, 10.0f), 1e-6f);
+  EXPECT_NEAR(9.0f, Wrap(5.0f, -1.0f, 10.0f), 1e-6f);
+
+  EXPECT_NEAR(10.0f, Wrap(5.0f, 10.0f, 10.0f), 1e-6f);
+  EXPECT_NEAR(1.0f, Wrap(5.0f, -9.0f, 10.0f), 1e-6f);
+}
+
 }  // namespace testing
 }  // namespace zeroing
 }  // namespace frc971
diff --git a/motors/BUILD b/motors/BUILD
index 292a40f..0a63da0 100644
--- a/motors/BUILD
+++ b/motors/BUILD
@@ -172,9 +172,10 @@
         "//motors/core",
         "//motors/peripheral:adc",
         "//motors/peripheral:can",
+        "//motors/seems_reasonable:drivetrain_lib",
+        "//motors/seems_reasonable:spring",
         "//motors/usb",
         "//motors/usb:cdc",
-        "//motors/seems_reasonable:drivetrain_lib",
     ],
 )
 
diff --git a/motors/seems_reasonable/BUILD b/motors/seems_reasonable/BUILD
index 47ca065..a1613eb 100644
--- a/motors/seems_reasonable/BUILD
+++ b/motors/seems_reasonable/BUILD
@@ -74,3 +74,23 @@
         "//frc971/control_loops/drivetrain:polydrivetrain_uc",
     ],
 )
+
+cc_library(
+    name = "spring",
+    srcs = ["spring.cc"],
+    hdrs = ["spring.h"],
+    compatible_with = mcu_cpus,
+    visibility = ["//visibility:public"],
+    deps = ["//frc971/zeroing:wrap"],
+)
+
+cc_test(
+    name = "spring_test",
+    srcs = [
+        "spring_test.cc",
+    ],
+    deps = [
+        ":spring",
+        "//aos/testing:googletest",
+    ],
+)
diff --git a/motors/seems_reasonable/spring.cc b/motors/seems_reasonable/spring.cc
new file mode 100644
index 0000000..d66e950
--- /dev/null
+++ b/motors/seems_reasonable/spring.cc
@@ -0,0 +1,152 @@
+#include "motors/seems_reasonable/spring.h"
+
+#include "frc971/zeroing/wrap.h"
+
+#include <cmath>
+
+namespace motors {
+namespace seems_reasonable {
+namespace {
+
+constexpr float kTwoPi = 2.0 * M_PI;
+
+}  // namespace
+
+float NextGoal(float current_goal, float goal) {
+  float remainder = remainderf(current_goal - goal, kTwoPi);
+  if (remainder >= 0.0f) {
+    remainder -= kTwoPi;
+  }
+  return -remainder + current_goal;
+}
+
+float PreviousGoal(float current_goal, float goal) {
+  float remainder = remainderf(current_goal - goal, kTwoPi);
+  if (remainder <= 0.0f) {
+    remainder += kTwoPi;
+  }
+  return -remainder + current_goal;
+}
+
+void Spring::Iterate(bool unload, bool prime, bool fire, bool force_reset,
+                     bool encoder_valid, float angle) {
+  // Angle is +- M_PI.  So, we need to find the nearest angle to the previous
+  // one, and that's our new angle.
+  angle_ = ::frc971::zeroing::Wrap(angle_, angle, kTwoPi);
+
+  switch (state_) {
+    case State::UNINITIALIZED:
+      // Go to the previous unload from where we are.
+      goal_ = angle_;
+      goal_ = PreviousGoal(kUnloadGoal);
+      if (prime && fire) {
+        Unload();
+      }
+      break;
+    case State::UNLOAD:
+      if (!encoder_valid) {
+        state_ = State::STUCK_UNLOAD;
+      } else if (!unload && prime && fire) {
+        // Go to the next goal from the current location.  This handles if we
+        // fired or didn't on the previous cycle.
+        goal_ = angle_;
+        goal_ = NextGoal(kLoadGoal);
+        Load();
+      }
+    case State::STUCK_UNLOAD:
+      if (force_reset && encoder_valid && state_ == State::STUCK_UNLOAD) {
+        state_ = State::UNINITIALIZED;
+      } else if (timeout_ > 0) {
+        --timeout_;
+      }
+      break;
+    case State::LOAD:
+      if (!encoder_valid) {
+        goal_ = PreviousGoal(kUnloadGoal);
+        StuckUnload();
+      } else if (unload) {
+        goal_ = PreviousGoal(kUnloadGoal);
+        Unload();
+      } else if (!Near()) {
+        if (timeout_ > 0) {
+          --timeout_;
+        } else {
+          StuckUnload();
+        }
+      } else if (prime) {
+        goal_ = NextGoal(kPrimeGoal);
+        Prime();
+      }
+      break;
+    case State::PRIME:
+      if (!encoder_valid) {
+        goal_ = PreviousGoal(kUnloadGoal);
+        StuckUnload();
+      } else if (unload) {
+        goal_ = PreviousGoal(kUnloadGoal);
+        Unload();
+      } else if (!prime) {
+        goal_ = PreviousGoal(kLoadGoal);
+        Load();
+      } else if (!Near()) {
+        if (timeout_ > 0) {
+          --timeout_;
+        } else {
+          StuckUnload();
+        }
+      } else if (fire) {
+        goal_ = NextGoal(kFireGoal);
+        Fire();
+      }
+      break;
+
+    case State::FIRE:
+      if (!encoder_valid) {
+        goal_ = PreviousGoal(kUnloadGoal);
+        StuckUnload();
+      } else if (!Near()) {
+        if (timeout_ > 0) {
+          --timeout_;
+        } else {
+          StuckUnload();
+        }
+      } else {
+        // TODO(austin): Maybe have a different timeout for success.
+        if (timeout_ > 0) {
+          timeout_--;
+        } else {
+          Load();
+          goal_ = NextGoal(kLoadGoal);
+        }
+      }
+      break;
+  }
+  const float error = goal_ - angle_;
+  const float derror = (error - last_error_) * 200.0f;
+
+  switch (state_) {
+    case State::UNINITIALIZED:
+      output_ = 0.0f;
+      break;
+    case State::STUCK_UNLOAD:
+    case State::UNLOAD:
+      if (timeout_ > 0) {
+        output_ = -0.1f;
+      } else {
+        output_ = 0.0f;
+      }
+      break;
+
+    case State::LOAD:
+    case State::PRIME:
+    case State::FIRE: {
+      constexpr float kP = 3.00f;
+      constexpr float kD = 0.00f;
+      output_ = kP * error + kD * derror;
+    } break;
+  }
+  last_error_ = error;
+}
+
+}  // namespace seems_reasonable
+}  // namespace motors
diff --git a/motors/seems_reasonable/spring.h b/motors/seems_reasonable/spring.h
new file mode 100644
index 0000000..908ecc4
--- /dev/null
+++ b/motors/seems_reasonable/spring.h
@@ -0,0 +1,102 @@
+#ifndef MOTORS_SEEMS_REASONABLE_SPRING_H_
+#define MOTORS_SEEMS_REASONABLE_SPRING_H_
+
+#include <cmath>
+
+namespace motors {
+namespace seems_reasonable {
+
+float NextGoal(float current_goal, float goal);
+float PreviousGoal(float current_goal, float goal);
+
+class Spring {
+ public:
+  Spring() = default;
+  Spring(const Spring &) = delete;
+  Spring &operator=(const Spring &) = delete;
+
+  // Iterates the loop.
+  // If unload is true, unload.
+  // If the encoder isn't valid, unload.
+  // If prime is true, go to primed state.
+  // If prime and fire are true, fire.
+  void Iterate(bool unload, bool prime, bool fire, bool force_reset,
+               bool encoder_valid, float angle);
+
+  enum class State {
+    UNINITIALIZED = 0,
+    STUCK_UNLOAD = 1,
+    UNLOAD = 2,
+    LOAD = 3,
+    PRIME = 4,
+    FIRE = 5,
+  };
+
+  // Returns the current to output to the spring motors.
+  float output() const { return output_; }
+
+  // Returns true if the motor is near the goal.
+  bool Near() { return ::std::abs(angle_ - goal_) < 0.2f; }
+
+  State state() const { return state_; }
+
+  float angle() const { return angle_; }
+  float goal() const { return goal_; }
+
+  int timeout() const { return timeout_; }
+
+ private:
+  void Load() {
+    timeout_ = 5 * 200;
+    state_ = State::LOAD;
+  }
+
+  void Prime() {
+    timeout_ = 1 * 200;
+    state_ = State::PRIME;
+  }
+
+  void Unload() {
+    timeout_ = 10 * 200;
+    state_ = State::UNLOAD;
+  }
+
+  void StuckUnload() {
+    timeout_ = 10 * 200;
+    state_ = State::STUCK_UNLOAD;
+  }
+
+  void Fire() {
+    timeout_ = 100;
+    state_ = State::FIRE;
+  }
+
+  float NextGoal(float goal) {
+    return ::motors::seems_reasonable::NextGoal(goal_, goal);
+  }
+
+  float PreviousGoal(float goal) {
+    return ::motors::seems_reasonable::PreviousGoal(goal_, goal);
+  }
+
+  State state_ = State::UNINITIALIZED;
+
+  // Note, these need to be (-M_PI, M_PI]
+  constexpr static float kLoadGoal = -0.18f;
+  constexpr static float kPrimeGoal = -0.10f;
+  constexpr static float kFireGoal = -0.0f;
+  constexpr static float kUnloadGoal = kFireGoal;
+
+  float angle_ = 0.0f;
+  float goal_ = 0.0f;
+
+  int timeout_ = 0;
+
+  float output_ = 0.0f;
+  float last_error_ = 0.0f;
+};
+
+}  // namespace seems_reasonable
+}  // namespace motors
+
+#endif  //  MOTORS_SEEMS_REASONABLE_SPRING_H_
diff --git a/motors/seems_reasonable/spring_test.cc b/motors/seems_reasonable/spring_test.cc
new file mode 100644
index 0000000..bcc2a8a
--- /dev/null
+++ b/motors/seems_reasonable/spring_test.cc
@@ -0,0 +1,35 @@
+#include "motors/seems_reasonable/spring.h"
+
+#include "gtest/gtest.h"
+
+namespace motors {
+namespace seems_reasonable {
+namespace testing {
+
+// Tests that NextGoal always returns the next goal.
+TEST(GoalTest, TestNextGoal) {
+  EXPECT_NEAR(1.0, NextGoal(0.0, 1.0), 1e-6);
+
+  EXPECT_NEAR(2.0 * M_PI + 1.0, NextGoal(1.1, 1.0), 1e-6);
+
+  EXPECT_NEAR(6.0 * M_PI + 1.0, NextGoal(1.1 + 4.0 * M_PI, 1.0), 1e-6);
+
+  EXPECT_NEAR(2.0 * M_PI + 1.0, NextGoal(1.0, 1.0), 1e-6);
+
+  EXPECT_NEAR(2.0 * M_PI + 6.0, NextGoal(6.1, 6.0), 1e-6);
+}
+
+// Tests that PreviousGoal always returns the previous goal.
+TEST(GoalTest, TestPreviousGoal) {
+  EXPECT_NEAR(-2.0 * M_PI + 1.0, PreviousGoal(0.0, 1.0), 1e-6);
+
+  EXPECT_NEAR(1.0, PreviousGoal(1.1, 1.0), 1e-6);
+
+  EXPECT_NEAR(4.0 * M_PI + 1.0, PreviousGoal(1.1 + 4.0 * M_PI, 1.0), 1e-6);
+
+  EXPECT_NEAR(-2.0 * M_PI + 1.0, PreviousGoal(1.0, 1.0), 1e-6);
+}
+
+}  // namespace testing
+}  // namespace seems_reasonable
+}  // namespace motors
diff --git a/motors/simple_receiver.cc b/motors/simple_receiver.cc
index c6d70b8..678319a 100644
--- a/motors/simple_receiver.cc
+++ b/motors/simple_receiver.cc
@@ -5,17 +5,18 @@
 #include <atomic>
 #include <cmath>
 
+#include "frc971/control_loops/drivetrain/polydrivetrain.h"
 #include "motors/core/kinetis.h"
 #include "motors/core/time.h"
 #include "motors/peripheral/adc.h"
 #include "motors/peripheral/can.h"
 #include "motors/peripheral/configuration.h"
+#include "motors/seems_reasonable/drivetrain_dog_motor_plant.h"
+#include "motors/seems_reasonable/polydrivetrain_dog_motor_plant.h"
+#include "motors/seems_reasonable/spring.h"
 #include "motors/usb/cdc.h"
 #include "motors/usb/usb.h"
 #include "motors/util.h"
-#include "frc971/control_loops/drivetrain/polydrivetrain.h"
-#include "motors/seems_reasonable/drivetrain_dog_motor_plant.h"
-#include "motors/seems_reasonable/polydrivetrain_dog_motor_plant.h"
 
 namespace frc971 {
 namespace motors {
@@ -26,6 +27,7 @@
 using ::frc971::constants::ShifterHallEffect;
 using ::frc971::control_loops::DrivetrainQueue_Goal;
 using ::frc971::control_loops::DrivetrainQueue_Output;
+using ::motors::seems_reasonable::Spring;
 
 const ShifterHallEffect kThreeStateDriveShifter{0.0, 0.0, 0.25, 0.75};
 
@@ -58,6 +60,7 @@
 ::std::atomic<teensy::AcmTty *> global_stdout{nullptr};
 
 ::std::atomic<PolyDrivetrain<float> *> global_polydrivetrain{nullptr};
+::std::atomic<Spring *> global_spring{nullptr};
 
 // Last width we received on each channel.
 uint16_t pwm_input_widths[5];
@@ -373,7 +376,7 @@
 
     const float magnitude = hypot(sin, cos);
     const float magnitude_error = ::std::abs(magnitude - 1.0f);
-    valid = magnitude_error < 0.15f;
+    valid = magnitude_error < 0.30f;
 
     angle = ::std::atan2(sin, cos);
   }
@@ -388,6 +391,7 @@
   PIT_TFLG3 = 1;
   PolyDrivetrain<float> *polydrivetrain =
       global_polydrivetrain.load(::std::memory_order_acquire);
+  Spring *spring = global_spring.load(::std::memory_order_acquire);
 
   salsa::SimpleAdcReadings adc_readings;
   {
@@ -396,13 +400,21 @@
   }
 
   EncoderReading encoder(adc_readings);
+  static float last_good_encoder = 0.0f;
+  static int invalid_encoder_count = 0;
+  if (encoder.valid) {
+    last_good_encoder = encoder.angle;
+    invalid_encoder_count = 0;
+  } else {
+    ++invalid_encoder_count;
+  }
 
   const bool lost_any_channel = lost_channel(0) || lost_channel(1) ||
                                 lost_channel(2) || lost_channel(3) ||
                                 lost_channel(4) ||
-                                (::std::abs(convert_input_width(4)) < 0.05f);
+                                (::std::abs(convert_input_width(4)) < 0.5f);
 
-  if (polydrivetrain != nullptr) {
+  if (polydrivetrain != nullptr && spring != nullptr) {
     DrivetrainQueue_Goal goal;
     goal.control_loop_driving = false;
     if (lost_any_channel) {
@@ -410,9 +422,8 @@
       goal.wheel = 0.0f;
     } else {
       goal.throttle = convert_input_width(1);
-      goal.wheel = convert_input_width(0);
+      goal.wheel = -convert_input_width(0);
     }
-    goal.throttle = convert_input_width(1);
     goal.quickturn = ::std::abs(polydrivetrain->velocity()) < 0.25f;
 
     DrivetrainQueue_Output output;
@@ -427,6 +438,33 @@
     vesc_set_duty(2, output.right_voltage / 12.0f);
     vesc_set_duty(3, output.right_voltage / 12.0f);
 
+    bool prime = convert_input_width(2) > 0.1f;
+    bool fire = convert_input_width(3) > 0.1f;
+
+    bool unload = lost_any_channel;
+    static bool was_lost = true;
+    bool force_reset = !lost_any_channel && was_lost;
+    was_lost = lost_any_channel;
+
+    spring->Iterate(unload, prime, fire, force_reset, invalid_encoder_count <= 2,
+                    last_good_encoder);
+
+    float spring_output = spring->output();
+
+    vesc_set_duty(4, -spring_output);
+    vesc_set_duty(5, spring_output);
+
+    /*
+     // Massive debug.  Turn on for useful bits.
+    if (!encoder.valid) {
+      printf("Stuck encoder: ADC %" PRIu16 " %" PRIu16
+             " enc %d/1000 %s mag %d\n",
+             adc_readings.sin, adc_readings.cos, (int)(encoder.angle * 1000),
+             encoder.valid ? "T" : "f",
+             (int)(hypot(ConvertEncoderChannel(adc_readings.sin),
+                         ConvertEncoderChannel(adc_readings.cos)) *
+                   1000));
+    }
     static int i = 0;
     ++i;
     if (i > 20) {
@@ -441,7 +479,8 @@
       } else {
         printf(
             "TODO(Austin): 200Hz loop %d %d %d %d %d, lr, %d, %d velocity %d "
-            "ADC %" PRIu16 " %" PRIu16 " enc %d/1000 %s from %d\n",
+            " state: %d, near %d angle %d goal %d to: %d ADC %" PRIu16
+            " %" PRIu16 " enc %d/1000 %s from %d\n",
             (int)(convert_input_width(0) * 1000),
             (int)(convert_input_width(1) * 1000),
             (int)(convert_input_width(2) * 1000),
@@ -450,7 +489,11 @@
             static_cast<int>(output.left_voltage * 100),
             static_cast<int>(output.right_voltage * 100),
             static_cast<int>(polydrivetrain->velocity() * 100),
-            adc_readings.sin, adc_readings.cos, (int)(encoder.angle * 1000),
+            static_cast<int>(spring->state()), static_cast<int>(spring->Near()),
+            static_cast<int>(spring->angle() * 1000),
+            static_cast<int>(spring->goal() * 1000),
+            static_cast<int>(spring->timeout()), adc_readings.sin,
+            adc_readings.cos, (int)(encoder.angle * 1000),
             encoder.valid ? "T" : "f",
             (int)(::std::sqrt(ConvertEncoderChannel(adc_readings.sin) *
                                   ConvertEncoderChannel(adc_readings.sin) +
@@ -459,6 +502,7 @@
                   1000));
       }
     }
+    */
   }
 }
 
@@ -550,6 +594,8 @@
 
   PolyDrivetrain<float> polydrivetrain(GetDrivetrainConfig(), nullptr);
   global_polydrivetrain.store(&polydrivetrain, ::std::memory_order_release);
+  Spring spring;
+  global_spring.store(&spring, ::std::memory_order_release);
 
   // Leave the LEDs on for a bit longer.
   delay(300);