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 "frc/smartdashboard/SendableChooser.h"
#include "valkyrie/Loggable.h"
#include "valkyrie/controllers/PIDF.h"
#include "valkyrie/util/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>;

    explicit BaseController(frc::DCMotor _motorSpec)
          : motorSpec{_motorSpec}, rotorToSensor(1), sensorToMech(1), voltageCompensation(units::volt_t{12.0}) {
        for (TuningMode mode : magic_enum::enum_values<TuningMode>())
            tuningModeChooser.AddOption(magic_enum::enum_name(mode), mode);
        tuningModeChooser.SetDefaultOption(magic_enum::enum_name(TuningMode::Voltage));
        opModeChooser.AddOption("Tuning", OperationalMode::Tuning);
        opModeChooser.AddOption("Non Operational", OperationalMode::NonOperational);
        opModeChooser.SetDefaultOption("Operational", OperationalMode::Operational);
 
        LogChild("Tuning Mode", &tuningModeChooser);
        LogChild("Operational Mode", &opModeChooser);
    }
 
    virtual ~BaseController() {
        wpi::SendableRegistry::Remove(&tuningModeChooser);
        wpi::SendableRegistry::Remove(&opModeChooser);
    }
 
    void OnLoggingStart() override {
        WriteLog("Tuning Setpoint", 0.0);
        WriteLog("Tune", false);
    }
 
    virtual void ApplyConfig() = 0;
 
    enum class OperationalMode { NonOperational, Tuning, Operational };
    enum class TuningMode { Voltage, Current, DutyCycle, Speed, Position };
 
    inline OperationalMode GetOperationalMode() const { return opModeChooser.GetSelected(); }
 
    inline TuningMode GetTuningMode() const { return tuningModeChooser.GetSelected(); }

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

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

    virtual constexpr void Reset() {}

    virtual units::turn_t GetPosition() const = 0;

    virtual units::ampere_t GetCurrent() const = 0;

    virtual units::turns_per_second_t GetSpeed() const = 0;
 
    virtual void SetEncoderPosition(units::turn_t position) = 0;
 
    void SetPosition(units::turn_t position, Output feedforward = 0, int slot = 0) {
        if (GetOperationalMode() == OperationalMode::Operational) {
            WriteLog("Requested Position", position);
            PositionSetter(position, feedforward, slot);
        } else if (GetOperationalMode() == OperationalMode::Tuning && ReadLog<bool>("Tune"))
            HandleTuningMode();
        else
            PowerSetter(0_V);
    }
 
    void SetSpeed(units::turns_per_second_t speed, Output feedforward = 0, int slot = 0) {
        if (GetOperationalMode() == OperationalMode::Operational) {
            WriteLog("Requested Speed", speed);
            SpeedSetter(speed, feedforward, slot);
        } else if (GetOperationalMode() == OperationalMode::Tuning && ReadLog<bool>("Tune"))
            HandleTuningMode();
        else
            PowerSetter(0_V);
    }
 
    void SetPower(Output power) {
        if (GetOperationalMode() == OperationalMode::Operational)
            PowerSetter(power);
        else if (GetOperationalMode() == OperationalMode::Tuning && ReadLog<bool>("Tune"))
            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); }
 
    constexpr void SetDefaultTuningMode(TuningMode mode) { tuningModeChooser.SetDefaultOption(magic_enum::enum_name(mode)); };
 
    frc::DCMotor motorSpec;
    double rotorToSensor;
    double sensorToMech;
 
   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;
 
    units::volt_t voltageCompensation;
 
   private:
    void HandleTuningMode() {
        double setpoint = ReadLog<double>("Tuning Setpoint");
        switch (tuningModeChooser.GetSelected()) {
            case TuningMode::Voltage:
                return PowerSetter(units::volt_t{setpoint});
            case TuningMode::Current:
                return PowerSetter(units::ampere_t{setpoint});
            case TuningMode::DutyCycle:
                return PowerSetter(setpoint);
            case TuningMode::Speed:
                WriteLog("Requested Speed", units::turns_per_second_t{setpoint});
                return SpeedSetter(units::turns_per_second_t{setpoint});
            case TuningMode::Position:
                WriteLog("Requested Position", units::turn_t{setpoint});
                return PositionSetter(units::turn_t{setpoint});
        }
    }
 
    valor::SendableChooser<TuningMode> tuningModeChooser;
    valor::SendableChooser<OperationalMode> opModeChooser;
};
}  // namespace valor