Publish the x, y, theta state of the drivetrain

This gives us a start of an on-field position to follow.

Change-Id: Ibdd3f06e1856efd451ceb504801ed32605b5e560
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index 432dd0e..577250a 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -173,6 +173,7 @@
         "//aos/logging:matrix_logging",
         "//aos/logging:queue_logging",
         "//aos/util:log_interval",
+        "//frc971/control_loops:runge_kutta",
         "//frc971/queues:gyro",
         "//frc971/wpilib:imu_queue",
     ],
diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index 62f31fa..36f7c82 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -10,11 +10,12 @@
 #include "aos/logging/queue_logging.h"
 #include "aos/logging/matrix_logging.h"
 
+#include "frc971/control_loops/drivetrain/down_estimator.h"
 #include "frc971/control_loops/drivetrain/drivetrain.q.h"
+#include "frc971/control_loops/drivetrain/drivetrain_config.h"
 #include "frc971/control_loops/drivetrain/polydrivetrain.h"
 #include "frc971/control_loops/drivetrain/ssdrivetrain.h"
-#include "frc971/control_loops/drivetrain/down_estimator.h"
-#include "frc971/control_loops/drivetrain/drivetrain_config.h"
+#include "frc971/control_loops/runge_kutta.h"
 #include "frc971/queues/gyro.q.h"
 #include "frc971/shifter_hall_effect.h"
 #include "frc971/wpilib/imu.q.h"
@@ -78,6 +79,51 @@
   }
 }
 
+::Eigen::Matrix<double, 3, 1> DrivetrainLoop::PredictState(
+    const ::Eigen::Matrix<double, 3, 1> &xytheta_state,
+    const ::Eigen::Matrix<double, 7, 1> &state,
+    const ::Eigen::Matrix<double, 7, 1> &previous_state) const {
+  const double dt =
+      ::std::chrono::duration_cast<::std::chrono::duration<double>>(
+          dt_config_.dt)
+          .count();
+
+  const double distance_traveled =
+      (state(0) + state(2)) / 2.0 -
+      (previous_state(0) + previous_state(2)) / 2.0;
+
+  const double omega0 =
+      (previous_state(3) - previous_state(1)) / (dt_config_.robot_radius * 2.0);
+  const double omega1 = (state(3) - state(1)) / (dt_config_.robot_radius * 2.0);
+  const double alpha = (omega1 - omega0) / dt;
+
+  const double velocity_start = (previous_state(3) + previous_state(1)) / 2.0;
+  const double velocity_end = (state(3) + state(1)) / 2.0;
+
+  const double acceleration = (velocity_end - velocity_start) / dt;
+  const double velocity_offset =
+      distance_traveled / dt - 0.5 * acceleration * dt - velocity_start;
+  const double velocity0 = velocity_start + velocity_offset;
+
+  // TODO(austin): Substep 10x here.  This is super important! ?
+  return RungeKutta(
+      [&dt, &velocity0, &acceleration, &omega0, &alpha](
+          double t, const ::Eigen::Matrix<double, 3, 1> &X) {
+        const double velocity1 = velocity0 + acceleration * t;
+        const double omega1 = omega0 + alpha * t;
+        const double theta = X(2);
+
+        return (::Eigen::Matrix<double, 3, 1>()
+                    << ::std::cos(theta) * velocity1,
+                ::std::sin(theta) * velocity1, omega1)
+            .finished();
+      },
+      xytheta_state, 0.0,
+      ::std::chrono::duration_cast<::std::chrono::duration<double>>(
+          dt_config_.dt)
+          .count());
+}
+
 void DrivetrainLoop::RunIteration(
     const ::frc971::control_loops::DrivetrainQueue::Goal *goal,
     const ::frc971::control_loops::DrivetrainQueue::Position *position,
@@ -163,7 +209,7 @@
     LOG(DEBUG,
         "New IMU value from ADIS16448, rate is %f, angle %f, fused %f, bias "
         "%f\n",
-        rate, angle, down_estimator_.X_hat(0, 0), down_estimator_.X_hat(1, 0));
+        rate, angle, down_estimator_.X_hat(0), down_estimator_.X_hat(1));
     down_U_(0, 0) = rate;
   }
   down_estimator_.UpdateObserver(down_U_, ::aos::controls::kLoopFrequency);
@@ -227,9 +273,44 @@
     Y << position->left_encoder, position->right_encoder, last_gyro_rate_,
         last_accel_;
     kf_.Correct(Y);
+
+    // We are going to choose to integrate velocity to get position by assuming
+    // that velocity is a linear function of time.  For drivetrains with large
+    // amounts of mass, we won't get large changes in acceleration over a 5 ms
+    // timestep.  Do note, the only place that this matters is when we are
+    // talking about the curvature errors introduced by integration.  The
+    // velocities are scaled such that the distance traveled is correct.
+    //
+    // We want to do this after the kalman filter runs so we take into account
+    // the encoder and gyro corrections.
+    //
+    // Start by computing the beginning and ending linear and angular
+    // velocities.
+    // To handle 0 velocity well, compute the offset required to be added to
+    // both velocities to make the robot travel the correct distance.
+
+    xytheta_state_.block<3, 1>(0, 0) = PredictState(
+        xytheta_state_.block<3, 1>(0, 0), kf_.X_hat(), last_state_);
+
+    // Use trapezoidal integration for the gyro heading since it's more
+    // accurate.
+    const double average_angular_velocity =
+        ((kf_.X_hat(3) - kf_.X_hat(1)) + (last_state_(3) - last_state_(1))) /
+        2.0 / (dt_config_.robot_radius * 2.0);
+
     integrated_kf_heading_ +=
-        chrono::duration_cast<chrono::duration<double>>(dt_config_.dt).count() *
-        (kf_.X_hat(3, 0) - kf_.X_hat(1, 0)) / (dt_config_.robot_radius * 2.0);
+        ::std::chrono::duration_cast<::std::chrono::duration<double>>(
+            dt_config_.dt)
+            .count() *
+        average_angular_velocity;
+
+    // Copy over the gyro heading.
+    xytheta_state_(2) = integrated_kf_heading_;
+    // Copy over the velocities heading.
+    xytheta_state_(3) = kf_.X_hat(1);
+    xytheta_state_(4) = kf_.X_hat(3);
+    // Copy over the voltage errors.
+    xytheta_state_.block<2, 1>(5, 0) = kf_.X_hat().block<2, 1>(4, 0);
 
     // gyro_heading = (real_right - real_left) / width
     // wheel_heading = (wheel_right - wheel_left) / width
@@ -267,7 +348,7 @@
 
   // set the output status of the control loop state
   if (status) {
-    status->robot_speed = (kf_.X_hat(1, 0) + kf_.X_hat(3, 0)) / 2.0;
+    status->robot_speed = (kf_.X_hat(1) + kf_.X_hat(3)) / 2.0;
 
     Eigen::Matrix<double, 2, 1> linear =
         dt_config_.LeftRightToLinear(kf_.X_hat());
@@ -288,11 +369,16 @@
     status->uncapped_left_voltage = kf_.U_uncapped(0, 0);
     status->uncapped_right_voltage = kf_.U_uncapped(1, 0);
 
-    status->left_voltage_error = kf_.X_hat(4, 0);
-    status->right_voltage_error = kf_.X_hat(5, 0);
-    status->estimated_angular_velocity_error = kf_.X_hat(6, 0);
+    status->left_voltage_error = kf_.X_hat(4);
+    status->right_voltage_error = kf_.X_hat(5);
+    status->estimated_angular_velocity_error = kf_.X_hat(6);
     status->estimated_heading = integrated_kf_heading_;
-    status->ground_angle = down_estimator_.X_hat(0, 0) + dt_config_.down_offset;
+
+    status->x = xytheta_state_(0);
+    status->y = xytheta_state_(1);
+    status->theta = xytheta_state_(2);
+
+    status->ground_angle = down_estimator_.X_hat(0) + dt_config_.down_offset;
 
     dt_openloop_.PopulateStatus(status);
     dt_closedloop_.PopulateStatus(status);
@@ -326,7 +412,8 @@
   last_left_voltage_ = left_voltage;
   last_right_voltage_ = right_voltage;
 
-  kf_.UpdateObserver(U, ::aos::controls::kLoopFrequency);
+  last_state_ = kf_.X_hat();
+  kf_.UpdateObserver(U, dt_config_.dt);
 }
 
 void DrivetrainLoop::Zero(
diff --git a/frc971/control_loops/drivetrain/drivetrain.h b/frc971/control_loops/drivetrain/drivetrain.h
index 3f6c32f..9396aed 100644
--- a/frc971/control_loops/drivetrain/drivetrain.h
+++ b/frc971/control_loops/drivetrain/drivetrain.h
@@ -38,6 +38,12 @@
 
   void Zero(::frc971::control_loops::DrivetrainQueue::Output *output) override;
 
+  // Computes the xy state change given the change in the lr state.
+  ::Eigen::Matrix<double, 3, 1> PredictState(
+      const ::Eigen::Matrix<double, 3, 1> &xytheta_state,
+      const ::Eigen::Matrix<double, 7, 1> &state,
+      const ::Eigen::Matrix<double, 7, 1> &previous_state) const;
+
   double last_gyro_rate_ = 0.0;
 
   const DrivetrainConfig<double> dt_config_;
@@ -66,6 +72,16 @@
   bool has_been_enabled_ = false;
 
   double last_accel_ = 0.0;
+
+  // Current xytheta state of the robot.  This is essentially the kalman filter
+  // integrated up in a direction.
+  // [x, y, theta, vl, vr, left_error, right_error]
+  ::Eigen::Matrix<double, 7, 1> xytheta_state_ =
+      ::Eigen::Matrix<double, 7, 1>::Zero();
+
+  // Last kalman filter state.
+  ::Eigen::Matrix<double, 7, 1> last_state_ =
+      ::Eigen::Matrix<double, 7, 1>::Zero();
 };
 
 }  // namespace drivetrain
diff --git a/frc971/control_loops/drivetrain/drivetrain.q b/frc971/control_loops/drivetrain/drivetrain.q
index 7fb859d..2e78e51 100644
--- a/frc971/control_loops/drivetrain/drivetrain.q
+++ b/frc971/control_loops/drivetrain/drivetrain.q
@@ -145,8 +145,11 @@
     double estimated_angular_velocity_error;
     // The KF estimated heading.
     double estimated_heading;
-    // The KF wheel estimated heading.
-    //double estimated_wheel_heading;
+
+    // xytheta of the robot.
+    double x;
+    double y;
+    double theta;
 
     // True if the output voltage was capped last cycle.
     bool output_was_capped;