SwerveDrive.hpp

Go to the documentation of this file.

// Copyright (c) 2026 Yet Another Software Suite
// SPDX-License-Identifier: LGPL-3.0-or-later
 
#pragma once
 
#include <cassert>
#include <cstdio>
#include <functional>
#include <optional>
#include <span>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
#include <wpi/commands2/CommandPtr.hpp>
#include <wpi/commands2/Commands.hpp>
#include <wpi/framework/RobotBase.hpp>
#include <wpi/math/estimator/SwerveDrivePoseEstimator.hpp>
#include <wpi/math/geometry/Pose2d.hpp>
#include <wpi/math/geometry/Rotation2d.hpp>
#include <wpi/math/geometry/Rotation3d.hpp>
#include <wpi/math/geometry/Translation2d.hpp>
#include <wpi/math/geometry/Twist2d.hpp>
#include <wpi/math/kinematics/ChassisVelocities.hpp>
#include <wpi/math/kinematics/SwerveDriveKinematics.hpp>
#include <wpi/math/kinematics/SwerveModulePosition.hpp>
#include <wpi/math/kinematics/SwerveModuleVelocity.hpp>
#include <wpi/math/kinematics/struct/ChassisVelocitiesStruct.hpp>
#include <wpi/nt/NetworkTableInstance.hpp>
#include <wpi/smartdashboard/Field2d.hpp>
#include <wpi/system/Timer.hpp>
#include <wpi/tunables/Tunables.hpp>
#include <wpi/units/angle.hpp>
#include <wpi/units/force.hpp>
#include <wpi/units/length.hpp>
#include <wpi/units/math.hpp>
#include <wpi/units/time.hpp>
#include <wpi/units/velocity.hpp>
#include <wpi/util/array.hpp>
 
#include "yams/exceptions.hpp"
#include "yams/mechanisms/swerve/SwerveDriveConfig.hpp"
#include "yams/mechanisms/swerve/SwerveModule.hpp"
#include "yams/telemetry/MechanismTelemetry.hpp"
#include "yams/telemetry/NetworkTablesBackends.hpp"
#include "yams/telemetry/SwerveDriveTelemetry.hpp"
#include "yams/telemetry/SwerveDriveTelemetryConfig.hpp"
 
namespace yams::mechanisms::swerve {
 
template <size_t NumModules = 4>
class SwerveDrive {
 public:
  explicit SwerveDrive(SwerveDriveConfig* config)
      : m_config{config},
        m_kinematics{BuildKinematics(*m_config)},
        m_poseEstimator{m_kinematics, ComputeInitialRotation(*m_config),
                        ComputeInitialPositions(*m_config), m_config->GetInitialPose()} {
    assert(m_config->GetModules().size() == NumModules &&
           "Module count in SwerveDriveConfig must equal NumModules template parameter.");
 
    if (auto optimization = m_config->GetModuleStateOptimization()) {
      for (auto* module : m_config->GetModules()) {
        module->GetMutableConfig().WithOptimization(*optimization);
      }
    }
    m_simPose = m_config->GetInitialPose();
    // Start with the gyro reading the starting pose's heading so field relative driving matches it.
    ResetOdometry(m_config->GetInitialPose());
 
    SetupTelemetry();
 
    m_field2d.SetRobotPose(GetPose());
    telemetry::EnsureMechanismsTunableBackend();
    wpi::tunables::Publish("Mechanisms/" + GetName() + "/field", m_field2d);
  }
 
  // ---- Drive commands --------------------------------------------------------
 
  wpi::cmd::CommandPtr Drive(std::function<wpi::math::ChassisVelocities()> robotRelativeSpeeds) {
    return wpi::cmd::Run(
               [this, robotRelativeSpeeds] {
                 SetRobotRelativeChassisSpeeds(robotRelativeSpeeds());
               },
               {m_config->GetSubsystem()})
        .WithName("Drive");
  }
 
  wpi::cmd::CommandPtr DriveToPose(wpi::math::Pose2d pose) {
    return wpi::cmd::StartRun(
               [this] {
                 ResetTranslationPID();
                 ResetRotationPID();
               },
               [this, pose] { SetRobotRelativeChassisSpeeds(DriveToPoseSetpoint(pose)); },
               {m_config->GetSubsystem()})
        .FinallyDo([this] { Stop(); })
        .WithName("Drive to Pose");
  }
 
  wpi::cmd::CommandPtr DriveToPose(wpi::math::Pose2d pose,
                                   wpi::units::meter_t translationTolerance,
                                   wpi::units::degree_t rotationTolerance) {
    return wpi::cmd::StartRun(
               [this] {
                 ResetTranslationPID();
                 ResetRotationPID();
               },
               [this, pose] { SetRobotRelativeChassisSpeeds(DriveToPoseSetpoint(pose)); },
               {m_config->GetSubsystem()})
        .Until([this, pose, translationTolerance, rotationTolerance] {
          return IsNear(pose, translationTolerance, rotationTolerance);
        })
        .FinallyDo([this] { Stop(); })
        .WithName("Drive to Pose");
  }
 
  bool IsNear(wpi::math::Pose2d pose, wpi::units::meter_t translationTolerance,
              wpi::units::degree_t rotationTolerance) {
    return GetDistanceFromPose(pose) <= translationTolerance &&
           wpi::units::math::abs(GetAngleDifferenceFromPose(pose)) <= rotationTolerance;
  }
 
  void Stop() { SetRobotRelativeChassisSpeeds(wpi::math::ChassisVelocities{}); }
 
  wpi::math::ChassisVelocities DriveToPoseSetpoint(wpi::math::Pose2d targetPose) {
    return DriveToPoseSetpoint(targetPose, GetPose(),
                               wpi::math::Rotation2d{wpi::units::radian_t{GetGyroAngle()}});
  }
 
  wpi::math::ChassisVelocities DriveToPoseSetpoint(wpi::math::Pose2d targetPose,
                                                   wpi::math::Pose2d currentPose,
                                                   wpi::math::Rotation2d heading) {
    auto rotationPIDRef = m_config->GetRotationPID();
    if (!rotationPIDRef) {
      throw exceptions::SwerveDriveConfigurationException(
          "No rotation PID controller configured, cannot drive to pose. Use "
          "SwerveDriveConfig::WithRotationController(PIDController) to fix this error.");
    }
    auto translationPIDRef = m_config->GetTranslationPID();
    if (!translationPIDRef) {
      throw exceptions::SwerveDriveConfigurationException(
          "No translation PID controller configured, cannot drive to pose. Use "
          "SwerveDriveConfig::WithTranslationController(PIDController) to fix this error.");
    }
    auto& rotationPID = rotationPIDRef->get();
    auto& translationPID = translationPIDRef->get();
    // Plain field-frame translation delta (not RelativeTo/operator-, which expresses the result
    // in targetPose's rotated frame and would skew the commanded direction whenever targetPose's
    // heading is non-zero).
    auto toTarget = targetPose.Translation() - currentPose.Translation();
    double distance = toTarget.Norm().value();
    // The controller sees the distance as its measurement with a setpoint of 0, so its output is
    // negative while there is distance left: the speed toward the target is the negative of it.
    double speed = -translationPID.Calculate(distance, 0.0);
    // Drive at that speed along the unit vector toward the target.
    double vx = distance > 1e-9 ? toTarget.X().value() / distance * speed : 0.0;
    double vy = distance > 1e-9 ? toTarget.Y().value() / distance * speed : 0.0;
    return (wpi::math::ChassisVelocities{
                wpi::units::meters_per_second_t{vx}, wpi::units::meters_per_second_t{vy},
                wpi::units::radians_per_second_t{
                    rotationPID.Calculate(currentPose.Rotation().Radians().value(),
                                          targetPose.Rotation().Radians().value())}})
        .ToRobotRelative(heading);
  }
 
  // ---- Core drive methods ---------------------------------------------------
 
  void SetSwerveModuleStates(wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules> states) {
    for (size_t i = 0; i < NumModules; ++i) {
      m_desiredModuleStates[i] = m_config->GetModules()[i]->SetSwerveModuleState(states[i]);
    }
  }
 
  void SetSwerveModuleStates(wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules> states,
                             std::span<const wpi::units::newton_t> feedforwardForces) {
    for (size_t i = 0; i < NumModules; ++i) {
      m_desiredModuleStates[i] =
          m_config->GetModules()[i]->SetSwerveModuleState(states[i], feedforwardForces[i]);
    }
  }
 
  wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules>
  GetStateFromRobotRelativeChassisSpeeds(wpi::math::ChassisVelocities robotRelativeSpeeds) {
    robotRelativeSpeeds = m_config->OptimizeRobotRelativeChassisSpeeds(robotRelativeSpeeds);
    auto cor = m_config->GetCenterOfRotation();
    auto states = cor ? m_kinematics.ToSwerveModuleVelocities(robotRelativeSpeeds, *cor)
                      : m_kinematics.ToSwerveModuleVelocities(robotRelativeSpeeds);
    // Scale the wheels down together when one is asked to go faster than it can.
    if (auto maxModuleSpeed = m_config->GetMaximumModuleLinearVelocity()) {
      auto maxLinear = m_config->GetMaximumChassisLinearVelocity();
      auto maxAngular = m_config->GetMaximumChassisAngularVelocity();
      if (maxLinear && maxAngular) {
        states = wpi::math::SwerveDriveKinematics<NumModules>::DesaturateWheelVelocities(
            states, robotRelativeSpeeds, *maxModuleSpeed, *maxLinear,
            wpi::units::radians_per_second_t{*maxAngular});
      } else {
        states = wpi::math::SwerveDriveKinematics<NumModules>::DesaturateWheelVelocities(
            states, *maxModuleSpeed);
      }
    }
    return states;
  }
 
  wpi::math::ChassisVelocities GetRobotRelativeChassisSpeedsFromState(
      const wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules>& states) {
    return m_kinematics.ToChassisVelocities(states);
  }
 
  void SetRobotRelativeChassisSpeeds(wpi::math::ChassisVelocities robotRelativeSpeeds) {
    m_desiredChassisSpeeds = robotRelativeSpeeds;
    SetSwerveModuleStates(GetStateFromRobotRelativeChassisSpeeds(robotRelativeSpeeds));
  }
 
  void SetRobotRelativeChassisSpeeds(wpi::math::ChassisVelocities robotRelativeSpeeds,
                                     std::span<const wpi::units::newton_t> feedforwardForces) {
    m_desiredChassisSpeeds = robotRelativeSpeeds;
    auto states = GetStateFromRobotRelativeChassisSpeeds(robotRelativeSpeeds);
    if (feedforwardForces.empty()) {
      SetSwerveModuleStates(states);
      return;
    }
    if (feedforwardForces.size() != NumModules) {
      throw std::invalid_argument("feedforwardForces must be empty or have one entry per module (" +
                                  std::to_string(NumModules) + "), in FL, FR, BL, BR order.");
    }
    SetSwerveModuleStates(states, feedforwardForces);
  }
 
  void SetFieldRelativeChassisSpeeds(wpi::math::ChassisVelocities fieldRelativeSpeeds) {
    SetRobotRelativeChassisSpeeds(
        (fieldRelativeSpeeds)
            .ToRobotRelative(wpi::math::Rotation2d{wpi::units::radian_t{GetGyroAngle()}}));
  }
 
  void LockPose() {
    // Turn each wheel toward the robot center without driving it. SwerveModule::SetAzimuthAngle
    // skips the velocity deadband, which would otherwise hold these zero speed states where the
    // wheels are.
    for (size_t i = 0; i < NumModules; ++i) {
      auto* module = m_config->GetModules()[i];
      m_desiredModuleStates[i] = module->SetAzimuthAngle(
          module->GetConfig().GetLocation()->Angle().value_or(wpi::math::Rotation2d{}));
    }
    m_desiredChassisSpeeds = wpi::math::ChassisVelocities{};
  }
 
  void SetAzimuthAngles(wpi::math::Rotation2d angle) {
    for (size_t i = 0; i < NumModules; ++i) {
      m_desiredModuleStates[i] = m_config->GetModules()[i]->SetAzimuthAngle(angle);
    }
    m_desiredChassisSpeeds = wpi::math::ChassisVelocities{};
  }
 
  // ---- Odometry & pose -------------------------------------------------------
 
  wpi::math::Pose2d GetPose() { return m_poseEstimator.GetEstimatedPosition(); }
 
  void ResetOdometry(wpi::math::Pose2d pose) {
    SetGyroAngle(pose.Rotation().Degrees());
    m_poseEstimator.ResetPosition(wpi::math::Rotation2d{wpi::units::radian_t{GetGyroAngle()}},
                                  GetModulePositions(), pose);
    m_desiredChassisSpeeds = wpi::math::ChassisVelocities{};
    m_desiredModuleStates = m_kinematics.ToSwerveModuleVelocities(wpi::math::ChassisVelocities{});
    m_simPose = pose;
  }
 
  void ZeroGyro() {
    // ResetOdometry also sets the gyro to read the new heading.
    ResetOdometry(wpi::math::Pose2d{GetPose().Translation(), wpi::math::Rotation2d{}});
  }
 
  void AddVisionMeasurement(wpi::math::Pose2d robotPose, wpi::units::second_t timestamp) {
    m_poseEstimator.AddVisionMeasurement(robotPose, timestamp);
  }
 
  void AddVisionMeasurement(wpi::math::Pose2d robotPose, wpi::units::second_t timestamp,
                            const wpi::util::array<double, 3>& stdDevs) {
    m_poseEstimator.AddVisionMeasurement(robotPose, timestamp, stdDevs);
  }
 
  // ---- Telemetry & simulation -----------------------------------------------
 
  void UpdateTelemetry() {
    UpdatePoseEstimator();
    auto pose = GetPose();
    auto currentStates = GetModuleStates();
 
    m_swerveTelemetry->Publish(this);
 
    for (auto* mod : m_config->GetModules()) {
      mod->UpdateTelemetry();
    }
    m_telemetry.UpdateLoopTime();
 
    m_field2d.SetRobotPose(pose);
    std::vector<wpi::math::Pose2d> modulePoses;
    modulePoses.reserve(NumModules);
    for (size_t i = 0; i < NumModules; ++i) {
      auto* mod = m_config->GetModules()[i];
      auto location = *mod->GetConfig().GetLocation();
      auto moduleTranslation = pose.Translation() + location.RotateBy(pose.Rotation());
      auto moduleHeading = pose.Rotation() + currentStates[i].angle;
      modulePoses.push_back(wpi::math::Pose2d{moduleTranslation, moduleHeading});
    }
    m_field2d.GetObject("modules")->SetPoses(modulePoses);
  }
 
  void SimIterate() {
    if (!m_simTimer.IsRunning()) m_simTimer.Start();
    for (auto* mod : m_config->GetModules()) {
      mod->SimIterate();
    }
    auto desired = m_kinematics.ToChassisVelocities(m_desiredModuleStates);
 
    auto dt = m_simTimer.Get();
    wpi::math::Twist2d twist{desired.vx * dt, desired.vy * dt, desired.omega * dt};
    m_simPose = m_simPose + twist.Exp();
 
    auto speeds = m_kinematics.ToChassisVelocities(GetModuleStates());
    m_simGyroAngle +=
        wpi::units::degree_t{wpi::units::degrees_per_second_t{speeds.omega}.value() * dt.value()};
    m_simTimer.Reset();
  }
 
  // ---- Queries ---------------------------------------------------------------
 
  wpi::units::degree_t GetGyroAngle() {
    if (wpi::RobotBase::IsSimulation()) return m_simGyroAngle;
    return m_config->GetGyroAngle();
  }
 
  wpi::math::Rotation3d GetGyroRotation3d() {
    if (wpi::RobotBase::IsSimulation()) {
      return wpi::math::Rotation3d{wpi::units::radian_t{0}, wpi::units::radian_t{0},
                                   wpi::units::radian_t{m_simGyroAngle}};
    }
    return m_config->GetGyroRotation3d();
  }
 
  void StartDriveToPoseTuning() {
    ResetTranslationPID();
    ResetRotationPID();
  }
 
  void ApplyTuningValues() {
    if (m_swerveTelemetry) m_swerveTelemetry->ApplyTuningValues(this);
  }
 
  wpi::math::Pose2d GetSimPose() { return m_simPose; }
 
  wpi::math::ChassisVelocities GetRobotRelativeSpeed() {
    return m_kinematics.ToChassisVelocities(GetModuleStates());
  }
 
  wpi::math::ChassisVelocities GetFieldRelativeSpeed() {
    return GetRobotRelativeSpeed().ToFieldRelative(
        wpi::math::Rotation2d{wpi::units::radian_t{GetGyroAngle()}});
  }
 
  wpi::util::array<wpi::math::SwerveModulePosition, NumModules> GetModulePositions() {
    wpi::util::array<wpi::math::SwerveModulePosition, NumModules> positions{wpi::util::empty_array};
    for (size_t i = 0; i < NumModules; ++i) {
      positions[i] = m_config->GetModules()[i]->GetPosition();
    }
    return positions;
  }
 
  wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules> GetModuleStates() {
    wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules> states{wpi::util::empty_array};
    for (size_t i = 0; i < NumModules; ++i) {
      states[i] = m_config->GetModules()[i]->GetState();
    }
    return states;
  }
 
  wpi::units::meter_t GetDistanceFromPose(wpi::math::Pose2d pose) {
    return wpi::units::meter_t{GetPose().Translation().Distance(pose.Translation())};
  }
 
  wpi::units::degree_t GetAngleDifferenceFromPose(wpi::math::Pose2d pose) {
    return (GetPose() - pose).Rotation().Degrees();
  }
 
  std::optional<SwerveModule*> GetModule(const std::string& name) {
    for (auto* mod : m_config->GetModules()) {
      if (mod->GetName() == name) return mod;
    }
    return std::nullopt;
  }
 
  void ResetRotationPID() {
    if (auto pid = m_config->GetRotationPID()) pid->get().Reset();
  }
 
  void ResetTranslationPID() {
    if (auto pid = m_config->GetTranslationPID()) pid->get().Reset();
  }
 
  void SetRotationPID(wpi::math::PIDController controller) {
    auto currentRotationPID = m_config->GetRotationPID();
    if (!currentRotationPID || currentRotationPID->get().GetP() != controller.GetP() ||
        currentRotationPID->get().GetI() != controller.GetI() ||
        currentRotationPID->get().GetD() != controller.GetD()) {
      controller.Reset();
      m_config->WithRotationController(std::move(controller));
    }
  }
 
  void SetTranslationPID(wpi::math::PIDController controller) {
    auto currentTranslationPID = m_config->GetTranslationPID();
    if (!currentTranslationPID || currentTranslationPID->get().GetP() != controller.GetP() ||
        currentTranslationPID->get().GetI() != controller.GetI() ||
        currentTranslationPID->get().GetD() != controller.GetD()) {
      controller.Reset();
      m_config->WithTranslationController(std::move(controller));
    }
  }
 
  void SetVisionMeasurementStdDevs(const wpi::util::array<double, 3>& stdDevs) {
    m_poseEstimator.SetVisionMeasurementStdDevs(stdDevs);
  }
 
  std::string GetName() const { return m_config->GetTelemetryName(); }
 
  SwerveDriveConfig& GetConfig() { return *m_config; }
 
  wpi::math::SwerveDriveKinematics<NumModules>& GetKinematics() { return m_kinematics; }
 
  wpi::Field2d& GetField2d() { return m_field2d; }
 
  wpi::math::ChassisVelocities GetDesiredChassisSpeeds() { return m_desiredChassisSpeeds; }
 
  wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules> GetDesiredModuleStates() {
    return m_desiredModuleStates;
  }
 
 private:
  SwerveDriveConfig* m_config{nullptr};
  wpi::math::SwerveDriveKinematics<NumModules> m_kinematics;
  wpi::math::SwerveDrivePoseEstimator<NumModules> m_poseEstimator;
 
  telemetry::MechanismTelemetry m_telemetry;
  std::optional<telemetry::SwerveDriveTelemetry> m_swerveTelemetry;
  std::optional<wpi::cmd::CommandPtr> m_driveToPoseTuningCommand;
 
  wpi::Field2d m_field2d;
  wpi::Timer m_simTimer;
  wpi::units::degree_t m_simGyroAngle{0};
  wpi::math::Pose2d m_simPose;
 
  wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules> m_desiredModuleStates{
      wpi::util::empty_array};
  wpi::math::ChassisVelocities m_desiredChassisSpeeds{};
 
  void SetupTelemetry() {
    using TelemetryVerbosity = motorcontrollers::SmartMotorControllerConfig::TelemetryVerbosity;
    telemetry::SwerveDriveTelemetryConfig telemetryCfg;
    if (auto specified = m_config->GetSwerveDriveTelemetryConfig()) {
      telemetryCfg = std::move(*specified);
    } else {
      telemetryCfg.WithTelemetryVerbosity(
          m_config->GetTelemetryVerbosity().value_or(TelemetryVerbosity::HIGH));
    }
 
    if (auto dataLogName = telemetryCfg.GetDataLogName()) {
      m_telemetry.SetupTelemetry(GetName(), *dataLogName);
    } else {
      m_telemetry.SetupTelemetry(GetName());
    }
 
    m_swerveTelemetry.emplace(std::move(telemetryCfg));
    m_swerveTelemetry->SetupTelemetry(this);
 
    // Drive to pose tuning needs both controllers, so only offer it when they are configured.
    if (m_config->GetTranslationPID() && m_config->GetRotationPID()) {
      auto tuningCommand = wpi::cmd::StartRun(
          [this] {
            std::puts(
                "================= Starting SwerveDrive.driveToPoseTuning() =================\n");
            ResetTranslationPID();
            ResetRotationPID();
          },
          [this] { m_swerveTelemetry->ApplyTuningValues(this); });
      // Stop any tuning voltage when the command ends.
      m_driveToPoseTuningCommand.emplace(
          std::move(tuningCommand).FinallyDo([this] { m_swerveTelemetry->StopTuning(this); }));
      telemetry::EnsureMechanismsTunableBackend();
      wpi::tunables::Publish("Mechanisms/" + GetName() + "/tuning",
                             *m_driveToPoseTuningCommand->get());
    }
  }
 
  void UpdatePoseEstimator() {
    m_poseEstimator.Update(wpi::math::Rotation2d{wpi::units::radian_t{GetGyroAngle()}},
                           GetModulePositions());
  }
 
  static wpi::math::SwerveDriveKinematics<NumModules> BuildKinematics(
      const SwerveDriveConfig& config) {
    assert(config.GetModules().size() == NumModules);
    wpi::util::array<wpi::math::Translation2d, NumModules> locations{wpi::util::empty_array};
    for (size_t i = 0; i < NumModules; ++i) {
      locations[i] = *config.GetModules()[i]->GetConfig().GetLocation();
    }
    return wpi::math::SwerveDriveKinematics<NumModules>{locations};
  }
 
  void SetGyroAngle(wpi::units::degree_t heading) {
    if (wpi::RobotBase::IsSimulation()) {
      m_simGyroAngle = heading;
      return;
    }
    m_config->WithGyroOffset(GetGyroAngle() + m_config->GetGyroOffset() - heading);
  }
 
  static wpi::math::Rotation2d ComputeInitialRotation(const SwerveDriveConfig& config) {
    if (wpi::RobotBase::IsSimulation()) return wpi::math::Rotation2d{};
    return wpi::math::Rotation2d{wpi::units::radian_t{config.GetGyroAngle()}};
  }
 
  static wpi::util::array<wpi::math::SwerveModulePosition, NumModules> ComputeInitialPositions(
      const SwerveDriveConfig& config) {
    wpi::util::array<wpi::math::SwerveModulePosition, NumModules> positions{wpi::util::empty_array};
    for (size_t i = 0; i < NumModules; ++i) {
      positions[i] = config.GetModules()[i]->GetPosition();
    }
    return positions;
  }
};
 
}  // namespace yams::mechanisms::swerve
 
namespace yams::telemetry {
 
// SwerveDriveTelemetry::SetupTelemetry()/Publish()/ApplyTuningValues() are templated on the
// SwerveDrive's module count so that SwerveDriveTelemetry itself need not be a template. Their
// bodies live here after SwerveDrive<NumModules> is fully defined above because this is the
// only place both SwerveDriveTelemetry's private members and SwerveDrive<NumModules>'s interface
// are simultaneously visible.
 
template <std::size_t NumModules>
void SwerveDriveTelemetry::SetupTelemetry(mechanisms::swerve::SwerveDrive<NumModules>* drive) {
  auto driveName = drive->GetName();
  for (auto* module : drive->GetConfig().GetModules()) {
    module->SetupTelemetry(driveName);
  }
 
  auto inst = wpi::nt::NetworkTableInstance::GetDefault();
  m_dataTable = inst.GetTable("Mechanisms")->GetSubTable(driveName);
  m_tuningTable = inst.GetTable("Tuning")->GetSubTable(driveName);
 
  bool nt4Enabled = m_config.GetNT4Enabled();
  auto dataLogName = m_config.GetDataLogName();
 
  {
    // Only publish the drive to pose gains and auto-align fields for controllers that exist.
    auto translationPID = drive->GetConfig().GetTranslationPID();
    auto rotationPID = drive->GetConfig().GetRotationPID();
    auto& doubleFields = m_config.GetDoubleFields();
    if (translationPID) {
      doubleFields.at(DoubleTelemetryField::TranslationP)
          .SetDefaultValue(translationPID->get().GetP());
      doubleFields.at(DoubleTelemetryField::TranslationI)
          .SetDefaultValue(translationPID->get().GetI());
      doubleFields.at(DoubleTelemetryField::TranslationD)
          .SetDefaultValue(translationPID->get().GetD());
    } else {
      doubleFields.at(DoubleTelemetryField::TranslationP).Disable();
      doubleFields.at(DoubleTelemetryField::TranslationI).Disable();
      doubleFields.at(DoubleTelemetryField::TranslationD).Disable();
    }
    if (rotationPID) {
      doubleFields.at(DoubleTelemetryField::RotationP).SetDefaultValue(rotationPID->get().GetP());
      doubleFields.at(DoubleTelemetryField::RotationI).SetDefaultValue(rotationPID->get().GetI());
      doubleFields.at(DoubleTelemetryField::RotationD).SetDefaultValue(rotationPID->get().GetD());
    } else {
      doubleFields.at(DoubleTelemetryField::RotationP).Disable();
      doubleFields.at(DoubleTelemetryField::RotationI).Disable();
      doubleFields.at(DoubleTelemetryField::RotationD).Disable();
    }
    // Auto-align drives to a pose, which needs both controllers.
    if (!translationPID || !rotationPID) {
      doubleFields.at(DoubleTelemetryField::AutoAlignPoseX).Disable();
      doubleFields.at(DoubleTelemetryField::AutoAlignPoseY).Disable();
      doubleFields.at(DoubleTelemetryField::AutoAlignPoseRotation).Disable();
      m_config.GetBoolFields().at(BooleanTelemetryField::AutoAlignEnabled).Disable();
    }
 
    auto& modules = drive->GetConfig().GetModules();
    if (!modules.empty()) {
      auto* driveMotor = modules[0]->GetDriveMotorController();
      auto driveGains =
          driveMotor->GetConfig().GetSlotGains(driveMotor->GetClosedLoopControllerSlot());
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesDriveP)
          .SetDefaultValue(driveGains.kP);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesDriveI)
          .SetDefaultValue(driveGains.kI);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesDriveD)
          .SetDefaultValue(driveGains.kD);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesDriveKs)
          .SetDefaultValue(driveGains.kS);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesDriveKv)
          .SetDefaultValue(driveGains.kV);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesDriveKa)
          .SetDefaultValue(driveGains.kA);
 
      auto* azimuthMotor = modules[0]->GetAzimuthMotorController();
      auto azimuthGains =
          azimuthMotor->GetConfig().GetSlotGains(azimuthMotor->GetClosedLoopControllerSlot());
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesAzimuthP)
          .SetDefaultValue(azimuthGains.kP);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesAzimuthI)
          .SetDefaultValue(azimuthGains.kI);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesAzimuthD)
          .SetDefaultValue(azimuthGains.kD);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesAzimuthKs)
          .SetDefaultValue(azimuthGains.kS);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesAzimuthKv)
          .SetDefaultValue(azimuthGains.kV);
      m_config.GetDoubleFields()
          .at(DoubleTelemetryField::ModulesAzimuthKa)
          .SetDefaultValue(azimuthGains.kA);
    }
  }
 
  for (auto& [field, dt] : m_config.GetDoubleFields()) {
    if (!dt.IsEnabled()) continue;
    if (nt4Enabled) dt.SetupNetworkTables(m_dataTable, m_tuningTable);
    if (dataLogName) dt.SetupDataLog(*dataLogName);
  }
  for (auto& [field, bt] : m_config.GetBoolFields()) {
    if (!bt.IsEnabled()) continue;
    if (nt4Enabled) bt.SetupNetworkTables(m_dataTable, m_tuningTable);
    if (dataLogName) bt.SetupDataLog(*dataLogName);
  }
  for (auto& [field, stt] : m_config.GetPoseFields()) {
    if (!stt.IsEnabled()) continue;
    if (nt4Enabled) stt.SetupNetworkTables(m_dataTable, m_tuningTable);
    if (dataLogName) stt.SetupDataLog(*dataLogName);
  }
  for (auto& [field, stt] : m_config.GetChassisSpeedsFields()) {
    if (!stt.IsEnabled()) continue;
    if (nt4Enabled) stt.SetupNetworkTables(m_dataTable, m_tuningTable);
    if (dataLogName) stt.SetupDataLog(*dataLogName);
  }
  for (auto& [field, stat] : m_config.GetModuleStatesFields()) {
    if (!stat.IsEnabled()) continue;
    if (nt4Enabled) stat.SetupNetworkTables(m_dataTable, m_tuningTable);
    if (dataLogName) stat.SetupDataLog(*dataLogName);
  }
}
 
template <std::size_t NumModules>
void SwerveDriveTelemetry::Publish(mechanisms::swerve::SwerveDrive<NumModules>* drive) {
  for (auto& [field, stt] : m_config.GetPoseFields()) {
    if (!stt.IsEnabled()) continue;
    switch (field) {
      case StructTelemetryField::Pose:
        stt.Set(drive->GetPose());
        break;
      default:
        break;
    }
  }
  for (auto& [field, stt] : m_config.GetChassisSpeedsFields()) {
    if (!stt.IsEnabled()) continue;
    switch (field) {
      case StructTelemetryField::DesiredRobotRelativeChassisSpeeds:
        stt.Set(drive->GetDesiredChassisSpeeds());
        break;
      case StructTelemetryField::CurrentRobotRelativeChassisSpeeds:
        stt.Set(drive->GetRobotRelativeSpeed());
        break;
      case StructTelemetryField::FieldRelativeChassisSpeeds:
        stt.Set(drive->GetFieldRelativeSpeed());
        break;
      default:
        break;
    }
  }
  for (auto& [field, dt] : m_config.GetDoubleFields()) {
    if (!dt.IsEnabled()) continue;
    switch (field) {
      case DoubleTelemetryField::Gyro:
        dt.Set(drive->GetGyroAngle().value());
        break;
      default:
        break;
    }
  }
  for (auto& [field, stat] : m_config.GetModuleStatesFields()) {
    if (!stat.IsEnabled()) continue;
    switch (field) {
      case StructArrayTelemetryField::DesiredModuleStates: {
        auto states = drive->GetDesiredModuleStates();
        stat.Set(std::vector<wpi::math::SwerveModuleVelocity>(states.begin(), states.end()));
        break;
      }
      case StructArrayTelemetryField::CurrentModuleStates: {
        auto states = drive->GetModuleStates();
        stat.Set(std::vector<wpi::math::SwerveModuleVelocity>(states.begin(), states.end()));
        break;
      }
      default:
        break;
    }
  }
}
 
template <std::size_t NumModules>
void SwerveDriveTelemetry::ApplyTuningValues(mechanisms::swerve::SwerveDrive<NumModules>* drive) {
  if (!m_config.GetNT4Enabled()) return;
  auto& modules = drive->GetConfig().GetModules();
  auto& boolFields = m_config.GetBoolFields();
  auto& doubleFields = m_config.GetDoubleFields();
  // Copies of the configured controllers; gain fields for missing controllers are disabled.
  std::optional<wpi::math::PIDController> translationPID;
  if (auto pid = drive->GetConfig().GetTranslationPID()) translationPID = pid->get();
  std::optional<wpi::math::PIDController> rotationPID;
  if (auto pid = drive->GetConfig().GetRotationPID()) rotationPID = pid->get();
 
  // Read which tuning modes are on. Only one may drive the chassis at a time: auto-align, then
  // drive tuning, then azimuth tuning. A mode that is on turns the ones it outranks back off, in
  // NetworkTables too so the dashboard reflects what's actually happening, so the order the fields
  // are read in does not matter. Auto-align fields are disabled when either drive to pose
  // controller is missing.
  std::set<TuningMode> modes;
  for (auto& [field, bt] : boolFields) {
    if (!bt.IsEnabled() || !bt.Get()) continue;
    switch (field) {
      case BooleanTelemetryField::AutoAlignEnabled:
        modes.insert(TuningMode::AutoAlign);
        modes.erase(TuningMode::Drive);
        modes.erase(TuningMode::Azimuth);
        boolFields.at(BooleanTelemetryField::ModulesDriveTuningEnabled).ForceSet(false);
        boolFields.at(BooleanTelemetryField::ModulesAzimuthTuningEnabled).ForceSet(false);
        break;
      case BooleanTelemetryField::ModulesDriveTuningEnabled:
        if (!modes.contains(TuningMode::AutoAlign)) {
          modes.insert(TuningMode::Drive);
          modes.erase(TuningMode::Azimuth);
          boolFields.at(BooleanTelemetryField::ModulesAzimuthTuningEnabled).ForceSet(false);
        }
        break;
      case BooleanTelemetryField::ModulesDriveInPlace:
        modes.insert(TuningMode::DriveInPlace);
        break;
      case BooleanTelemetryField::ModulesDriveVoltageEnabled:
        modes.insert(TuningMode::DriveVoltage);
        break;
      case BooleanTelemetryField::ModulesAzimuthTuningEnabled:
        if (!modes.contains(TuningMode::AutoAlign) && !modes.contains(TuningMode::Drive)) {
          modes.insert(TuningMode::Azimuth);
        }
        break;
      case BooleanTelemetryField::ModulesAzimuthVoltageEnabled:
        modes.insert(TuningMode::AzimuthVoltage);
        break;
    }
  }
  // Spinning in place and a voltage only apply while their tuning mode is on.
  if (!modes.contains(TuningMode::Drive)) {
    modes.erase(TuningMode::DriveInPlace);
    modes.erase(TuningMode::DriveVoltage);
  }
  if (!modes.contains(TuningMode::Azimuth)) modes.erase(TuningMode::AzimuthVoltage);
  // The angle each module points while drive tuning: forward, or, to spin in place, tangent to its
  // position around the robot center so a positive velocity spins the robot counter-clockwise
  // (WPILib's positive rotation direction).
  wpi::util::array<wpi::math::Rotation2d, NumModules> driveAngles{wpi::util::empty_array};
  for (size_t i = 0; i < NumModules; ++i) {
    driveAngles[i] = wpi::math::Rotation2d{};
    if (modes.contains(TuningMode::DriveInPlace)) {
      auto moduleLocation = *modules[i]->GetConfig().GetLocation();
      driveAngles[i] = moduleLocation.Angle().value_or(wpi::math::Rotation2d{}) +
                       wpi::math::Rotation2d{wpi::units::degree_t{90}};
    }
  }
 
  // Auto-align drives to its target pose every loop.
  if (modes.contains(TuningMode::AutoAlign)) {
    wpi::math::Pose2d targetPose{
        wpi::units::meter_t{doubleFields.at(DoubleTelemetryField::AutoAlignPoseX).Get()},
        wpi::units::meter_t{doubleFields.at(DoubleTelemetryField::AutoAlignPoseY).Get()},
        wpi::math::Rotation2d{
            wpi::units::degree_t{doubleFields.at(DoubleTelemetryField::AutoAlignPoseRotation).Get()}}};
    drive->SetRobotRelativeChassisSpeeds(drive->DriveToPoseSetpoint(targetPose));
  }
 
  // When the tuning modes change, switch the motors over, and send the setpoints of the modes now
  // on.
  if (modes != m_lastModes) {
    const std::set<TuningMode> lastModes = m_lastModes.value_or(std::set<TuningMode>{});
    m_lastModes = modes;
    // Back to the closed loops when a voltage turns off.
    for (auto* module : modules) {
      if (lastModes.contains(TuningMode::DriveVoltage) && !modes.contains(TuningMode::DriveVoltage)) {
        module->GetDriveMotorController()->StartClosedLoopController();
      }
      if (lastModes.contains(TuningMode::AzimuthVoltage) &&
          !modes.contains(TuningMode::AzimuthVoltage)) {
        module->GetAzimuthMotorController()->StartClosedLoopController();
      }
    }
    wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules> states{wpi::util::empty_array};
    if (modes.empty()) {
      // No tuning mode is driving the chassis; make sure the modules don't keep spinning at
      // whatever velocity was last commanded while drive tuning was on.
      for (size_t i = 0; i < NumModules; ++i) {
        states[i] = wpi::math::SwerveModuleVelocity{wpi::units::meters_per_second_t{0},
                                                    modules[i]->GetState().angle};
      }
      drive->SetSwerveModuleStates(states);
    }
    for (auto mode : modes) {
      switch (mode) {
        case TuningMode::AutoAlign:
          // Driven every loop, above.
          [[fallthrough]];
        case TuningMode::DriveInPlace:
          // Sent with drive tuning's setpoint.
          break;
        case TuningMode::Drive:
          if (!modes.contains(TuningMode::DriveVoltage)) {
            for (size_t i = 0; i < NumModules; ++i) {
              states[i] = wpi::math::SwerveModuleVelocity{
                  wpi::units::meters_per_second_t{
                      doubleFields.at(DoubleTelemetryField::ModulesDriveVelocity).Get()},
                  driveAngles[i]};
            }
            drive->SetSwerveModuleStates(states);
          }
          break;
        case TuningMode::DriveVoltage:
          for (size_t i = 0; i < NumModules; ++i) {
            // A closed loop run on the SystemCore would replace the voltage on its next loop.
            if (!lastModes.contains(TuningMode::DriveVoltage)) {
              modules[i]->GetDriveMotorController()->StopClosedLoopController();
            }
            // The wheels are not optimized, so they always drive the same way.
            modules[i]->GetDriveMotorController()->SetVoltage(
                wpi::units::volt_t{doubleFields.at(DoubleTelemetryField::ModulesDriveVoltage).Get()});
            modules[i]->GetAzimuthMotorController()->SetPosition(
                wpi::units::degree_t{driveAngles[i].Degrees()});
          }
          break;
        case TuningMode::Azimuth:
          if (!modes.contains(TuningMode::AzimuthVoltage)) {
            // Zero speed states would be held where the wheels are by the velocity deadband, so turn
            // the wheels without driving them.
            drive->SetAzimuthAngles(wpi::math::Rotation2d{wpi::units::degree_t{
                doubleFields.at(DoubleTelemetryField::ModulesAzimuthAngle).Get()}});
          }
          break;
        case TuningMode::AzimuthVoltage:
          for (auto* module : modules) {
            // A closed loop run on the SystemCore would replace the voltage on its next loop.
            if (!lastModes.contains(TuningMode::AzimuthVoltage)) {
              module->GetAzimuthMotorController()->StopClosedLoopController();
            }
            // Hold the drive motors still while the azimuth motors turn.
            module->GetDriveMotorController()->SetVelocity(wpi::units::meters_per_second_t{0} +
                                                           module->GetCouplingVelocity());
            module->GetAzimuthMotorController()->SetVoltage(wpi::units::volt_t{
                doubleFields.at(DoubleTelemetryField::ModulesAzimuthVoltage).Get()});
          }
          break;
      }
    }
  }
 
  // Set every gain and setpoint that changed.
  for (auto field : {DoubleTelemetryField::TranslationP,
                     DoubleTelemetryField::TranslationI,
                     DoubleTelemetryField::TranslationD,
                     DoubleTelemetryField::RotationP,
                     DoubleTelemetryField::RotationI,
                     DoubleTelemetryField::RotationD,
                     DoubleTelemetryField::AutoAlignPoseX,
                     DoubleTelemetryField::AutoAlignPoseY,
                     DoubleTelemetryField::AutoAlignPoseRotation,
                     DoubleTelemetryField::ModulesDriveP,
                     DoubleTelemetryField::ModulesDriveI,
                     DoubleTelemetryField::ModulesDriveD,
                     DoubleTelemetryField::ModulesDriveKs,
                     DoubleTelemetryField::ModulesDriveKv,
                     DoubleTelemetryField::ModulesDriveKa,
                     DoubleTelemetryField::ModulesDriveVelocity,
                     DoubleTelemetryField::ModulesDriveVoltage,
                     DoubleTelemetryField::ModulesAzimuthP,
                     DoubleTelemetryField::ModulesAzimuthI,
                     DoubleTelemetryField::ModulesAzimuthD,
                     DoubleTelemetryField::ModulesAzimuthKs,
                     DoubleTelemetryField::ModulesAzimuthKv,
                     DoubleTelemetryField::ModulesAzimuthKa,
                     DoubleTelemetryField::ModulesAzimuthAngle,
                     DoubleTelemetryField::ModulesAzimuthVoltage}) {
    auto& dt = doubleFields.at(field);
    if (!dt.IsEnabled() || !dt.IsTunable()) continue;
    switch (field) {
      case DoubleTelemetryField::TranslationP:
        if (translationPID) {
          translationPID->SetP(dt.Get());
          drive->SetTranslationPID(*translationPID);
        }
        break;
      case DoubleTelemetryField::TranslationI:
        if (translationPID) {
          translationPID->SetI(dt.Get());
          drive->SetTranslationPID(*translationPID);
        }
        break;
      case DoubleTelemetryField::TranslationD:
        if (translationPID) {
          translationPID->SetD(dt.Get());
          drive->SetTranslationPID(*translationPID);
        }
        break;
      case DoubleTelemetryField::RotationP:
        if (rotationPID) {
          rotationPID->SetP(dt.Get());
          drive->SetRotationPID(*rotationPID);
        }
        break;
      case DoubleTelemetryField::RotationI:
        if (rotationPID) {
          rotationPID->SetI(dt.Get());
          drive->SetRotationPID(*rotationPID);
        }
        break;
      case DoubleTelemetryField::RotationD:
        if (rotationPID) {
          rotationPID->SetD(dt.Get());
          drive->SetRotationPID(*rotationPID);
        }
        break;
      case DoubleTelemetryField::AutoAlignPoseX:
      case DoubleTelemetryField::AutoAlignPoseY:
      case DoubleTelemetryField::AutoAlignPoseRotation:
        // Auto-align reads its target pose every loop, above.
        break;
      case DoubleTelemetryField::ModulesDriveP:
        for (auto* module : modules) module->GetDriveMotorController()->SetKp(dt.Get());
        break;
      case DoubleTelemetryField::ModulesDriveI:
        for (auto* module : modules) module->GetDriveMotorController()->SetKi(dt.Get());
        break;
      case DoubleTelemetryField::ModulesDriveD:
        for (auto* module : modules) module->GetDriveMotorController()->SetKd(dt.Get());
        break;
      case DoubleTelemetryField::ModulesDriveKs:
        for (auto* module : modules) module->GetDriveMotorController()->SetKs(dt.Get());
        break;
      case DoubleTelemetryField::ModulesDriveKv:
        for (auto* module : modules) module->GetDriveMotorController()->SetKv(dt.Get());
        break;
      case DoubleTelemetryField::ModulesDriveKa:
        for (auto* module : modules) module->GetDriveMotorController()->SetKa(dt.Get());
        break;
      case DoubleTelemetryField::ModulesAzimuthP:
        for (auto* module : modules) module->GetAzimuthMotorController()->SetKp(dt.Get());
        break;
      case DoubleTelemetryField::ModulesAzimuthI:
        for (auto* module : modules) module->GetAzimuthMotorController()->SetKi(dt.Get());
        break;
      case DoubleTelemetryField::ModulesAzimuthD:
        for (auto* module : modules) module->GetAzimuthMotorController()->SetKd(dt.Get());
        break;
      case DoubleTelemetryField::ModulesAzimuthKs:
        for (auto* module : modules) module->GetAzimuthMotorController()->SetKs(dt.Get());
        break;
      case DoubleTelemetryField::ModulesAzimuthKv:
        for (auto* module : modules) module->GetAzimuthMotorController()->SetKv(dt.Get());
        break;
      case DoubleTelemetryField::ModulesAzimuthKa:
        for (auto* module : modules) module->GetAzimuthMotorController()->SetKa(dt.Get());
        break;
      case DoubleTelemetryField::ModulesDriveVelocity:
        if (modes.contains(TuningMode::Drive) && !modes.contains(TuningMode::DriveVoltage)) {
          wpi::util::array<wpi::math::SwerveModuleVelocity, NumModules> states{
              wpi::util::empty_array};
          for (size_t i = 0; i < NumModules; ++i) {
            states[i] = wpi::math::SwerveModuleVelocity{wpi::units::meters_per_second_t{dt.Get()},
                                                        driveAngles[i]};
          }
          drive->SetSwerveModuleStates(states);
        }
        break;
      case DoubleTelemetryField::ModulesDriveVoltage:
        if (modes.contains(TuningMode::DriveVoltage)) {
          for (size_t i = 0; i < NumModules; ++i) {
            // The wheels are not optimized, so they always drive the same way.
            modules[i]->GetDriveMotorController()->SetVoltage(wpi::units::volt_t{dt.Get()});
            modules[i]->GetAzimuthMotorController()->SetPosition(
                wpi::units::degree_t{driveAngles[i].Degrees()});
          }
        }
        break;
      case DoubleTelemetryField::ModulesAzimuthAngle:
        if (modes.contains(TuningMode::Azimuth) && !modes.contains(TuningMode::AzimuthVoltage)) {
          // Zero speed states would be held where the wheels are by the velocity deadband, so turn
          // the wheels without driving them.
          drive->SetAzimuthAngles(wpi::math::Rotation2d{wpi::units::degree_t{dt.Get()}});
        }
        break;
      case DoubleTelemetryField::ModulesAzimuthVoltage:
        if (modes.contains(TuningMode::AzimuthVoltage)) {
          for (auto* module : modules) {
            // Hold the drive motors still while the azimuth motors turn.
            module->GetDriveMotorController()->SetVelocity(wpi::units::meters_per_second_t{0} +
                                                           module->GetCouplingVelocity());
            module->GetAzimuthMotorController()->SetVoltage(wpi::units::volt_t{dt.Get()});
          }
        }
        break;
      default:
        break;
    }
  }
}
 
template <std::size_t NumModules>
void SwerveDriveTelemetry::StopTuning(mechanisms::swerve::SwerveDrive<NumModules>* drive) {
  // Turn every tuning mode off and apply it, so a voltage stops and the wheels are held where they
  // are.
  for (auto& [field, bt] : m_config.GetBoolFields()) {
    if (bt.IsEnabled()) bt.ForceSet(false);
  }
  ApplyTuningValues(drive);
  // Send everything again when tuning starts again.
  m_lastModes.reset();
}
 
}  // namespace yams::telemetry