Use enums for drivetrain KF states

This makes things a lot easier to follow.

Change-Id: Ib3d2b60fd8f77851cf842340447d92905b69cb75
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index d42b97c..cf6b19f 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -285,6 +285,7 @@
         ":drivetrain_config",
         ":drivetrain_goal_fbs",
         ":drivetrain_output_fbs",
+        ":drivetrain_states",
         ":drivetrain_status_fbs",
         ":gear",
         ":localizer",
@@ -315,6 +316,7 @@
     }),
     deps = [
         ":drivetrain_config",
+        ":drivetrain_states",
         ":gear",
         "//aos:math",
         "//aos/controls:polytope",
@@ -368,6 +370,11 @@
 )
 
 cc_library(
+    name = "drivetrain_states",
+    hdrs = ["drivetrain_states.h"],
+)
+
+cc_library(
     name = "drivetrain_lib",
     srcs = [
         "drivetrain.cc",
@@ -381,6 +388,7 @@
         ":drivetrain_goal_fbs",
         ":drivetrain_output_fbs",
         ":drivetrain_position_fbs",
+        ":drivetrain_states",
         ":drivetrain_status_fbs",
         ":gear",
         ":improved_down_estimator",
diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index 5f2d3c4..edde3c7 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -301,7 +301,7 @@
     // currently unusable--either don't use voltage error at all for the spline
     // following code, or use the EKF's voltage error estimates.
     const Eigen::Matrix<double, 2, 1> voltage_error =
-        0 * kf_.X_hat().block<2, 1>(4, 0);
+        0 * kf_.X_hat().block<2, 1>(kLeftError, 0);
     dt_spline_.Update(
         output != nullptr && controller_type == ControllerType::SPLINE_FOLLOWER,
         trajectory_state, voltage_error);
@@ -369,24 +369,25 @@
 
     dt_closedloop_.PopulateStatus(&builder);
 
-    builder.add_estimated_left_position(gyro_left_right(0, 0));
-    builder.add_estimated_right_position(gyro_left_right(2, 0));
+    builder.add_estimated_left_position(gyro_left_right(kLeftPosition));
+    builder.add_estimated_right_position(gyro_left_right(kRightPosition));
 
-    builder.add_estimated_left_velocity(gyro_left_right(1, 0));
-    builder.add_estimated_right_velocity(gyro_left_right(3, 0));
+    builder.add_estimated_left_velocity(gyro_left_right(kLeftVelocity));
+    builder.add_estimated_right_velocity(gyro_left_right(kRightVelocity));
 
     if (dt_spline_.enable()) {
       dt_spline_.PopulateStatus(&builder);
     } else {
-      builder.add_robot_speed((kf_.X_hat(1) + kf_.X_hat(3)) / 2.0);
+      builder.add_robot_speed(
+          (kf_.X_hat(kLeftVelocity) + kf_.X_hat(kRightVelocity)) / 2.0);
       builder.add_output_was_capped(dt_closedloop_.output_was_capped());
-      builder.add_uncapped_left_voltage(kf_.U_uncapped(0, 0));
-      builder.add_uncapped_right_voltage(kf_.U_uncapped(1, 0));
+      builder.add_uncapped_left_voltage(kf_.U_uncapped(kLeftVoltage));
+      builder.add_uncapped_right_voltage(kf_.U_uncapped(kRightVoltage));
     }
 
-    builder.add_left_voltage_error(kf_.X_hat(4));
-    builder.add_right_voltage_error(kf_.X_hat(5));
-    builder.add_estimated_angular_velocity_error(kf_.X_hat(6));
+    builder.add_left_voltage_error(kf_.X_hat(kLeftError));
+    builder.add_right_voltage_error(kf_.X_hat(kRightError));
+    builder.add_estimated_angular_velocity_error(kf_.X_hat(kAngularError));
     builder.add_estimated_heading(localizer_->theta());
 
     builder.add_x(localizer_->x());
@@ -429,8 +430,8 @@
   last_last_left_voltage_ = last_left_voltage_;
   last_last_right_voltage_ = last_right_voltage_;
   Eigen::Matrix<double, 2, 1> U;
-  U(0, 0) = last_left_voltage_;
-  U(1, 0) = last_right_voltage_;
+  U(kLeftVoltage) = last_left_voltage_;
+  U(kRightVoltage) = last_right_voltage_;
   last_left_voltage_ = left_voltage;
   last_right_voltage_ = right_voltage;
 
diff --git a/frc971/control_loops/drivetrain/drivetrain.h b/frc971/control_loops/drivetrain/drivetrain.h
index d350321..f282cc0 100644
--- a/frc971/control_loops/drivetrain/drivetrain.h
+++ b/frc971/control_loops/drivetrain/drivetrain.h
@@ -11,6 +11,7 @@
 #include "frc971/control_loops/drivetrain/drivetrain_goal_generated.h"
 #include "frc971/control_loops/drivetrain/drivetrain_output_generated.h"
 #include "frc971/control_loops/drivetrain/drivetrain_position_generated.h"
+#include "frc971/control_loops/drivetrain/drivetrain_states.h"
 #include "frc971/control_loops/drivetrain/drivetrain_status_generated.h"
 #include "frc971/control_loops/drivetrain/gear.h"
 #include "frc971/control_loops/drivetrain/improved_down_estimator.h"
diff --git a/frc971/control_loops/drivetrain/drivetrain_states.h b/frc971/control_loops/drivetrain/drivetrain_states.h
new file mode 100644
index 0000000..9ea2518
--- /dev/null
+++ b/frc971/control_loops/drivetrain/drivetrain_states.h
@@ -0,0 +1,24 @@
+#ifndef FRC971_CONTROL_LOOPS_DRIVETRAIN_DRIVETRAIN_STATES_H_
+#define FRC971_CONTROL_LOOPS_DRIVETRAIN_DRIVETRAIN_STATES_H_
+
+namespace frc971 {
+namespace control_loops {
+namespace drivetrain {
+
+enum KalmanState {
+  kLeftPosition = 0,
+  kLeftVelocity = 1,
+  kRightPosition = 2,
+  kRightVelocity = 3,
+  kLeftError = 4,
+  kRightError = 5,
+  kAngularError = 6
+};
+
+enum OutputState { kLeftVoltage = 0, kRightVoltage = 1 };
+
+}  // namespace drivetrain
+}  // namespace control_loops
+}  // namespace frc971
+
+#endif  // FRC971_CONTROL_LOOPS_DRIVETRAIN_DRIVETRAIN_STATES_H_
diff --git a/frc971/control_loops/drivetrain/polydrivetrain.h b/frc971/control_loops/drivetrain/polydrivetrain.h
index d6cd2b5..17a32b8 100644
--- a/frc971/control_loops/drivetrain/polydrivetrain.h
+++ b/frc971/control_loops/drivetrain/polydrivetrain.h
@@ -20,6 +20,7 @@
 #include "frc971/control_loops/drivetrain/drivetrain_status_float_generated.h"
 #endif  // __linux__
 #include "frc971/control_loops/drivetrain/drivetrain_config.h"
+#include "frc971/control_loops/drivetrain/drivetrain_states.h"
 #include "frc971/control_loops/state_feedback_loop.h"
 
 namespace frc971 {
@@ -294,8 +295,8 @@
 template <typename Scalar>
 void PolyDrivetrain<Scalar>::Update(Scalar voltage_battery) {
   if (dt_config_.loop_type == LoopType::CLOSED_LOOP) {
-    loop_->mutable_X_hat()(0, 0) = kf_->X_hat()(1, 0);
-    loop_->mutable_X_hat()(1, 0) = kf_->X_hat()(3, 0);
+    loop_->mutable_X_hat()(0, 0) = kf_->X_hat()(kLeftVelocity);
+    loop_->mutable_X_hat()(1, 0) = kf_->X_hat()(kRightVelocity);
   }
 
   // TODO(austin): Observer for the current velocity instead of difference
diff --git a/frc971/control_loops/drivetrain/ssdrivetrain.cc b/frc971/control_loops/drivetrain/ssdrivetrain.cc
index 82c83d2..3eacef1 100644
--- a/frc971/control_loops/drivetrain/ssdrivetrain.cc
+++ b/frc971/control_loops/drivetrain/ssdrivetrain.cc
@@ -62,21 +62,25 @@
   const Eigen::Matrix<double, 7, 1> error = kf_->R() - kf_->X_hat();
 
   Eigen::Matrix<double, 2, 2> position_K;
-  position_K << kf_->controller().K(0, 0), kf_->controller().K(0, 2),
-      kf_->controller().K(1, 0), kf_->controller().K(1, 2);
+  position_K << kf_->controller().K(kLeftVoltage, kLeftPosition),
+      kf_->controller().K(kLeftVoltage, kRightPosition),
+      kf_->controller().K(kRightVoltage, kLeftPosition),
+      kf_->controller().K(kRightVoltage, kRightPosition);
   Eigen::Matrix<double, 2, 2> velocity_K;
-  velocity_K << kf_->controller().K(0, 1), kf_->controller().K(0, 3),
-      kf_->controller().K(1, 1), kf_->controller().K(1, 3);
+  velocity_K << kf_->controller().K(kLeftVoltage, kLeftVelocity),
+      kf_->controller().K(kLeftVoltage, kRightVelocity),
+      kf_->controller().K(kRightVoltage, kLeftVelocity),
+      kf_->controller().K(kRightVoltage, kRightVelocity);
 
   Eigen::Matrix<double, 2, 1> position_error;
-  position_error << error(0, 0), error(2, 0);
+  position_error << error(kLeftPosition), error(kRightPosition);
   // drive_error = [total_distance_error, left_error - right_error]
   const auto drive_error = T_inverse_ * position_error;
   Eigen::Matrix<double, 2, 1> velocity_error;
-  velocity_error << error(1, 0), error(3, 0);
+  velocity_error << error(kLeftVelocity), error(kRightVelocity);
 
   Eigen::Matrix<double, 2, 1> U_integral;
-  U_integral << kf_->X_hat(4, 0), kf_->X_hat(5, 0);
+  U_integral << kf_->X_hat(kLeftError), kf_->X_hat(kRightError);
 
   const ::aos::controls::HVPolytope<2, 4, 4> pos_poly_hv(
       U_poly_.static_H() * position_K * T_,
@@ -201,11 +205,11 @@
         (gyro_to_wheel_offset - last_gyro_to_wheel_offset_) *
         dt_config_.robot_radius;
 
-    kf_->mutable_next_R().block<4, 1>(0, 0) =
+    kf_->mutable_next_R().block<4, 1>(kLeftPosition, 0) =
         dt_config_.AngularLinearToLeftRight(next_linear, next_angular);
 
-    kf_->mutable_next_R().block<3, 1>(4, 0) =
-        unprofiled_goal_.block<3, 1>(4, 0);
+    kf_->mutable_next_R().block<3, 1>(kLeftError, 0) =
+        unprofiled_goal_.block<3, 1>(kLeftError, 0);
 
     kf_->mutable_next_R(0, 0) -= wheel_compensation_offset;
     kf_->mutable_next_R(2, 0) += wheel_compensation_offset;
@@ -263,8 +267,8 @@
 void DrivetrainMotorsSS::SetOutput(
     ::frc971::control_loops::drivetrain::OutputT *output) const {
   if (output) {
-    output->left_voltage = kf_->U(0, 0);
-    output->right_voltage = kf_->U(1, 0);
+    output->left_voltage = kf_->U(kLeftVoltage);
+    output->right_voltage = kf_->U(kRightVoltage);
     output->left_high = true;
     output->right_high = true;
   }
@@ -282,10 +286,14 @@
   Eigen::Matrix<double, 4, 1> profiled_gyro_left_right =
       dt_config_.AngularLinearToLeftRight(profiled_linear, profiled_angular);
 
-  builder->add_profiled_left_position_goal(profiled_gyro_left_right(0, 0));
-  builder->add_profiled_left_velocity_goal(profiled_gyro_left_right(1, 0));
-  builder->add_profiled_right_position_goal(profiled_gyro_left_right(2, 0));
-  builder->add_profiled_right_velocity_goal(profiled_gyro_left_right(3, 0));
+  builder->add_profiled_left_position_goal(
+      profiled_gyro_left_right(kLeftPosition));
+  builder->add_profiled_left_velocity_goal(
+      profiled_gyro_left_right(kLeftVelocity));
+  builder->add_profiled_right_position_goal(
+      profiled_gyro_left_right(kRightPosition));
+  builder->add_profiled_right_velocity_goal(
+      profiled_gyro_left_right(kRightVelocity));
 }
 
 }  // namespace drivetrain
diff --git a/frc971/control_loops/drivetrain/ssdrivetrain.h b/frc971/control_loops/drivetrain/ssdrivetrain.h
index 4d91807..d83473d 100644
--- a/frc971/control_loops/drivetrain/ssdrivetrain.h
+++ b/frc971/control_loops/drivetrain/ssdrivetrain.h
@@ -11,6 +11,7 @@
 #include "frc971/control_loops/drivetrain/drivetrain_config.h"
 #include "frc971/control_loops/drivetrain/drivetrain_goal_generated.h"
 #include "frc971/control_loops/drivetrain/drivetrain_output_generated.h"
+#include "frc971/control_loops/drivetrain/drivetrain_states.h"
 #include "frc971/control_loops/drivetrain/drivetrain_status_generated.h"
 #include "frc971/control_loops/drivetrain/localizer.h"
 #include "frc971/control_loops/state_feedback_loop.h"
diff --git a/frc971/control_loops/state_feedback_loop.h b/frc971/control_loops/state_feedback_loop.h
index 168df43..50edf0d 100644
--- a/frc971/control_loops/state_feedback_loop.h
+++ b/frc971/control_loops/state_feedback_loop.h
@@ -413,38 +413,38 @@
   }
   Scalar X_hat(int i, int j = 0) const { return X_hat()(i, j); }
   const Eigen::Matrix<Scalar, number_of_states, 1> &R() const { return R_; }
-  Scalar R(int i, int j) const { return R()(i, j); }
+  Scalar R(int i, int j = 0) const { return R()(i, j); }
   const Eigen::Matrix<Scalar, number_of_states, 1> &next_R() const {
     return next_R_;
   }
-  Scalar next_R(int i, int j) const { return next_R()(i, j); }
+  Scalar next_R(int i, int j = 0) const { return next_R()(i, j); }
   const Eigen::Matrix<Scalar, number_of_inputs, 1> &U() const { return U_; }
-  Scalar U(int i, int j) const { return U()(i, j); }
+  Scalar U(int i, int j = 0) const { return U()(i, j); }
   const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_uncapped() const {
     return U_uncapped_;
   }
-  Scalar U_uncapped(int i, int j) const { return U_uncapped()(i, j); }
+  Scalar U_uncapped(int i, int j = 0) const { return U_uncapped()(i, j); }
   const Eigen::Matrix<Scalar, number_of_inputs, 1> &ff_U() const {
     return ff_U_;
   }
-  Scalar ff_U(int i, int j) const { return ff_U()(i, j); }
+  Scalar ff_U(int i, int j = 0) const { return ff_U()(i, j); }
 
   Eigen::Matrix<Scalar, number_of_states, 1> &mutable_X_hat() {
     return observer_.mutable_X_hat();
   }
   Scalar &mutable_X_hat(int i, int j = 0) { return mutable_X_hat()(i, j); }
   Eigen::Matrix<Scalar, number_of_states, 1> &mutable_R() { return R_; }
-  Scalar &mutable_R(int i, int j) { return mutable_R()(i, j); }
+  Scalar &mutable_R(int i, int j = 0) { return mutable_R()(i, j); }
   Eigen::Matrix<Scalar, number_of_states, 1> &mutable_next_R() {
     return next_R_;
   }
-  Scalar &mutable_next_R(int i, int j) { return mutable_next_R()(i, j); }
+  Scalar &mutable_next_R(int i, int j = 0) { return mutable_next_R()(i, j); }
   Eigen::Matrix<Scalar, number_of_inputs, 1> &mutable_U() { return U_; }
-  Scalar &mutable_U(int i, int j) { return mutable_U()(i, j); }
+  Scalar &mutable_U(int i, int j = 0) { return mutable_U()(i, j); }
   Eigen::Matrix<Scalar, number_of_inputs, 1> &mutable_U_uncapped() {
     return U_uncapped_;
   }
-  Scalar &mutable_U_uncapped(int i, int j) {
+  Scalar &mutable_U_uncapped(int i, int j = 0) {
     return mutable_U_uncapped()(i, j);
   }