BaseController.h

#pragma once
 
#include <memory>
#include <string>
 
#include <ctre/phoenix6/CANBus.hpp>
#include <ctre/phoenix6/CANcoder.hpp>
#include <ctre/phoenix6/TalonFX.hpp>
#include <frc/simulation/DCMotorSim.h>
#include <frc/simulation/FlywheelSim.h>
#include <frc/simulation/SingleJointedArmSim.h>
#include <frc/smartdashboard/SmartDashboard.h>
#include <frc/system/plant/DCMotor.h>
#include <networktables/NetworkTable.h>
#include <networktables/NetworkTableEntry.h>
#include <networktables/NetworkTableInstance.h>
#include <units/angle.h>
#include <units/angular_velocity.h>
#include <units/base.h>
#include <units/current.h>
#include <units/velocity.h>
#include <units/voltage.h>
#include <valkyrie/Loggable.h>
#include <valkyrie/controllers/PIDF.h>
 
#include "frc/smartdashboard/SendableChooser.h"
#include "wpi/sendable/SendableRegistry.h"
 
namespace valor {

enum NeutralMode { Brake, Coast };

class BaseController : public valor::Loggable {
   public:
    using Output = std::variant<units::volt_t, units::ampere_t, units::scalar_t>;
 
    template <typename T>
    class Signal {
       public:
        constexpr Signal(T value, units::second_t timestamp) : value{value}, timestamp{timestamp} {}
 
        inline const T* operator->() const { return &value; }
 
        inline const T& operator*() const { return value; }
 
        constexpr units::second_t GetTimestamp() const { return timestamp; }
 
       private:
        T value;
        units::second_t timestamp;
    };
 
    enum class OperationalMode {
        NON_OPERATIONAL,
        TUNING,
        OPERATIONAL,
    };
 
    // Should match the order in BaseController::Output
    enum class ControlMode {
        VOLTAGE,
        CURRENT,
        DUTY_CYCLE,
        SPEED,
        POSITION,
    };

    explicit BaseController(frc::DCMotor _motorSpec)
          : motorSpec{_motorSpec}, rotorToSensor(1), sensorToMech(1), voltageCompensation(units::volt_t{12.0}) {
        tuningModeChooser.SetDefaultOption(magic_enum::enum_name(ControlMode::VOLTAGE), ControlMode::VOLTAGE);
        for (ControlMode mode : magic_enum::enum_values<ControlMode>())
            tuningModeChooser.AddOption(magic_enum::enum_name(mode), mode);
 
        opModeChooser.SetDefaultOption(magic_enum::enum_name(OperationalMode::OPERATIONAL), OperationalMode::OPERATIONAL);
        for (OperationalMode mode : magic_enum::enum_values<OperationalMode>())
            opModeChooser.AddOption(magic_enum::enum_name(mode), mode);
 
        LogChild("Tuning Mode", &tuningModeChooser);
        LogChild("Operational Mode", &opModeChooser);
    }
 
    void OnLoggingStart() override { WriteLog("Setpoint", 0.0); }
 
    virtual void ApplyConfig() = 0;

    constexpr units::turns_per_second_t GetMaxMechanismSpeed() { return motorSpec.freeSpeed / GetGearRatio(); }
 
    constexpr double GetGearRatio() const { return rotorToSensor * sensorToMech; }

    constexpr void SetVoltageCompensation(units::volt_t _voltageCompensation) { voltageCompensation = _voltageCompensation; }
 
    constexpr units::volt_t GetVoltageCompensation() const { return voltageCompensation; }

    virtual constexpr void Reset() {}

    virtual Signal<units::turn_t> GetPosition() const = 0;

    virtual Signal<units::ampere_t> GetCurrent() const = 0;
 
    virtual Signal<units::volt_t> GetVoltage() const = 0;

    virtual Signal<units::turns_per_second_t> GetSpeed() const = 0;
 
    virtual Signal<units::turns_per_second_squared_t> GetAcceleration() const = 0;
 
    virtual void SetEncoderPosition(units::turn_t position) = 0;
 
    void SetPosition(units::turn_t position, Output feedforward = 0, int slot = 0) {
        if (opModeChooser.GetSelected() == OperationalMode::OPERATIONAL) {
            setpoint = position.value();
            controlMode = ControlMode::POSITION;
            PositionSetter(position, feedforward, slot);
        } else if (opModeChooser.GetSelected() == OperationalMode::TUNING)
            HandleTuningMode();
        else
            PowerSetter(0_V);
    }
 
    void SetSpeed(units::turns_per_second_t speed, Output feedforward = 0, int slot = 0) {
        if (opModeChooser.GetSelected() == OperationalMode::OPERATIONAL) {
            setpoint = speed.value();
            controlMode = ControlMode::SPEED;
            SpeedSetter(speed, feedforward, slot);
        } else if (opModeChooser.GetSelected() == OperationalMode::TUNING)
            HandleTuningMode();
        else
            PowerSetter(0_V);
    }
 
    void SetPower(Output power) {
        if (opModeChooser.GetSelected() == OperationalMode::OPERATIONAL) {
            std::visit([this](auto&& power) { setpoint = power.value(); }, power);
            controlMode = static_cast<ControlMode>(power.index());
            PowerSetter(power);
        } else if (opModeChooser.GetSelected() == OperationalMode::TUNING)
            HandleTuningMode();
        else
            PowerSetter(0_V);
    }

    virtual void SetupFollower(int canID, bool followerInverted = false) = 0;

    virtual void SetPIDF(valor::PIDF<> pidf, int slot = 0, bool saveImmediately = false) = 0;

    void SetLimits(units::turn_t reverse, units::turn_t forward, bool saveImmediately = false) {
        SetForwardLimit(forward, saveImmediately);
        SetReverseLimit(reverse, saveImmediately);
    }

    virtual void SetForwardLimit(units::turn_t forward, bool saveImmediately = false) = 0;

    virtual void SetReverseLimit(units::turn_t reverse, bool saveImmediately = false) = 0;

    virtual void SetNeutralMode(valor::NeutralMode mode, bool saveImmediately = false) = 0;

    virtual void SetGearRatios(double rotorToSensor, double sensorToMech, bool saveImmediately = false) = 0;
 
    virtual void SetOpenLoopRamp(units::second_t time, bool saveImmediately = false) = 0;
 
    virtual units::turn_t GetAbsoluteEncoderPosition() const = 0;
 
    virtual units::volt_t CalculateAppliedVoltage() = 0;
    virtual void SetSimState(units::turn_t pos, units::turns_per_second_t vel, units::turns_per_second_squared_t accel) = 0;
 
    virtual inline void SetSimState(frc::sim::DCMotorSim& sim) {
        SetSimState(sim.GetAngularPosition(), sim.GetAngularVelocity(), sim.GetAngularAcceleration());
    }
 
    virtual inline void SetSimState(frc::sim::FlywheelSim& sim) {
        SetSimState(0_tr, sim.GetAngularVelocity(), sim.GetAngularAcceleration());
    }
 
    virtual inline void SetSimState(frc::sim::SingleJointedArmSim& sim) { SetSimState(sim.GetAngle(), sim.GetVelocity(), 0_tr_per_s_sq); }
 
    void LoggablePeriodic() override {
        WriteLog("Control Mode", controlMode);
        if (opModeChooser.GetSelected() == OperationalMode::TUNING)
            setpoint = ReadLog("Setpoint", setpoint);
        else
            WriteLog("Setpoint", setpoint);
    }
 
    frc::DCMotor motorSpec;
 
   protected:
    virtual void PositionSetter(units::turn_t position, Output feedforward = 0, int slot = 0) = 0;

    virtual void SpeedSetter(units::turns_per_second_t speed, Output feedforward = 0, int slot = 0) = 0;
 
    virtual void PowerSetter(Output out) = 0;
 
    double rotorToSensor;
    double sensorToMech;
    units::volt_t voltageCompensation;
 
   private:
    void HandleTuningMode() {
        controlMode = tuningModeChooser.GetSelected();
        switch (tuningModeChooser.GetSelected()) {
            case ControlMode::VOLTAGE:
                return PowerSetter(units::volt_t{setpoint});
            case ControlMode::CURRENT:
                return PowerSetter(units::ampere_t{setpoint});
            case ControlMode::DUTY_CYCLE:
                return PowerSetter(setpoint);
            case ControlMode::SPEED:
                return SpeedSetter(units::turns_per_second_t{setpoint});
            case ControlMode::POSITION:
                return PositionSetter(units::turn_t{setpoint});
        }
    }
 
    frc::SendableChooser<ControlMode> tuningModeChooser;
    frc::SendableChooser<OperationalMode> opModeChooser;
    ControlMode controlMode = ControlMode::DUTY_CYCLE;
    double setpoint = 0.0;
};
}  // namespace valor