// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.

#include <numbers>

#include "wpi/driverstation/Gamepad.hpp"
#include "wpi/framework/TimedRobot.hpp"
#include "wpi/hardware/motor/PWMSparkMax.hpp"
#include "wpi/hardware/rotation/Encoder.hpp"
#include "wpi/math/controller/LinearQuadraticRegulator.hpp"
#include "wpi/math/estimator/KalmanFilter.hpp"
#include "wpi/math/system/DCMotor.hpp"
#include "wpi/math/system/LinearSystemLoop.hpp"
#include "wpi/math/system/Models.hpp"
#include "wpi/math/trajectory/TrapezoidProfile.hpp"
#include "wpi/units/angle.hpp"
#include "wpi/units/moment_of_inertia.hpp"

/**
 * This is a sample program to demonstrate how to use a state-space controller
 * to control an arm.
 */
class Robot : public wpi::TimedRobot {
  static constexpr int kMotorPort = 0;
  static constexpr int kEncoderAChannel = 0;
  static constexpr int kEncoderBChannel = 1;
  static constexpr int kJoystickPort = 0;

  static constexpr wpi::units::radian_t kRaisedPosition = 90_deg;
  static constexpr wpi::units::radian_t kLoweredPosition = 0_deg;

  // Moment of inertia of the arm. Can be estimated with CAD. If finding this
  // constant is difficult, wpi::math::LinearSystem.identifyPositionSystem may
  // be better.
  static constexpr wpi::units::kilogram_square_meter_t kArmMOI = 1.2_kg_sq_m;

  // Reduction between motors and encoder, as output over input. If the arm
  // spins slower than the motors, this number should be greater than one.
  static constexpr double kArmGearing = 10.0;

  // The plant holds a state-space model of our arm. This system has the
  // following properties:
  //
  // States: [position, velocity], in radians and radians per second.
  // Inputs (what we can "put in"): [voltage], in volts.
  // Outputs (what we can measure): [position], in radians.
  wpi::math::LinearSystem<2, 1, 1> armPlant =
      wpi::math::Models::SingleJointedArmFromPhysicalConstants(
          wpi::math::DCMotor::NEO(2), kArmMOI, kArmGearing)
          .Slice(0);

  // The observer fuses our encoder data and voltage inputs to reject noise.
  wpi::math::KalmanFilter<2, 1, 1> observer{
      armPlant,
      {0.015, 0.17},  // How accurate we think our model is
      {0.01},         // How accurate we think our encoder position
      // data is. In this case we very highly trust our encoder position
      // reading.
      20_ms};

  // A LQR uses feedback to create voltage commands.
  wpi::math::LinearQuadraticRegulator<2, 1> controller{
      armPlant,
      // qelms. Velocity error tolerance, in radians and radians per second.
      // Decrease this to more heavily penalize state excursion, or make the
      // controller behave more aggressively.
      {1.0 * 2.0 * std::numbers::pi / 360.0,
       10.0 * 2.0 * std::numbers::pi / 360.0},
      // relms. Control effort (voltage) tolerance. Decrease this to more
      // heavily penalize control effort, or make the controller less
      // aggressive. 12 is a good starting point because that is the
      // (approximate) maximum voltage of a battery.
      {12.0},
      // Nominal time between loops. 20ms for TimedRobot, but can be lower if
      // using notifiers.
      20_ms};

  // The state-space loop combines a controller, observer, feedforward and plant
  // for easy control.
  wpi::math::LinearSystemLoop<2, 1, 1> loop{armPlant, controller, observer,
                                            12_V, 20_ms};

  // An encoder set up to measure arm position in radians per second.
  wpi::Encoder encoder{kEncoderAChannel, kEncoderBChannel};

  wpi::PWMSparkMax motor{kMotorPort};
  wpi::Gamepad joystick{kJoystickPort};

  wpi::math::TrapezoidProfile<wpi::units::radians> profile{
      {45_deg_per_s, 90_deg_per_s / 1_s}};

  wpi::math::TrapezoidProfile<wpi::units::radians>::State lastProfiledReference;

 public:
  Robot() {
    // We go 2 pi radians per 4096 clicks.
    encoder.SetDistancePerPulse(2.0 * std::numbers::pi / 4096.0);
  }

  void TeleopInit() override {
    loop.Reset(wpi::math::Vectord<2>{encoder.GetDistance(), encoder.GetRate()});

    lastProfiledReference = {
        wpi::units::radian_t{encoder.GetDistance()},
        wpi::units::radians_per_second_t{encoder.GetRate()}};
  }

  void TeleopPeriodic() override {
    // Sets the target position of our arm. This is similar to setting the
    // setpoint of a PID controller.
    wpi::math::TrapezoidProfile<wpi::units::radians>::State goal;
    if (joystick.GetRightBumperButton()) {
      // We pressed the bumper, so let's set our next reference
      goal = {kRaisedPosition, 0_rad_per_s};
    } else {
      // We released the bumper, so let's spin down
      goal = {kLoweredPosition, 0_rad_per_s};
    }
    lastProfiledReference =
        profile.Calculate(20_ms, lastProfiledReference, goal);

    loop.SetNextR(
        wpi::math::Vectord<2>{lastProfiledReference.position.value(),
                              lastProfiledReference.velocity.value()});

    // Correct our Kalman filter's state vector estimate with encoder data.
    loop.Correct(wpi::math::Vectord<1>{encoder.GetDistance()});

    // Update our LQR to generate new voltage commands and use the voltages to
    // predict the next state with out Kalman filter.
    loop.Predict(20_ms);

    // Send the new calculated voltage to the motors.
    // voltage = duty cycle * battery voltage, so
    // duty cycle = voltage / battery voltage
    motor.SetVoltage(wpi::units::volt_t{loop.U(0)});
  }
};

#ifndef RUNNING_WPILIB_TESTS
int main() {
  return wpi::StartRobot<Robot>();
}
#endif
