// 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.

package org.wpilib.snippets.flywheelbangbangcontroller;

import org.wpilib.driverstation.Joystick;
import org.wpilib.framework.TimedRobot;
import org.wpilib.hardware.motor.PWMSparkMax;
import org.wpilib.hardware.rotation.Encoder;
import org.wpilib.math.controller.BangBangController;
import org.wpilib.math.controller.SimpleMotorFeedforward;
import org.wpilib.math.numbers.N1;
import org.wpilib.math.system.DCMotor;
import org.wpilib.math.system.LinearSystem;
import org.wpilib.math.system.Models;
import org.wpilib.math.util.Units;
import org.wpilib.simulation.EncoderSim;
import org.wpilib.simulation.FlywheelSim;
import org.wpilib.smartdashboard.SmartDashboard;
import org.wpilib.system.RobotController;

/**
 * This is a sample program to demonstrate the use of a BangBangController with a flywheel to
 * control RPM.
 */
public class Robot extends TimedRobot {
  private static final int kMotorPort = 0;
  private static final int kEncoderAChannel = 0;
  private static final int kEncoderBChannel = 1;

  // Max setpoint for joystick control in RPM
  private static final double kMaxSetpointValue = 6000.0;

  // Joystick to control setpoint
  private final Joystick joystick = new Joystick(0);

  private final PWMSparkMax flywheelMotor = new PWMSparkMax(kMotorPort);
  private final Encoder encoder = new Encoder(kEncoderAChannel, kEncoderBChannel);

  private final BangBangController bangBangController = new BangBangController();

  // Gains are for example purposes only - must be determined for your own robot!
  public static final double kFlywheelKs = 0.0001; // V
  public static final double kFlywheelKv = 0.000195; // V/RPM
  public static final double kFlywheelKa = 0.0003; // V/(RPM/s)
  private final SimpleMotorFeedforward feedforward =
      new SimpleMotorFeedforward(kFlywheelKs, kFlywheelKv, kFlywheelKa);

  // Simulation classes help us simulate our robot

  // Reduction between motors and encoder, as output over input. If the flywheel
  // spins slower than the motors, this number should be greater than one.
  private static final double kFlywheelGearing = 1.0;

  // 1/2 MR²
  private static final double kFlywheelMomentOfInertia =
      0.5 * Units.lbsToKilograms(1.5) * Math.pow(Units.inchesToMeters(4), 2);

  private final DCMotor gearbox = DCMotor.getNEO(1);

  private final LinearSystem<N1, N1, N1> plant =
      Models.flywheelFromPhysicalConstants(gearbox, kFlywheelGearing, kFlywheelMomentOfInertia);

  private final FlywheelSim flywheelSim = new FlywheelSim(plant, gearbox);
  private final EncoderSim encoderSim = new EncoderSim(encoder);

  public Robot() {
    // Add bang-bang controller to SmartDashboard and networktables.
    SmartDashboard.putData(bangBangController);
  }

  /** Controls flywheel to a set velocity (RPM) controlled by a joystick. */
  @Override
  public void teleopPeriodic() {
    // Scale setpoint value between 0 and maxSetpointValue
    double setpoint =
        Math.max(
            0.0,
            joystick.getRawAxis(0) * Units.rotationsPerMinuteToRadiansPerSecond(kMaxSetpointValue));

    // Set setpoint and measurement of the bang-bang controller
    double bangOutput = bangBangController.calculate(encoder.getRate(), setpoint) * 12.0;

    // Controls a motor with the output of the BangBang controller and a
    // feedforward. The feedforward is reduced slightly to avoid overspeeding
    // the shooter.
    flywheelMotor.setVoltage(bangOutput + 0.9 * feedforward.calculate(setpoint));
  }

  /** Update our simulation. This should be run every robot loop in simulation. */
  @Override
  public void simulationPeriodic() {
    // To update our simulation, we set motor voltage inputs, update the
    // simulation, and write the simulated velocities to our simulated encoder
    flywheelSim.setInputVoltage(flywheelMotor.getThrottle() * RobotController.getInputVoltage());
    flywheelSim.update(0.02);
    encoderSim.setRate(flywheelSim.getAngularVelocity());
  }
}
