Add imu reciever for the pi

Change-Id: I7f08b37ee0ecd321e4eb41d2a259addee12704fc
Signed-off-by: Ravago Jones <ravagojones@gmail.com>
diff --git a/aos/events/logging/BUILD b/aos/events/logging/BUILD
index d190120..ec13f92 100644
--- a/aos/events/logging/BUILD
+++ b/aos/events/logging/BUILD
@@ -134,10 +134,10 @@
     visibility = ["//visibility:public"],
     deps = [
         ":buffer_encoder",
-        ":crc32",
         ":logger_fbs",
         "//aos:configuration_fbs",
         "//aos/containers:resizeable_buffer",
+        "//aos/util:crc32",
         "@com_github_google_flatbuffers//:flatbuffers",
         "@com_github_google_glog//:glog",
         "@com_google_absl//absl/types:span",
@@ -465,12 +465,3 @@
     gen_reflections = 1,
     target_compatible_with = ["@platforms//os:linux"],
 )
-
-cc_library(
-    name = "crc32",
-    srcs = ["crc32.cc"],
-    hdrs = ["crc32.h"],
-    deps = [
-        "@com_google_absl//absl/types:span",
-    ],
-)
diff --git a/aos/events/logging/snappy_encoder.cc b/aos/events/logging/snappy_encoder.cc
index 4621273..8f46cf3 100644
--- a/aos/events/logging/snappy_encoder.cc
+++ b/aos/events/logging/snappy_encoder.cc
@@ -1,6 +1,6 @@
 #include "aos/events/logging/snappy_encoder.h"
 
-#include "aos/events/logging/crc32.h"
+#include "aos/util/crc32.h"
 #include "external/snappy/snappy.h"
 
 namespace aos::logger {
diff --git a/aos/util/BUILD b/aos/util/BUILD
index 0314069..8df16e2 100644
--- a/aos/util/BUILD
+++ b/aos/util/BUILD
@@ -257,6 +257,15 @@
     ],
 )
 
+cc_library(
+    name = "crc32",
+    srcs = ["crc32.cc"],
+    hdrs = ["crc32.h"],
+    deps = [
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
 py_library(
     name = "python_init",
     srcs = ["__init__.py"],
diff --git a/aos/events/logging/crc32.cc b/aos/util/crc32.cc
similarity index 95%
rename from aos/events/logging/crc32.cc
rename to aos/util/crc32.cc
index cd9d9fb..7f13f30 100644
--- a/aos/events/logging/crc32.cc
+++ b/aos/util/crc32.cc
@@ -1,4 +1,4 @@
-#include "aos/events/logging/crc32.h"
+#include "aos/util/crc32.h"
 
 namespace aos {
 
@@ -53,9 +53,8 @@
   return AccumulateCrc32(data, std::nullopt);
 }
 
-uint32_t AccumulateCrc32(
-    const absl::Span<uint8_t> data,
-    std::optional<uint32_t> current_checksum) {
+uint32_t AccumulateCrc32(const absl::Span<uint8_t> data,
+                         std::optional<uint32_t> current_checksum) {
   uint32_t crc =
       current_checksum.has_value() ? current_checksum.value() : 0xFF'FF'FF'FF;
   for (const uint8_t n : data) {
diff --git a/aos/events/logging/crc32.h b/aos/util/crc32.h
similarity index 100%
rename from aos/events/logging/crc32.h
rename to aos/util/crc32.h
diff --git a/frc971/wpilib/imu.fbs b/frc971/wpilib/imu.fbs
index abd01e2..33bad72 100644
--- a/frc971/wpilib/imu.fbs
+++ b/frc971/wpilib/imu.fbs
@@ -45,6 +45,11 @@
   // Operation section for more details on conditions that may cause this bit to
   // be set to 1.
   data_path_overrun:bool (id: 6);
+
+  // True indicates that the Raspberry Pi Pico recieved a packet
+  // from the imu that had a bad checksum but still sent a message
+  // containing a timestamp and encoder values.
+  checksum_mismatch:bool (id:7);
 }
 
 // Values returned from an IMU.
@@ -87,6 +92,24 @@
   // converted from fpga_timestamp.
   monotonic_timestamp_ns:long (id: 12);
 
+  // The timestamp when the values were captured by the Raspberry Pi Pico.
+  // This has microsecond precision.
+  pico_timestamp_us:int (id:20);
+
+  // The number of this reading produced from a 16-bit counter.
+  data_counter:int (id:19);
+
+  // The number of messages recieved by the Raspberry Pi that had bad checksums
+  failed_checksums:int (id:21);
+  // True if this packet has a bad checksum
+  // from the Raspberry Pi Pico to the Raspberry Pi.
+  checksum_failed:bool (id:22);
+
+  // The position of the left drivetrain encoder in encoder ticks
+  left_encoder:int (id: 17);
+  // The position of the right drivetrain encoder in encoder ticks
+  right_encoder:int (id: 18);
+
   // For an ADIS16470, the DIAG_STAT value immediately after reset.
   start_diag_stat:ADIS16470DiagStat (id: 13);
   // For an ADIS16470, the DIAG_STAT value after the initial sensor self test we
diff --git a/y2022/BUILD b/y2022/BUILD
index 76b3903..78cb86c 100644
--- a/y2022/BUILD
+++ b/y2022/BUILD
@@ -28,6 +28,7 @@
     binaries = [
         "//y2020/vision:calibration",
         "//y2022/vision:viewer",
+        "//y2022/localizer:imu_main",
     ],
     data = [
         ":config",
diff --git a/y2022/localizer/BUILD b/y2022/localizer/BUILD
new file mode 100644
index 0000000..c2afa55
--- /dev/null
+++ b/y2022/localizer/BUILD
@@ -0,0 +1,32 @@
+cc_library(
+    name = "imu",
+    srcs = [
+        "imu.cc",
+    ],
+    hdrs = [
+        "imu.h",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        "//aos/events:epoll",
+        "//aos/events:shm_event_loop",
+        "//aos/util:crc32",
+        "//frc971/wpilib:imu_batch_fbs",
+        "//frc971/wpilib:imu_fbs",
+        "//y2022:constants",
+        "@com_github_google_glog//:glog",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_binary(
+    name = "imu_main",
+    srcs = ["imu_main.cc"],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":imu",
+        "//aos:init",
+        "//aos/events:shm_event_loop",
+    ],
+)
diff --git a/y2022/localizer/imu.cc b/y2022/localizer/imu.cc
new file mode 100644
index 0000000..dbdc1a1
--- /dev/null
+++ b/y2022/localizer/imu.cc
@@ -0,0 +1,165 @@
+#include "y2022/localizer/imu.h"
+
+#include "aos/util/crc32.h"
+#include "glog/logging.h"
+
+namespace y2022::localizer {
+
+namespace {
+
+constexpr size_t kReadSize = 50;
+constexpr double kGyroScale = 1 / 655360.0 / 360.0 * (2 * M_PI);
+constexpr double kAccelScale = 1 / 26756268.0 / 9.80665;
+constexpr double kTempScale = 0.1;
+
+}  // namespace
+
+Imu::Imu(aos::ShmEventLoop *event_loop)
+    : event_loop_(event_loop),
+      imu_sender_(
+          event_loop_->MakeSender<frc971::IMUValuesBatch>("/drivetrain")) {
+  event_loop->SetRuntimeRealtimePriority(30);
+  imu_fd_ = open("/dev/adis16505", O_RDONLY | O_NONBLOCK);
+  PCHECK(imu_fd_ != -1) << ": Failed to open SPI device for IMU.";
+  aos::internal::EPoll *epoll = event_loop_->epoll();
+  epoll->OnReadable(imu_fd_, [this]() {
+    uint8_t buf[kReadSize];
+    ssize_t read_len = read(imu_fd_, buf, kReadSize);
+    // TODO: Do we care about gracefully handling EAGAIN or anything else?
+    // This should only get called when there is data.
+    PCHECK(read_len != -1);
+    CHECK_EQ(read_len, static_cast<ssize_t>(kReadSize))
+        << ": Read incorrect number of bytes.";
+
+    auto sender = imu_sender_.MakeBuilder();
+
+    const flatbuffers::Offset<frc971::IMUValues> values_offset =
+        ProcessReading(sender.fbb(), absl::Span(buf, kReadSize));
+    const flatbuffers::Offset<
+        flatbuffers::Vector<flatbuffers::Offset<frc971::IMUValues>>>
+        readings_offset = sender.fbb()->CreateVector(&values_offset, 1);
+    frc971::IMUValuesBatch::Builder batch_builder =
+        sender.MakeBuilder<frc971::IMUValuesBatch>();
+    batch_builder.add_readings(readings_offset);
+    imu_sender_.CheckOk(sender.Send(batch_builder.Finish()));
+  });
+}
+
+flatbuffers::Offset<frc971::IMUValues> Imu::ProcessReading(
+    flatbuffers::FlatBufferBuilder *fbb, const absl::Span<uint8_t> message) {
+  absl::Span<const uint8_t> buf = message;
+
+  uint64_t driver_timestamp;
+  memcpy(&driver_timestamp, buf.data(), sizeof(driver_timestamp));
+  buf = buf.subspan(8);
+
+  uint16_t diag_stat;
+  memcpy(&diag_stat, buf.data(), sizeof(diag_stat));
+  buf = buf.subspan(2);
+
+  double x_gyro = ConvertValue32(buf, kGyroScale);
+  buf = buf.subspan(4);
+  double y_gyro = ConvertValue32(buf, kGyroScale);
+  buf = buf.subspan(4);
+  double z_gyro = ConvertValue32(buf, kGyroScale);
+  buf = buf.subspan(4);
+  double x_accel = ConvertValue32(buf, kAccelScale);
+  buf = buf.subspan(4);
+  double y_accel = ConvertValue32(buf, kAccelScale);
+  buf = buf.subspan(4);
+  double z_accel = ConvertValue32(buf, kAccelScale);
+  buf = buf.subspan(4);
+  double temp = ConvertValue16(buf, kTempScale);
+  buf = buf.subspan(2);
+  uint16_t data_counter;
+  memcpy(&data_counter, buf.data(), sizeof(data_counter));
+  buf = buf.subspan(2);
+  uint32_t pico_timestamp;
+  memcpy(&pico_timestamp, buf.data(), sizeof(pico_timestamp));
+  buf = buf.subspan(4);
+  int16_t encoder1_count;
+  memcpy(&encoder1_count, buf.data(), sizeof(encoder1_count));
+  buf = buf.subspan(2);
+  int16_t encoder2_count;
+  memcpy(&encoder2_count, buf.data(), sizeof(encoder2_count));
+  buf = buf.subspan(2);
+  uint32_t checksum;
+  memcpy(&checksum, buf.data(), sizeof(checksum));
+  buf = buf.subspan(4);
+
+  CHECK(buf.empty()) << "Have leftover bytes: " << buf.size();
+
+  u_int32_t calculated_checksum = aos::ComputeCrc32(message.subspan(8, 38));
+
+  if (checksum != calculated_checksum) {
+    this->failed_checksums_++;
+  }
+
+  const auto diag_stat_offset = PackDiagStat(fbb, diag_stat);
+
+  frc971::IMUValues::Builder imu_builder(*fbb);
+
+  if (checksum == calculated_checksum) {
+    constexpr uint16_t kChecksumMismatch = 1 << 0;
+    bool imu_checksum_matched = !(diag_stat & kChecksumMismatch);
+
+    // data from the IMU packet
+    if (imu_checksum_matched) {
+      imu_builder.add_gyro_x(x_gyro);
+      imu_builder.add_gyro_y(y_gyro);
+      imu_builder.add_gyro_z(z_gyro);
+
+      imu_builder.add_accelerometer_x(x_accel);
+      imu_builder.add_accelerometer_y(y_accel);
+      imu_builder.add_accelerometer_z(z_accel);
+
+      imu_builder.add_temperature(temp);
+
+      imu_builder.add_data_counter(data_counter);
+    }
+
+    // extra data from the pico
+    imu_builder.add_pico_timestamp_us(pico_timestamp);
+    imu_builder.add_left_encoder(encoder1_count);
+    imu_builder.add_right_encoder(encoder2_count);
+    imu_builder.add_previous_reading_diag_stat(diag_stat_offset);
+  }
+
+  // extra data from us
+  imu_builder.add_monotonic_timestamp_ns(driver_timestamp);
+  imu_builder.add_failed_checksums(failed_checksums_);
+  imu_builder.add_checksum_failed(checksum != calculated_checksum);
+
+  return imu_builder.Finish();
+}
+
+flatbuffers::Offset<frc971::ADIS16470DiagStat> Imu::PackDiagStat(
+    flatbuffers::FlatBufferBuilder *fbb, uint16_t value) {
+  frc971::ADIS16470DiagStat::Builder diag_stat_builder(*fbb);
+  diag_stat_builder.add_clock_error(value & (1 << 7));
+  diag_stat_builder.add_memory_failure(value & (1 << 6));
+  diag_stat_builder.add_sensor_failure(value & (1 << 5));
+  diag_stat_builder.add_standby_mode(value & (1 << 4));
+  diag_stat_builder.add_spi_communication_error(value & (1 << 3));
+  diag_stat_builder.add_flash_memory_update_error(value & (1 << 2));
+  diag_stat_builder.add_data_path_overrun(value & (1 << 1));
+  diag_stat_builder.add_checksum_mismatch(value & (1 << 0));
+  return diag_stat_builder.Finish();
+}
+
+double Imu::ConvertValue32(absl::Span<const uint8_t> data,
+                           double lsb_per_output) {
+  int32_t value;
+  memcpy(&value, data.data(), sizeof(value));
+  return static_cast<double>(value) * lsb_per_output;
+}
+
+double Imu::ConvertValue16(absl::Span<const uint8_t> data,
+                           double lsb_per_output) {
+  int16_t value;
+  memcpy(&value, data.data(), sizeof(value));
+  return static_cast<double>(value) * lsb_per_output;
+}
+
+Imu::~Imu() { PCHECK(0 == close(imu_fd_)); }
+}  // namespace y2022::localizer
diff --git a/y2022/localizer/imu.h b/y2022/localizer/imu.h
new file mode 100644
index 0000000..cd45710
--- /dev/null
+++ b/y2022/localizer/imu.h
@@ -0,0 +1,31 @@
+#ifndef Y2022_LOCALIZER_IMU_H_
+#define Y2022_LOCALIZER_IMU_H_
+#include "aos/events/shm_event_loop.h"
+#include "frc971/wpilib/imu_batch_generated.h"
+#include "y2022/constants.h"
+
+namespace y2022::localizer {
+
+// Reads IMU packets from the kernel driver which reads them over spi
+// from the Raspberry Pi Pico on the IMU board.
+class Imu {
+ public:
+  Imu(aos::ShmEventLoop *event_loop);
+  ~Imu();
+
+ private:
+  flatbuffers::Offset<frc971::ADIS16470DiagStat> PackDiagStat(
+      flatbuffers::FlatBufferBuilder *fbb, uint16_t value);
+  flatbuffers::Offset<frc971::IMUValues> ProcessReading(
+      flatbuffers::FlatBufferBuilder *fbb, absl::Span<uint8_t> buf);
+  double ConvertValue32(absl::Span<const uint8_t> data, double lsb_per_output);
+  double ConvertValue16(absl::Span<const uint8_t> data, double lsb_per_output);
+
+  aos::ShmEventLoop *event_loop_;
+  aos::Sender<frc971::IMUValuesBatch> imu_sender_;
+  int imu_fd_;
+
+  uint failed_checksums_ = 0;
+};
+}  // namespace y2022::localizer
+#endif  // Y2022_LOCALIZER_IMU_H_
diff --git a/y2022/localizer/imu_main.cc b/y2022/localizer/imu_main.cc
new file mode 100644
index 0000000..bba2dd7
--- /dev/null
+++ b/y2022/localizer/imu_main.cc
@@ -0,0 +1,19 @@
+#include "aos/events/shm_event_loop.h"
+#include "aos/init.h"
+#include "y2022/localizer/imu.h"
+
+DEFINE_string(config, "config.json", "Path to the config file to use.");
+
+int main(int argc, char *argv[]) {
+  aos::InitGoogle(&argc, &argv);
+
+  aos::FlatbufferDetachedBuffer<aos::Configuration> config =
+      aos::configuration::ReadConfig(FLAGS_config);
+
+  aos::ShmEventLoop event_loop(&config.message());
+  y2022::localizer::Imu imu(&event_loop);
+
+  event_loop.Run();
+
+  return 0;
+}