Add drivetrain_replay binary for y2024

This lets us debug drivetrain/controller issues offline.

Change-Id: Ifad269cf953370f2173de32162b00548638225c8
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/y2024/constants/simulated_constants_sender.cc b/y2024/constants/simulated_constants_sender.cc
index 28189f8..9b3ee5b 100644
--- a/y2024/constants/simulated_constants_sender.cc
+++ b/y2024/constants/simulated_constants_sender.cc
@@ -6,8 +6,13 @@
 
 namespace y2024 {
 bool SendSimulationConstants(aos::SimulatedEventLoopFactory *factory, int team,
-                             std::string constants_path) {
+                             std::string constants_path,
+                             const std::set<std::string_view> &node_names) {
   for (const aos::Node *node : factory->nodes()) {
+    if (!node_names.empty() &&
+        !node_names.contains(node->name()->string_view())) {
+      continue;
+    }
     std::unique_ptr<aos::EventLoop> event_loop =
         factory->MakeEventLoop("constants_sender", node);
     frc971::constants::ConstantSender<Constants, ConstantsList> sender(
diff --git a/y2024/constants/simulated_constants_sender.h b/y2024/constants/simulated_constants_sender.h
index 5ecc4da..40a5bdd 100644
--- a/y2024/constants/simulated_constants_sender.h
+++ b/y2024/constants/simulated_constants_sender.h
@@ -1,16 +1,21 @@
 #ifndef Y2024_CONSTANTS_SIMULATED_CONFIG_SENDER_H_
 #define Y2024_CONSTANTS_SIMULATED_CONFIG_SENDER_H_
 
+#include <set>
+
 #include "aos/events/simulated_event_loop.h"
 #include "aos/testing/path.h"
 
 namespace y2024 {
 // Returns true, to allow this to be easily called in the initializer list of a
 // constructor.
+// If node_names is specified, we limit ourselves to sending constants on the
+// specified nodes.
 bool SendSimulationConstants(
     aos::SimulatedEventLoopFactory *factory, int team,
     std::string constants_path =
-        aos::testing::ArtifactPath("y2024/constants/test_constants.json"));
+        aos::testing::ArtifactPath("y2024/constants/test_constants.json"),
+    const std::set<std::string_view> &node_names = {});
 }  // namespace y2024
 
 #endif  // Y2024_CONSTANTS_SIMULATED_CONFIG_SENDER_H_
diff --git a/y2024/control_loops/drivetrain/BUILD b/y2024/control_loops/drivetrain/BUILD
index e9f93ca..2a97a7b 100644
--- a/y2024/control_loops/drivetrain/BUILD
+++ b/y2024/control_loops/drivetrain/BUILD
@@ -140,3 +140,29 @@
     parameters = {},
     visibility = ["//visibility:public"],
 )
+
+cc_binary(
+    name = "drivetrain_replay",
+    srcs = ["drivetrain_replay.cc"],
+    data = [
+        "//y2024:aos_config",
+        "//y2024/constants:constants.json",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        "//aos:configuration",
+        "//aos:init",
+        "//aos:json_to_flatbuffer",
+        "//aos/events:simulated_event_loop",
+        "//aos/events/logging:log_reader",
+        "//aos/events/logging:log_writer",
+        "//aos/util:simulation_logger",
+        "//frc971/constants:constants_sender_lib",
+        "//frc971/control_loops/drivetrain:drivetrain_lib",
+        "//frc971/control_loops/drivetrain:trajectory_generator",
+        "//frc971/control_loops/drivetrain/localization:puppet_localizer",
+        "//frc971/imu_fdcan:dual_imu_blender_lib",
+        "//y2024/constants:simulated_constants_sender",
+        "//y2024/control_loops/drivetrain:drivetrain_base",
+    ],
+)
diff --git a/y2024/control_loops/drivetrain/drivetrain_replay.cc b/y2024/control_loops/drivetrain/drivetrain_replay.cc
new file mode 100644
index 0000000..d43b2af
--- /dev/null
+++ b/y2024/control_loops/drivetrain/drivetrain_replay.cc
@@ -0,0 +1,102 @@
+#include "gflags/gflags.h"
+
+#include "aos/configuration.h"
+#include "aos/events/logging/log_reader.h"
+#include "aos/events/logging/log_writer.h"
+#include "aos/events/simulated_event_loop.h"
+#include "aos/init.h"
+#include "aos/json_to_flatbuffer.h"
+#include "aos/network/team_number.h"
+#include "aos/util/simulation_logger.h"
+#include "frc971/constants/constants_sender_lib.h"
+#include "frc971/control_loops/drivetrain/drivetrain.h"
+#include "frc971/control_loops/drivetrain/localization/puppet_localizer.h"
+#include "frc971/control_loops/drivetrain/trajectory_generator.h"
+#include "frc971/imu_fdcan/dual_imu_blender_lib.h"
+#include "y2024/constants/constants_generated.h"
+#include "y2024/constants/simulated_constants_sender.h"
+#include "y2024/control_loops/drivetrain/drivetrain_base.h"
+
+DEFINE_string(config, "y2024/aos_config.json",
+              "Name of the config file to replay using.");
+DEFINE_bool(override_config, false,
+            "If set, override the logged config with --config.");
+DEFINE_int32(team, 971, "Team number to use for logfile replay.");
+DEFINE_string(output_folder, "/tmp/replayed",
+              "Name of the folder to write replayed logs to.");
+DEFINE_string(constants_path, "y2024/constants/constants.json",
+              "Path to the constant file");
+
+int main(int argc, char **argv) {
+  aos::InitGoogle(&argc, &argv);
+
+  aos::network::OverrideTeamNumber(FLAGS_team);
+
+  const aos::FlatbufferDetachedBuffer<aos::Configuration> config =
+      aos::configuration::ReadConfig(FLAGS_config);
+
+  // sort logfiles
+  const std::vector<aos::logger::LogFile> logfiles =
+      aos::logger::SortParts(aos::logger::FindLogs(argc, argv));
+
+  // open logfiles
+  aos::logger::LogReader reader(
+      logfiles, FLAGS_override_config ? &config.message() : nullptr);
+
+  reader.RemapLoggedChannel("/imu/constants", "y2024.Constants");
+  reader.RemapLoggedChannel("/roborio/constants", "y2024.Constants");
+  reader.RemapLoggedChannel("/orin1/constants", "y2024.Constants");
+  reader.RemapLoggedChannel("/drivetrain",
+                            "frc971.control_loops.drivetrain.Status");
+  reader.RemapLoggedChannel("/drivetrain",
+                            "frc971.control_loops.drivetrain.Output");
+  reader.RemapLoggedChannel(
+      "/drivetrain", "frc971.control_loops.drivetrain.RioLocalizerInputs");
+
+  auto factory =
+      std::make_unique<aos::SimulatedEventLoopFactory>(reader.configuration());
+
+  reader.RegisterWithoutStarting(factory.get());
+
+  y2024::SendSimulationConstants(reader.event_loop_factory(), FLAGS_team,
+                                 FLAGS_constants_path, {"roborio"});
+
+  const aos::Node *node = nullptr;
+  if (aos::configuration::MultiNode(reader.configuration())) {
+    node = aos::configuration::GetNode(reader.configuration(), "roborio");
+  }
+  std::vector<std::unique_ptr<aos::util::LoggerState>> loggers;
+
+  std::unique_ptr<aos::EventLoop> drivetrain_event_loop;
+  std::unique_ptr<::frc971::control_loops::drivetrain::PuppetLocalizer>
+      localizer;
+  std::unique_ptr<frc971::control_loops::drivetrain::DrivetrainLoop> drivetrain;
+  reader.OnStart(node, [&factory, node, &loggers, &drivetrain_event_loop,
+                        &localizer, &drivetrain]() {
+    aos::NodeEventLoopFactory *node_factory =
+        factory->GetNodeEventLoopFactory(node);
+    drivetrain_event_loop = node_factory->MakeEventLoop("drivetrain");
+    const auto drivetrain_config =
+        ::y2024::control_loops::drivetrain::GetDrivetrainConfig(
+            drivetrain_event_loop.get());
+    localizer =
+        std::make_unique<::frc971::control_loops::drivetrain::PuppetLocalizer>(
+            drivetrain_event_loop.get(), drivetrain_config);
+    drivetrain =
+        std::make_unique<frc971::control_loops::drivetrain::DrivetrainLoop>(
+            drivetrain_config, drivetrain_event_loop.get(), localizer.get());
+    loggers.push_back(std::make_unique<aos::util::LoggerState>(
+        factory.get(), node, FLAGS_output_folder));
+    // The Trajectory information is NOT_LOGGED, so we need to rerun it against
+    // the SplineGoals that we have in the log.
+    node_factory
+        ->AlwaysStart<frc971::control_loops::drivetrain::TrajectoryGenerator>(
+            "trajectory_generator", drivetrain_config);
+  });
+
+  reader.event_loop_factory()->Run();
+
+  reader.Deregister();
+
+  return 0;
+}