Merge "Speed up downloading code."
diff --git a/frc971/control_loops/drivetrain/hybrid_ekf.h b/frc971/control_loops/drivetrain/hybrid_ekf.h
index ecdabb8..2a972b6 100644
--- a/frc971/control_loops/drivetrain/hybrid_ekf.h
+++ b/frc971/control_loops/drivetrain/hybrid_ekf.h
@@ -137,7 +137,7 @@
   };
   static constexpr int kNInputs = 4;
   // Number of previous samples to save.
-  static constexpr int kSaveSamples = 80;
+  static constexpr int kSaveSamples = 200;
   // Whether we should completely rerun the entire stored history of
   // kSaveSamples on every correction. Enabling this will increase overall CPU
   // usage substantially; however, leaving it disabled makes it so that we are
diff --git a/frc971/input/action_joystick_input.h b/frc971/input/action_joystick_input.h
index cdd0aa6..3284b15 100644
--- a/frc971/input/action_joystick_input.h
+++ b/frc971/input/action_joystick_input.h
@@ -65,6 +65,11 @@
     return drivetrain_input_reader_->robot_velocity();
   }
 
+  // Returns the current drivetrain status.
+  const control_loops::drivetrain::Status *drivetrain_status() const {
+    return drivetrain_input_reader_->drivetrain_status();
+  }
+
   // Returns the drivetrain config.
   const ::frc971::control_loops::drivetrain::DrivetrainConfig<double>
   dt_config() const {
diff --git a/frc971/input/drivetrain_input.h b/frc971/input/drivetrain_input.h
index 3419ea7..a33d449 100644
--- a/frc971/input/drivetrain_input.h
+++ b/frc971/input/drivetrain_input.h
@@ -94,6 +94,11 @@
   // Returns the current robot velocity in m/s.
   double robot_velocity() const { return robot_velocity_; }
 
+  // Returns the current drivetrain status.
+  const control_loops::drivetrain::Status *drivetrain_status() const {
+    return drivetrain_status_fetcher_.get();
+  }
+
   void set_vision_align_fn(
       ::std::function<bool(const ::frc971::input::driver_station::Data &data)>
           vision_align_fn) {
diff --git a/y2020/BUILD b/y2020/BUILD
index ef4e0f6..195e2f4 100644
--- a/y2020/BUILD
+++ b/y2020/BUILD
@@ -140,6 +140,7 @@
         "//frc971/input:action_joystick_input",
         "//frc971/input:drivetrain_input",
         "//frc971/input:joystick_input",
+        "//frc971/zeroing:wrap",
         "//y2020:constants",
         "//y2020/control_loops/drivetrain:drivetrain_base",
         "//y2020/control_loops/superstructure:superstructure_goal_fbs",
diff --git a/y2020/control_loops/drivetrain/drivetrain_main.cc b/y2020/control_loops/drivetrain/drivetrain_main.cc
index dec9118..24c876f 100644
--- a/y2020/control_loops/drivetrain/drivetrain_main.cc
+++ b/y2020/control_loops/drivetrain/drivetrain_main.cc
@@ -15,11 +15,13 @@
       aos::configuration::ReadConfig("config.json");
 
   ::aos::ShmEventLoop event_loop(&config.message());
-  ::y2020::control_loops::drivetrain::Localizer localizer(
-      &event_loop, ::y2020::control_loops::drivetrain::GetDrivetrainConfig());
+  std::unique_ptr<::y2020::control_loops::drivetrain::Localizer> localizer =
+      std::make_unique<y2020::control_loops::drivetrain::Localizer>(
+          &event_loop,
+          ::y2020::control_loops::drivetrain::GetDrivetrainConfig());
   std::unique_ptr<DrivetrainLoop> drivetrain = std::make_unique<DrivetrainLoop>(
       ::y2020::control_loops::drivetrain::GetDrivetrainConfig(), &event_loop,
-      &localizer);
+      localizer.get());
 
   event_loop.Run();
 
diff --git a/y2020/joystick_reader.cc b/y2020/joystick_reader.cc
index 9a080fb..cd4af11 100644
--- a/y2020/joystick_reader.cc
+++ b/y2020/joystick_reader.cc
@@ -15,6 +15,7 @@
 #include "frc971/input/driver_station_data.h"
 #include "frc971/input/drivetrain_input.h"
 #include "frc971/input/joystick_input.h"
+#include "frc971/zeroing/wrap.h"
 #include "y2020/constants.h"
 #include "y2020/control_loops/drivetrain/drivetrain_base.h"
 #include "y2020/control_loops/superstructure/superstructure_goal_generated.h"
@@ -48,6 +49,8 @@
 const ButtonLocation kFeedDriver(1, 2);
 const ButtonLocation kIntakeExtend(3, 9);
 const ButtonLocation kIntakeExtendDriver(1, 4);
+const ButtonLocation kRedLocalizerReset(3, 13);
+const ButtonLocation kBlueLocalizerReset(3, 14);
 const ButtonLocation kIntakeIn(4, 4);
 const ButtonLocation kSpit(4, 3);
 const ButtonLocation kLocalizerReset(3, 8);
@@ -76,7 +79,46 @@
     AOS_LOG(INFO, "Auto ended, assuming disc and have piece\n");
   }
 
+  void BlueResetLocalizer() {
+    auto builder = localizer_control_sender_.MakeBuilder();
+
+    frc971::control_loops::drivetrain::LocalizerControl::Builder
+        localizer_control_builder = builder.MakeBuilder<
+            frc971::control_loops::drivetrain::LocalizerControl>();
+    localizer_control_builder.add_x(7.4);
+    localizer_control_builder.add_y(-1.7);
+    localizer_control_builder.add_theta_uncertainty(10.0);
+    localizer_control_builder.add_theta(0.0);
+    localizer_control_builder.add_keep_current_theta(false);
+    if (!builder.Send(localizer_control_builder.Finish())) {
+      AOS_LOG(ERROR, "Failed to reset blue localizer.\n");
+    }
+  }
+
+  void RedResetLocalizer() {
+    auto builder = localizer_control_sender_.MakeBuilder();
+
+    frc971::control_loops::drivetrain::LocalizerControl::Builder
+        localizer_control_builder = builder.MakeBuilder<
+            frc971::control_loops::drivetrain::LocalizerControl>();
+    localizer_control_builder.add_x(-7.4);
+    localizer_control_builder.add_y(1.7);
+    localizer_control_builder.add_theta_uncertainty(10.0);
+    localizer_control_builder.add_theta(M_PI);
+    localizer_control_builder.add_keep_current_theta(false);
+    if (!builder.Send(localizer_control_builder.Finish())) {
+      AOS_LOG(ERROR, "Failed to reset red localizer.\n");
+    }
+  }
+
   void ResetLocalizer() {
+    const frc971::control_loops::drivetrain::Status *drivetrain_status =
+        this->drivetrain_status();
+    if (drivetrain_status == nullptr) {
+      return;
+    }
+    // Get the current position
+    // Snap to heading.
     auto builder = localizer_control_sender_.MakeBuilder();
 
     // Start roughly in front of the red-team goal, robot pointed away from
@@ -84,10 +126,12 @@
     frc971::control_loops::drivetrain::LocalizerControl::Builder
         localizer_control_builder = builder.MakeBuilder<
             frc971::control_loops::drivetrain::LocalizerControl>();
-    localizer_control_builder.add_x(5.0);
-    localizer_control_builder.add_y(-2.0);
-    localizer_control_builder.add_theta(M_PI);
-    localizer_control_builder.add_theta_uncertainty(0.01);
+    localizer_control_builder.add_x(drivetrain_status->x());
+    localizer_control_builder.add_y(drivetrain_status->y());
+    const double new_theta =
+        frc971::zeroing::Wrap(drivetrain_status->theta(), 0, M_PI);
+    localizer_control_builder.add_theta(new_theta);
+    localizer_control_builder.add_theta_uncertainty(10.0);
     if (!builder.Send(localizer_control_builder.Finish())) {
       AOS_LOG(ERROR, "Failed to reset localizer.\n");
     }
@@ -175,10 +219,17 @@
       climber_speed = 12.0f;
     }
 
-    if (data.IsPressed(kLocalizerReset)) {
+    if (data.PosEdge(kLocalizerReset)) {
       ResetLocalizer();
     }
 
+    if (data.PosEdge(kRedLocalizerReset)) {
+      RedResetLocalizer();
+    }
+    if (data.PosEdge(kBlueLocalizerReset)) {
+      BlueResetLocalizer();
+    }
+
     auto builder = superstructure_goal_sender_.MakeBuilder();
 
     flatbuffers::Offset<superstructure::Goal> superstructure_goal_offset;
diff --git a/y2020/vision/tools/python_code/camera_definition.py b/y2020/vision/tools/python_code/camera_definition.py
index 727cf70..b0ce5f3 100644
--- a/y2020/vision/tools/python_code/camera_definition.py
+++ b/y2020/vision/tools/python_code/camera_definition.py
@@ -34,6 +34,57 @@
         self.timestamp = 0
 
 
+def compute_extrinsic(camera_pitch, T_camera, is_turret):
+    # Compute the extrinsic calibration based on pitch and translation
+    # Includes camera rotation from robot x,y,z to opencv (z, -x, -y)
+
+    # Also, handle extrinsics for the turret
+    # The basic camera pose is relative to the center, base of the turret
+    # TODO<Jim>: Maybe store these to .json files, like with intrinsics?
+    base_cam_ext = CameraExtrinsics()
+    turret_cam_ext = CameraExtrinsics()
+
+    camera_pitch_matrix = np.array(
+        [[np.cos(camera_pitch), 0.0, -np.sin(camera_pitch)], [0.0, 1.0, 0.0],
+         [np.sin(camera_pitch), 0.0,
+          np.cos(camera_pitch)]])
+
+    robot_to_camera_rotation = np.array([[0., 0., 1.], [-1, 0, 0], [0, -1.,
+                                                                    0]])
+
+    if is_turret:
+        # Turret camera has default heading 180 deg away from the robot x
+        base_cam_ext.R = np.array([[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0],
+                                   [0.0, 0.0, 1.0]])
+        base_cam_ext.T = np.array([0.0, 0.0, 0.0])
+        turret_cam_ext.R = camera_pitch_matrix @ robot_to_camera_rotation
+        turret_cam_ext.T = T_camera
+    else:
+        base_cam_ext.R = camera_pitch_matrix @ robot_to_camera_rotation
+        base_cam_ext.T = T_camera
+        turret_cam_ext = None
+
+    return base_cam_ext, turret_cam_ext
+
+
+def compute_extrinsic_by_pi(pi_number):
+    # Defaults for non-turret camera
+    camera_pitch = 20.0 * np.pi / 180.0
+    is_turret = False
+    # Default camera location to robot origin
+    T = np.array([0.0, 0.0, 0.0])
+
+    if pi_number == "pi1":
+        # This is the turret camera
+        camera_pitch = 34.0 * np.pi / 180.0
+        is_turret = True
+        T = np.array([7.5 * 0.0254, -5.5 * 0.0254, 41.0 * 0.0254])
+    elif pi_number == "pi2":
+        T = np.array([4.5 * 0.0254, 3.75 * 0.0254, 26.0 * 0.0254])
+
+    return compute_extrinsic(camera_pitch, T, is_turret)
+
+
 def load_camera_definitions():
     ### CAMERA DEFINITIONS
     # We only load in cameras that have a calibration file
@@ -43,36 +94,11 @@
     # Or better yet, use //y2020/vision:calibration to calibrate the camera
     #   using a Charuco target board
 
-    # Extrinsic definition
-    # Camera rotation from robot x,y,z to opencv (z, -x, -y)
-    # This is extrinsics for the turret camera
-    # camera pose relative to center, base of the turret
-    # TODO<Jim>: Need to implement per-camera calibration, like with intrinsics
-    camera_pitch = 34.0 * np.pi / 180.0
-    camera_pitch_matrix = np.matrix(
-        [[np.cos(camera_pitch), 0.0, -np.sin(camera_pitch)], [0.0, 1.0, 0.0],
-         [np.sin(camera_pitch), 0.0,
-          np.cos(camera_pitch)]])
-    turret_cam_ext = CameraExtrinsics()
-    turret_cam_ext.R = np.array(
-        camera_pitch_matrix *
-        np.matrix([[0., 0., 1.], [-1, 0, 0], [0, -1., 0]]))
-    turret_cam_ext.T = np.array([7.5 * 0.0254, -5.5 * 0.0254, 41.0 * 0.0254])
-    default_cam_ext = CameraExtrinsics()
-    default_cam_ext.R = np.array([[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0],
-                                  [0.0, 0.0, 1.0]])
-    default_cam_ext.T = np.array([0.0, 0.0, 0.0])
-
-    default_cam_params = CameraParameters()
-    # Currently, all cameras have this same set of extrinsics
-    default_cam_params.camera_ext = default_cam_ext
-    default_cam_params.turret_ext = turret_cam_ext
-
     camera_list = []
 
     dir_name = dtd.bazel_name_fix('calib_files')
     glog.debug("Searching for calibration files in " + dir_name)
-    for filename in os.listdir(dir_name):
+    for filename in sorted(os.listdir(dir_name)):
         glog.debug("Inspecting %s", filename)
         if ("cam-calib-int" in filename
                 or 'calibration' in filename) and filename.endswith(".json"):
@@ -104,11 +130,15 @@
 
             glog.info("Found calib for " + node_name + ", team #" +
                       str(team_number))
-            camera_base = copy.deepcopy(default_cam_params)
-            camera_base.node_name = node_name
-            camera_base.team_number = team_number
-            camera_base.camera_int.camera_matrix = copy.copy(camera_matrix)
-            camera_base.camera_int.dist_coeffs = copy.copy(dist_coeffs)
-            camera_list.append(camera_base)
+
+            camera_params = CameraParameters()
+            camera_params.camera_ext, camera_params.turret_ext = compute_extrinsic_by_pi(
+                node_name)
+
+            camera_params.node_name = node_name
+            camera_params.team_number = team_number
+            camera_params.camera_int.camera_matrix = copy.copy(camera_matrix)
+            camera_params.camera_int.dist_coeffs = copy.copy(dist_coeffs)
+            camera_list.append(camera_params)
 
     return camera_list