ProfiledPIDController.h

#pragma once
#include <frc/controller/ProfiledPIDController.h>
#include <valkyrie/controllers/PIDController.h>
 
namespace valor {

template <typename INPUT, typename OUTPUT>
class ProfiledPIDController {
   public:
    using ErrorDerivative = units::unit_t<units::compound_unit<INPUT, units::inverse<units::second>>>;

    constexpr ProfiledPIDController(const PIDF<INPUT, OUTPUT>& pidf, units::second_t period = 20_ms)
          : controller{pidf}, profile{{pidf.maxVelocity, pidf.maxAcceleration}} {}

    constexpr void SetPID(const PIDF<INPUT, OUTPUT>& pidf) {
        controller.SetPID(pidf);
        profile = frc::TrapezoidProfile<INPUT>({pidf.maxVelocity, pidf.maxAcceleration});
    }

    constexpr void SetIZone(PIDF<INPUT, OUTPUT>::Output_t iZone) { controller.SetIZone(iZone); }

    PIDF<INPUT, OUTPUT> GetPID() const { return controller.GetPID(); }

    constexpr PIDF<INPUT, OUTPUT>::Output_t GetIZone() const { return controller.GetIZone(); }

    constexpr units::second_t GetPeriod() const { return controller.GetPeriod(); }

    constexpr PIDF<INPUT, OUTPUT>::Input_t GetPositionTolerance() const { return controller.GetErrorTolerance(); }

    constexpr ErrorDerivative GetVelocityTolerance() const { return controller.GetErrorDerivativeTolerance(); }

    constexpr PIDF<INPUT, OUTPUT>::Input_t GetAccumulatedError() const { return controller.GetAccumulatedError(); }

    constexpr void SetGoal(frc::TrapezoidProfile<INPUT>::State goal) { this->goal = goal; }

    constexpr frc::TrapezoidProfile<INPUT>::State GetGoal() const { return goal; }

    constexpr bool AtGoal() const { return AtSetpoint() && goal == setpoint; }

    constexpr frc::TrapezoidProfile<INPUT>::State GetSetpoint() const { return setpoint; }

    constexpr bool AtSetpoint() const { return controller.AtSetpoint(); }

    constexpr void EnableContinuousInput(PIDF<INPUT, OUTPUT>::Input_t minimumInput, PIDF<INPUT, OUTPUT>::Input_t maximumInput) {
        controller.EnableContinuousInput(minimumInput, maximumInput);
        this->minimumInput = minimumInput;
        this->maximumInput = maximumInput;
    }

    constexpr void DisableContinuousInput() { controller.DisableContinuousInput(); }

    constexpr void SetIntegratorRange(PIDF<INPUT, OUTPUT>::Output_t minimumIntegral, PIDF<INPUT, OUTPUT>::Output_t maximumIntegral) {
        controller.SetIntegratorRange(minimumIntegral, maximumIntegral);
    }

    constexpr void SetTolerance(PIDF<INPUT, OUTPUT>::Input_t positionTolerance,
                                ErrorDerivative velocityTolerance = ErrorDerivative{std::numeric_limits<double>::infinity()}) {
        controller.SetTolerance(positionTolerance, velocityTolerance);
    }

    constexpr PIDF<INPUT, OUTPUT>::Input_t GetPositionError() const { return controller.GetError(); }

    constexpr ErrorDerivative GetVelocityError() const { return controller.GetErrorDerivative(); }

    constexpr PIDF<INPUT, OUTPUT>::Output_t Calculate(PIDF<INPUT, OUTPUT>::Input_t measurement) {
        if (controller.IsContinuousInputEnabled()) {
            auto errorBound = (maximumInput - minimumInput) / 2.0;
            auto goalMinDistance = frc::InputModulus(goal.position - measurement, -errorBound, errorBound);
            auto setpointMinDistance = frc::InputModulus(setpoint.position - measurement, -errorBound, errorBound);
            goal.position = goalMinDistance + measurement;
            setpoint.position = setpointMinDistance + measurement;
        }
 
        setpoint = profile.Calculate(GetPeriod(), setpoint, goal);
        controller.SetSetpoint(setpoint.position);
        auto output = controller.Calculate(measurement);
        if constexpr (std::is_same_v<decltype(setpoint.velocity), typename PIDF<INPUT, OUTPUT>::Velocity_t>)
            if (controller.GetPID().kV)
                output += *controller.GetPID().kV * setpoint.velocity;
        return output;
    }

    constexpr void Reset(const frc::TrapezoidProfile<INPUT>::State& measurement) {
        controller.Reset();
        setpoint = measurement;
    }
 
   private:
    PIDController<INPUT, OUTPUT> controller;
    frc::TrapezoidProfile<INPUT>::State goal;
    frc::TrapezoidProfile<INPUT>::State setpoint;
    frc::TrapezoidProfile<INPUT> profile;
 
    PIDF<INPUT, OUTPUT>::Input_t minimumInput;
    PIDF<INPUT, OUTPUT>::Input_t maximumInput;
};
 
}  // namespace valor