Refactored HandleDrivetrain from joystick_reader to aos_input.

Change-Id: I4ed796cf90e698ed2198f0da4a2b75e88e28c434
diff --git a/aos/common/commonmath.h b/aos/common/commonmath.h
index f9337b2..acbcb1c 100644
--- a/aos/common/commonmath.h
+++ b/aos/common/commonmath.h
@@ -1,6 +1,8 @@
 #ifndef AOS_COMMON_MATH_H_
 #define AOS_COMMON_MATH_H_
 
+#include <cmath>
+
 namespace aos {
 
 // Clips a value so that it is in [min, max]
@@ -22,6 +24,20 @@
   }
 }
 
+// Adds deadband to provided value.  deadband is the region close to the origin
+// to add the deadband to, and max is the maximum input value used to re-scale
+// the output after adding the deadband.
+static inline double Deadband(double value, const double deadband,
+                              const double max) {
+  if (::std::abs(value) < deadband) {
+    value = 0.0;
+  } else if (value > 0.0) {
+    value = (value - deadband) / (max - deadband);
+  } else {
+    value = (value + deadband) / (max - deadband);
+  }
+  return value;
+}
 }  // namespace aos
 
 #endif  // AOS_COMMON_MATH_H_
diff --git a/aos/input/BUILD b/aos/input/BUILD
index 30cb053..d8fc2c8 100644
--- a/aos/input/BUILD
+++ b/aos/input/BUILD
@@ -16,3 +16,22 @@
     '//aos/common/logging:queue_logging',
   ],
 )
+
+cc_library(
+  name = 'drivetrain_input',
+  srcs = [
+    'drivetrain_input.cc'
+  ],
+  hdrs = [
+    'drivetrain_input.h',
+  ],
+  deps = [
+    '//aos/common/input:driver_station_data',
+    '//aos/common/logging',
+    '//aos/common/logging:queue_logging',
+    '//aos/common/messages:robot_state',
+    '//aos/common/network:socket',
+    '//aos/common:math',
+    '//frc971/control_loops/drivetrain:drivetrain_queue',
+  ],
+)
diff --git a/aos/input/drivetrain_input.cc b/aos/input/drivetrain_input.cc
new file mode 100644
index 0000000..a0d95cf
--- /dev/null
+++ b/aos/input/drivetrain_input.cc
@@ -0,0 +1,188 @@
+#include "aos/input/drivetrain_input.h"
+
+#include <math.h>
+#include <stdio.h>
+#include <string.h>
+#include <cmath>
+
+#include "aos/common/commonmath.h"
+#include "aos/common/input/driver_station_data.h"
+#include "aos/common/logging/logging.h"
+#include "frc971/control_loops/drivetrain/drivetrain.q.h"
+
+using ::frc971::control_loops::drivetrain_queue;
+using ::aos::input::driver_station::ButtonLocation;
+using ::aos::input::driver_station::ControlBit;
+using ::aos::input::driver_station::JoystickAxis;
+using ::aos::input::driver_station::POVLocation;
+
+namespace aos {
+namespace input {
+
+void DrivetrainInputReader::HandleDrivetrain(
+    const ::aos::input::driver_station::Data &data) {
+  bool is_control_loop_driving = false;
+
+  const auto wheel_and_throttle = GetWheelAndThrottle(data);
+  const double wheel = wheel_and_throttle.wheel;
+  const double throttle = wheel_and_throttle.throttle;
+
+  drivetrain_queue.status.FetchLatest();
+  if (drivetrain_queue.status.get()) {
+    robot_velocity_ = drivetrain_queue.status->robot_speed;
+  }
+
+  if (data.PosEdge(kTurn1) || data.PosEdge(kTurn2)) {
+    if (drivetrain_queue.status.get()) {
+      left_goal_ = drivetrain_queue.status->estimated_left_position;
+      right_goal_ = drivetrain_queue.status->estimated_right_position;
+    }
+  }
+  const double current_left_goal =
+      left_goal_ - wheel * wheel_multiplier_ + throttle * 0.3;
+  const double current_right_goal =
+      right_goal_ + wheel * wheel_multiplier_ + throttle * 0.3;
+  if (data.IsPressed(kTurn1) || data.IsPressed(kTurn2)) {
+    is_control_loop_driving = true;
+  }
+  auto new_drivetrain_goal = drivetrain_queue.goal.MakeMessage();
+  new_drivetrain_goal->steering = wheel;
+  new_drivetrain_goal->throttle = throttle;
+  new_drivetrain_goal->quickturn = data.IsPressed(kQuickTurn);
+  new_drivetrain_goal->control_loop_driving = is_control_loop_driving;
+  new_drivetrain_goal->left_goal = current_left_goal;
+  new_drivetrain_goal->right_goal = current_right_goal;
+  new_drivetrain_goal->left_velocity_goal = 0;
+  new_drivetrain_goal->right_velocity_goal = 0;
+
+  new_drivetrain_goal->linear.max_velocity = 3.0;
+  new_drivetrain_goal->linear.max_acceleration = 20.0;
+
+  if (!new_drivetrain_goal.Send()) {
+    LOG(WARNING, "sending stick values failed\n");
+  }
+}
+
+DrivetrainInputReader::WheelAndThrottle
+SteeringWheelDrivetrainInputReader::GetWheelAndThrottle(
+    const ::aos::input::driver_station::Data &data) {
+  const double wheel = -data.GetAxis(kSteeringWheel);
+  const double throttle = -data.GetAxis(kDriveThrottle);
+  return DrivetrainInputReader::WheelAndThrottle{wheel, throttle};
+}
+
+DrivetrainInputReader::WheelAndThrottle
+PistolDrivetrainInputReader::GetWheelAndThrottle(
+    const ::aos::input::driver_station::Data &data) {
+  const double unscaled_wheel = data.GetAxis(kSteeringWheel);
+  double wheel;
+  if (unscaled_wheel < 0.0) {
+    wheel = unscaled_wheel / 0.484375;
+  } else {
+    wheel = unscaled_wheel / 0.385827;
+  }
+
+  const double unscaled_throttle = -data.GetAxis(kDriveThrottle);
+  double unmodified_throttle;
+  if (unscaled_throttle < 0.0) {
+    unmodified_throttle = unscaled_throttle / 0.086614;
+  } else {
+    unmodified_throttle = unscaled_throttle / 0.265625;
+  }
+  unmodified_throttle = aos::Deadband(unmodified_throttle, 0.1, 1.0);
+
+  // Apply a sin function that's scaled to make it feel better.
+  constexpr double throttle_range = M_PI_2 * 0.5;
+
+  double throttle = ::std::sin(throttle_range * unmodified_throttle) /
+                    ::std::sin(throttle_range);
+  throttle = ::std::sin(throttle_range * throttle) / ::std::sin(throttle_range);
+  throttle = 2.0 * unmodified_throttle - throttle;
+  return DrivetrainInputReader::WheelAndThrottle{wheel, throttle};
+}
+
+DrivetrainInputReader::WheelAndThrottle
+XboxDrivetrainInputReader::GetWheelAndThrottle(
+    const ::aos::input::driver_station::Data &data) {
+  // xbox
+  constexpr double kWheelDeadband = 0.05;
+  constexpr double kThrottleDeadband = 0.05;
+  const double wheel =
+      aos::Deadband(-data.GetAxis(kSteeringWheel), kWheelDeadband, 1.0);
+
+  const double unmodified_throttle =
+      aos::Deadband(-data.GetAxis(kDriveThrottle), kThrottleDeadband, 1.0);
+
+  // Apply a sin function that's scaled to make it feel better.
+  constexpr double throttle_range = M_PI_2 * 0.9;
+
+  double throttle = ::std::sin(throttle_range * unmodified_throttle) /
+                    ::std::sin(throttle_range);
+  throttle = ::std::sin(throttle_range * throttle) / ::std::sin(throttle_range);
+  throttle = 2.0 * unmodified_throttle - throttle;
+  return DrivetrainInputReader::WheelAndThrottle{wheel, throttle};
+}
+
+std::unique_ptr<SteeringWheelDrivetrainInputReader>
+SteeringWheelDrivetrainInputReader::Make() {
+  const JoystickAxis kSteeringWheel(1, 1), kDriveThrottle(2, 2);
+  const ButtonLocation kQuickTurn(1, 5);
+  const ButtonLocation kTurn1(1, 7);
+  const ButtonLocation kTurn2(1, 11);
+  std::unique_ptr<SteeringWheelDrivetrainInputReader> result(
+      new SteeringWheelDrivetrainInputReader(kSteeringWheel, kDriveThrottle,
+                                             kQuickTurn, kTurn1, kTurn2));
+  return result;
+}
+
+std::unique_ptr<PistolDrivetrainInputReader>
+PistolDrivetrainInputReader::Make() {
+  // Pistol Grip controller
+  const JoystickAxis kSteeringWheel(1, 2), kDriveThrottle(1, 1);
+  const ButtonLocation kQuickTurn(1, 7);
+  const ButtonLocation kTurn1(1, 8);
+
+  // Nop
+  const ButtonLocation kTurn2(1, 9);
+  std::unique_ptr<PistolDrivetrainInputReader> result(
+      new PistolDrivetrainInputReader(kSteeringWheel, kDriveThrottle,
+                                      kQuickTurn, kTurn1, kTurn2));
+  result->set_wheel_multiplier(0.2);
+  return result;
+}
+
+std::unique_ptr<XboxDrivetrainInputReader> XboxDrivetrainInputReader::Make() {
+  // xbox
+  const JoystickAxis kSteeringWheel(1, 5), kDriveThrottle(1, 2);
+  const ButtonLocation kQuickTurn(1, 5);
+
+  // Nop
+  const ButtonLocation kTurn1(1, 1);
+  const ButtonLocation kTurn2(1, 2);
+
+  std::unique_ptr<XboxDrivetrainInputReader> result(
+      new XboxDrivetrainInputReader(kSteeringWheel, kDriveThrottle, kQuickTurn,
+                                    kTurn1, kTurn2));
+  return result;
+}
+::std::unique_ptr<DrivetrainInputReader> DrivetrainInputReader::Make(
+    InputType type) {
+  std::unique_ptr<DrivetrainInputReader> drivetrain_input_reader;
+
+  using InputType = DrivetrainInputReader::InputType;
+  switch (type) {
+    case InputType::kSteeringWheel:
+      drivetrain_input_reader = SteeringWheelDrivetrainInputReader::Make();
+      break;
+    case InputType::kPistol:
+      drivetrain_input_reader = PistolDrivetrainInputReader::Make();
+      break;
+    case InputType::kXbox:
+      drivetrain_input_reader = XboxDrivetrainInputReader::Make();
+      break;
+  }
+  return drivetrain_input_reader;
+}
+
+}  // namespace input
+}  // namespace aos
diff --git a/aos/input/drivetrain_input.h b/aos/input/drivetrain_input.h
new file mode 100644
index 0000000..4b13ac6
--- /dev/null
+++ b/aos/input/drivetrain_input.h
@@ -0,0 +1,144 @@
+#ifndef AOS_INPUT_DRIVETRAIN_INPUT_H_
+#define AOS_INPUT_DRIVETRAIN_INPUT_H_
+
+#include <math.h>
+#include <stdio.h>
+#include <string.h>
+#include <cmath>
+#include <memory>
+
+#include "aos/common/input/driver_station_data.h"
+#include "aos/common/logging/logging.h"
+#include "frc971/control_loops/drivetrain/drivetrain.q.h"
+
+namespace aos {
+namespace input {
+
+// We have a couple different joystick configurations used to drive our skid
+// steer robots.  These configurations don't change very often, and there are a
+// small, discrete, set of them.  The interface to the drivetrain is the same
+// across all of our drivetrains, so we can abstract that away from our users.
+//
+// We want to be able to re-use that code across years, and switch joystick
+// types quickly and efficiently.
+//
+// DrivetrainInputReader is the abstract base class which provides a consistent
+// API to control drivetrains.
+//
+// To construct the appropriate DrivetrainInputReader, call
+// DrivetrainInputReader::Make with the input type enum.
+// To use it, call HandleDrivetrain(data) every time a joystick packet is
+// received.
+//
+// Base class for the interface to the drivetrain implemented by multiple
+// joystick types.
+class DrivetrainInputReader {
+ public:
+  // Inputs driver station button and joystick locations
+  DrivetrainInputReader(driver_station::JoystickAxis kSteeringWheel,
+                        driver_station::JoystickAxis kDriveThrottle,
+                        driver_station::ButtonLocation kQuickTurn,
+                        driver_station::ButtonLocation kTurn1,
+                        driver_station::ButtonLocation kTurn2)
+      : kSteeringWheel(kSteeringWheel),
+        kDriveThrottle(kDriveThrottle),
+        kQuickTurn(kQuickTurn),
+        kTurn1(kTurn1),
+        kTurn2(kTurn2) {}
+
+  virtual ~DrivetrainInputReader() = default;
+
+  // List of driver joystick types.
+  enum class InputType {
+    kSteeringWheel,
+    kPistol,
+    kXbox,
+  };
+
+  // Constructs the appropriate DrivetrainInputReader.
+  static std::unique_ptr<DrivetrainInputReader> Make(InputType type);
+
+  // Processes new joystick data and publishes drivetrain goal messages.
+  void HandleDrivetrain(const ::aos::input::driver_station::Data &data);
+
+  // Sets the scalar for the steering wheel for closed loop mode converting
+  // steering ratio to meters displacement on the two wheels.
+  void set_wheel_multiplier(double wheel_multiplier) {
+    wheel_multiplier_ = wheel_multiplier;
+  }
+
+  // Returns the current robot velocity in m/s.
+  double robot_velocity() const { return robot_velocity_; }
+
+ protected:
+  const driver_station::JoystickAxis kSteeringWheel;
+  const driver_station::JoystickAxis kDriveThrottle;
+  const driver_station::ButtonLocation kQuickTurn;
+  const driver_station::ButtonLocation kTurn1;
+  const driver_station::ButtonLocation kTurn2;
+
+  // Structure containing the (potentially adjusted) steering and throttle
+  // values from the joysticks.
+  struct WheelAndThrottle {
+    double wheel;
+    double throttle;
+  };
+
+ private:
+  // Computes the steering and throttle from the provided driverstation data.
+  virtual WheelAndThrottle GetWheelAndThrottle(
+      const ::aos::input::driver_station::Data &data) = 0;
+
+  double robot_velocity_ = 0.0;
+  // Goals to send to the drivetrain in closed loop mode.
+  double left_goal_ = 0.0;
+  double right_goal_ = 0.0;
+  // The scale for the joysticks for closed loop mode converting
+  // joysticks to meters displacement on the two wheels.
+  double wheel_multiplier_ = 0.5;
+};
+
+// Implements DrivetrainInputReader for the original steering wheel.
+class SteeringWheelDrivetrainInputReader : public DrivetrainInputReader {
+ public:
+  using DrivetrainInputReader::DrivetrainInputReader;
+
+  // Creates a DrivetrainInputReader with the corresponding joystick ports and
+  // axis for the big steering wheel and throttle stick.
+  static std::unique_ptr<SteeringWheelDrivetrainInputReader> Make();
+
+ private:
+  WheelAndThrottle GetWheelAndThrottle(
+      const ::aos::input::driver_station::Data &data) override;
+};
+
+class PistolDrivetrainInputReader : public DrivetrainInputReader {
+ public:
+  using DrivetrainInputReader::DrivetrainInputReader;
+
+  // Creates a DrivetrainInputReader with the corresponding joystick ports and
+  // axis for the (cheap) pistol grip controller.
+  static std::unique_ptr<PistolDrivetrainInputReader> Make();
+
+ private:
+  WheelAndThrottle GetWheelAndThrottle(
+      const ::aos::input::driver_station::Data &data) override;
+};
+
+class XboxDrivetrainInputReader : public DrivetrainInputReader {
+ public:
+  using DrivetrainInputReader::DrivetrainInputReader;
+
+  // Creates a DrivetrainInputReader with the corresponding joystick ports and
+  // axis for the Xbox controller.
+  static std::unique_ptr<XboxDrivetrainInputReader> Make();
+
+ private:
+  WheelAndThrottle GetWheelAndThrottle(
+      const ::aos::input::driver_station::Data &data) override;
+};
+
+}  // namespace input
+}  // namespace aos
+
+#endif  // AOS_INPUT_DRIVETRAIN_INPUT_H_