PhoenixController.h

#pragma once
#ifdef _MSC_VER
#pragma warning(push)
#pragma warning(disable : 4250)  // Disable "inherits via dominance" warning
#else
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Woverloaded-virtual"
#endif
 
#include <ctime>
#include <cwchar>
#include <memory>
#include <string>
#include <unordered_map>
 
#include <ctre/phoenix6/CANBus.hpp>
#include <ctre/phoenix6/CANcoder.hpp>
#include <ctre/phoenix6/TalonFX.hpp>
#include <ctre/phoenix6/configs/MotionMagicConfigs.hpp>
#include <ctre/phoenix6/sim/ChassisReference.hpp>
#include <frc/RobotController.h>
#include <units/angular_acceleration.h>
#include <units/angular_jerk.h>
#include <units/angular_velocity.h>
#include <units/base.h>
#include <units/dimensionless.h>
#include <units/frequency.h>
#include <units/voltage.h>
 
#include "valkyrie/controllers/BaseController.h"
#include "valkyrie/util/PhoenixSignalManager.h"
 
// Conversion guide:
// https://v6.docs.ctr-electronics.com/en/latest/docs/migration/migration-guide/closed-loop-guide.html
 
namespace valor {
 
enum PhoenixControllerType { KRAKEN_X60_FOC, KRAKEN_X60, KRAKEN_X44, KRAKEN_X44_FOC, FALCON_FOC, FALCON };
 
class CorePhoenixController : public BaseController, public ctre::phoenix6::hardware::core::CoreTalonFX {
    static inline const std::unordered_map<PhoenixControllerType, frc::DCMotor> MOTOR_SPEC_MAP{
        {KRAKEN_X60, frc::DCMotor::KrakenX60()}, {KRAKEN_X60_FOC, frc::DCMotor::KrakenX60FOC()},
        {KRAKEN_X44, frc::DCMotor::KrakenX44()}, {KRAKEN_X44_FOC, frc::DCMotor::KrakenX44FOC()},
        {FALCON, frc::DCMotor::Falcon500()},     {FALCON_FOC, frc::DCMotor::Falcon500FOC()}};
 
   public:
    CorePhoenixController(valor::PhoenixControllerType controllerType, int canID, valor::NeutralMode neutralMode, bool inverted,
                          std::string canbus = "");
 
    explicit CorePhoenixController(valor::PhoenixControllerType controllerType)
          : BaseController{MOTOR_SPEC_MAP.at(controllerType)}, ctre::phoenix6::hardware::core::CoreTalonFX{-1} {}
 
    virtual ~CorePhoenixController();
 
    inline void ApplyConfig() override { GetConfigurator().Apply(config, 5_s); }
 
    inline void Reset() override { SetEncoderPosition(0_tr); }
 
    void SetNeutralMode(valor::NeutralMode mode, bool saveImmediately = false) override;
 
    void SetCurrentLimits(units::ampere_t statorCurrentLimit, units::ampere_t supplyCurrentLimit, units::ampere_t supplyCurrentThreshold,
                          units::second_t supplyTimeThreshold, bool saveImmediately = false);
 
    void SetVoltageLimits(units::volt_t reverseLimit, units::volt_t forwardLimit, bool saveImmediately = false);
 
    units::turn_t GetPosition() const override;
 
    inline ctre::phoenix6::StatusSignal<units::turn_t>& GetPositionSignal(bool refresh = true) {
        return ctre::phoenix6::hardware::core::CoreTalonFX::GetPosition(refresh);
    }
 
    units::turns_per_second_t GetSpeed() const override;
 
    units::ampere_t GetCurrent() const override;
 
    inline void SetEncoderPosition(units::turn_t position) override { ctre::phoenix6::hardware::core::CoreTalonFX::SetPosition(position); }
 
    void SetContinuousWrap(bool continuousWrap, bool saveImmediately = false);
 
    void PowerSetter(Output out) override;
 
    void SetupFollower(int canID, bool followerInverted = false) override;
 
    void SetupReverseHardwareLimit(int canID, ctre::phoenix6::signals::ReverseLimitTypeValue type, units::turn_t autosetPosition = 0_tr,
                                   bool saveImmediately = false);
 
    void SetupForwardHardwareLimit(int canID, ctre::phoenix6::signals::ForwardLimitTypeValue type, units::turn_t autosetPosition = 0_tr,
                                   bool saveImmediately = false);
 
    void SetForwardLimit(units::turn_t forward, bool saveImmediately = false) override;
 
    void SetReverseLimit(units::turn_t reverse, bool saveImmediately = false) override;
 
    void SetGearRatios(double _rotorToSensor, double _sensorToMech, bool saveImmediately = false) override;
 
    ctre::phoenix6::signals::MagnetHealthValue GetMagnetHealth();
 
    units::turn_t GetAbsoluteEncoderPosition() const override;
 
    void SetupCANcoder(int deviceId, units::turn_t offset, bool clockwise, std::string canbus = "", units::turn_t absoluteRange = 1_tr,
                       bool saveImmediately = false);
 
    inline ctre::phoenix6::hardware::CANcoder* GetCANcoder() { return cancoder.get(); }
 
    void SetOpenLoopRamp(units::second_t time, bool saveImmediately = false) override;
 
    units::volt_t CalculateAppliedVoltage() override;
 
    using BaseController::SetSimState;
    void SetSimState(units::turn_t pos, units::turns_per_second_t vel, units::turns_per_second_squared_t accel);
 
    void LoggablePeriodic() override;
 
   protected:
    static constexpr units::ampere_t STATOR_CURRENT_LIMIT = 40_A;
    static constexpr units::ampere_t SUPPLY_CURRENT_LIMIT = 30_A;
    static constexpr units::ampere_t SUPPLY_CURRENT_THRESHOLD = 30_A;
    static constexpr units::millisecond_t SUPPLY_TIME_THRESHOLD = 1_s;
    static constexpr units::ampere_t PEAK_TORQUE_CURRENT = 100_A;
    static constexpr units::turn_t DEADBAND = 0.01_tr;
 
    std::unique_ptr<CorePhoenixController> followerMotor;
 
    std::unique_ptr<ctre::phoenix6::hardware::CANcoder> cancoder;
    std::unique_ptr<ctre::phoenix6::configs::MagnetSensorConfigs> cancoderConfig;
 
    ctre::phoenix6::configs::TalonFXConfiguration config;
 
    ctre::phoenix6::controls::VoltageOut voltageRequest{0_V};
    ctre::phoenix6::controls::DutyCycleOut dutyCycleRequest{0};
    ctre::phoenix6::controls::TorqueCurrentFOC currentRequest{0_A};
};
 
template <class VelocityOutput = ctre::phoenix6::controls::VelocityTorqueCurrentFOC,
          class PositionOutput = ctre::phoenix6::controls::DynamicMotionMagicTorqueCurrentFOC>
class PhoenixController : public CorePhoenixController {
    static_assert(std::is_base_of_v<ctre::phoenix6::controls::ControlRequest, VelocityOutput>);
    static_assert(std::is_base_of_v<ctre::phoenix6::controls::ControlRequest, PositionOutput>);
 
   public:
    PhoenixController(valor::PhoenixControllerType controllerType, int canID, valor::NeutralMode neutralMode, bool inverted,
                      std::string canbus = "")
          : CorePhoenixController{controllerType, canID, neutralMode, inverted, canbus} {
        positionRequest.UpdateFreqHz = 0_Hz;
        velocityRequest.UpdateFreqHz = 0_Hz;
    };
 
    explicit PhoenixController(PhoenixControllerType controllerType = KRAKEN_X60_FOC) : CorePhoenixController{controllerType} {}
 
    inline void EnableFOC(bool enable) {
        TryEnableFOC(voltageRequest, enable);
        TryEnableFOC(dutyCycleRequest, enable);
        TryEnableFOC(currentRequest, enable);
        TryEnableFOC(velocityRequest, enable);
        TryEnableFOC(positionRequest, enable);
    }
 
    void SetNeutralMode(valor::NeutralMode mode, bool saveImmediately = false) override {
        if constexpr (requires { positionRequest.OverrideBrakeDurNeutral; })
            positionRequest.OverrideBrakeDurNeutral = mode == valor::NeutralMode::Brake;
        else
            positionRequest.OverrideCoastDurNeutral = mode == valor::NeutralMode::Coast;
 
        if constexpr (requires { velocityRequest.OverrideBrakeDurNeutral; })
            velocityRequest.OverrideBrakeDurNeutral = mode == valor::NeutralMode::Brake;
        else
            velocityRequest.OverrideCoastDurNeutral = mode == valor::NeutralMode::Coast;
 
        CorePhoenixController::SetNeutralMode(mode, saveImmediately);
    }
 
    void PositionSetter(units::turn_t position, Output feedforward, int slot) override {
        std::visit(
            [=, this](auto&& ff) {
                using FFType = decltype(positionRequest.FeedForward);
                if constexpr (std::is_same_v<std::decay_t<decltype(ff)>, FFType>)
                    SetControl(positionRequest.WithSlot(slot).WithPosition(position).WithFeedForward(ff));
                else {
                    if (ff.value())
                        FRC_ReportWarning("Position control feedforward is of the wrong unit, using zero instead...");
                    SetControl(positionRequest.WithSlot(slot).WithPosition(position).WithFeedForward(FFType{0}));
                }
            },
            feedforward);
    }
 
    using BaseController::SetPosition;
 
    void SetPosition(units::turn_t position, units::turns_per_second_t maxVelocity, units::turns_per_second_squared_t maxAcceleration,
                     units::turns_per_second_cubed_t maxJerk, Output feedforward = 0, int slot = 0) {
        positionRequest.Velocity = maxVelocity;
        positionRequest.Acceleration = maxAcceleration;
        positionRequest.Jerk = maxJerk;
        SetPosition(position, feedforward, slot);
    }
 
    void SpeedSetter(units::turns_per_second_t velocity, Output feedforward, int slot) override {
        std::visit(
            [=, this](auto&& ff) {
                using FFType = decltype(velocityRequest.FeedForward);
                if constexpr (std::is_same_v<std::decay_t<decltype(ff)>, FFType>)
                    SetControl(velocityRequest.WithSlot(slot).WithVelocity(velocity).WithFeedForward(ff));
                else {
                    if (ff.value())
                        FRC_ReportWarning("Velocity control feedforward is of the wrong unit, using zero instead...");
                    SetControl(velocityRequest.WithSlot(slot).WithVelocity(velocity).WithFeedForward(FFType{0}));
                }
            },
            feedforward);
    }
 
    void UseTimesync(units::hertz_t controlFreq, bool saveImmediately = false) {
        config.MotorOutput.ControlTimesyncFreqHz = controlFreq;
        positionRequest.UseTimesync = true;
        velocityRequest.UseTimesync = true;
 
        if (saveImmediately)
            GetConfigurator().Apply(config.MotorOutput);
    }
 
    void SetPIDF(valor::PIDF<> pidf, int slot = 0, bool saveImmediately = false) override {
        if constexpr (requires { positionRequest.Jerk; }) {
            positionRequest.Jerk = decltype(positionRequest.Jerk)(pidf.maxJerk.value());
            positionRequest.Acceleration = decltype(positionRequest.Acceleration)(pidf.maxAcceleration.value());
            positionRequest.Velocity = decltype(positionRequest.Velocity)(pidf.maxVelocity.value());
        }
 
        config.MotionMagic.MotionMagicJerk = decltype(config.MotionMagic.MotionMagicJerk)(pidf.maxJerk.value());
        config.MotionMagic.MotionMagicAcceleration = decltype(config.MotionMagic.MotionMagicAcceleration)(pidf.maxAcceleration.value());
        config.MotionMagic.MotionMagicCruiseVelocity = decltype(config.MotionMagic.MotionMagicCruiseVelocity)(pidf.maxVelocity.value());
 
        if (saveImmediately) {
            GetConfigurator().Apply(config.MotionMagic);
        }
 
        if (slot == 1)
            SetPIDFSlot(config.Slot1, pidf, saveImmediately);
        else if (slot == 2)
            SetPIDFSlot(config.Slot2, pidf, saveImmediately);
        else
            SetPIDFSlot(config.Slot0, pidf, saveImmediately);  // Default case
    }
 
   private:
    inline PositionOutput CreatePositionRequest() {
        if constexpr (requires { PositionOutput{0_tr}; })
            return {0_tr};
        else
            return {0_tr, 0_tps, 0_tr_per_s_sq};
    }
 
    // Helper function for enableFOC that only enables if the template has the field "EnableFOC"
    template <typename T>
    static constexpr void TryEnableFOC(T& request, bool enable) {
        if constexpr (requires { request.EnableFOC; })
            request.EnableFOC = enable;
    }
 
    template <class S>
    void SetPIDFSlot(S& slot, valor::PIDF<> pidf, bool saveImmediately) {
        // Generic PIDF configurations
        // Numerator for closed loop controls will be in volts
        slot.kP = pidf.P.value();
        slot.kI = pidf.I.value();
        slot.kD = pidf.D.value();
        slot.kS = pidf.kS.value();
        if (pidf.kV)
            slot.kV = (*pidf.kV).value();
        else if constexpr (std::is_same_v<decltype(positionRequest.FeedForward), units::ampere_t> ||
                           std::is_same_v<decltype(velocityRequest.FeedForward), units::ampere_t>)
            slot.kV = 0;
        else
            slot.kV = (voltageCompensation / GetMaxMechanismSpeed()).value();
 
        // Feedforward gain configuration
        slot.GravityArmPositionOffset = pidf.GTarget;
        slot.GravityType = pidf.GType == valor::GravityFFType::LINEAR ? ctre::phoenix6::signals::GravityTypeValue::Elevator_Static
                                                                      : ctre::phoenix6::signals::GravityTypeValue::Arm_Cosine;
        slot.kG = pidf.kG.value();
 
        if (saveImmediately)
            GetConfigurator().Apply(slot);
    }
 
    VelocityOutput velocityRequest{0_tps};
    PositionOutput positionRequest = CreatePositionRequest();
};
 
}  // namespace valor
 
#ifdef _MSC_VER
#pragma warning(pop)
#else
#pragma GCC diagnostic pop
#endif