package frc.robot.subsystems;

import org.littletonrobotics.junction.Logger;

import com.revrobotics.CANSparkMax;
import com.revrobotics.RelativeEncoder;
import com.revrobotics.CANSparkLowLevel.MotorType;

import edu.wpi.first.math.MathUtil;
import edu.wpi.first.math.controller.ArmFeedforward;
import edu.wpi.first.math.controller.PIDController;
import edu.wpi.first.math.controller.ProfiledPIDController;
import edu.wpi.first.units.Angle;
import edu.wpi.first.units.Measure;
import edu.wpi.first.units.MutableMeasure;
import edu.wpi.first.units.Velocity;
import edu.wpi.first.units.Voltage;
import edu.wpi.first.wpilibj2.command.Command;
import edu.wpi.first.wpilibj2.command.PIDCommand;
import edu.wpi.first.wpilibj2.command.SubsystemBase;
import edu.wpi.first.wpilibj2.command.sysid.SysIdRoutine;
import frc.robot.Constants;

import static edu.wpi.first.units.MutableMeasure.mutable;
import static edu.wpi.first.units.Units.Volts;
import static edu.wpi.first.units.Units.Radians;
import static edu.wpi.first.units.Units.RadiansPerSecond;
import static edu.wpi.first.units.Units.Seconds;

public class Arm extends SubsystemBase 
{
    // rad
    private double m_setpoint = Constants.arm_initial_position;
    private final CANSparkMax m_armL = new CANSparkMax(Constants.armL_ID, MotorType.kBrushless);
    private final CANSparkMax m_armR = new CANSparkMax(Constants.armR_ID, MotorType.kBrushless);

    private final RelativeEncoder m_encoderL = m_armL.getEncoder();
    private final RelativeEncoder m_encoderR = m_armR.getEncoder();

    private ArmFeedforward m_feedforward = new ArmFeedforward(Constants.ArmkS, Constants.ArmkG, Constants.ArmkV);

    private ProfiledPIDController m_pidL = new ProfiledPIDController(Constants.ArmkP, Constants.ArmkI, Constants.ArmkD, Constants.arm_constraints);
    private ProfiledPIDController m_pidR = new ProfiledPIDController(Constants.ArmkP, Constants.ArmkI, Constants.ArmkD, Constants.arm_constraints);
    
    // SYSID
    private final MutableMeasure<Voltage> m_voltage = mutable(Volts.of(0));
    private final MutableMeasure<Angle> m_angle = mutable(Radians.of(0));
    private final MutableMeasure<Velocity<Angle>> m_velocity = mutable(RadiansPerSecond.of(0));

    private final SysIdRoutine m_sysIdRoutine = new SysIdRoutine(
    new SysIdRoutine.Config(
        Volts.per(Seconds).of(0.5),
        Volts.of(3),
        Seconds.of(10)),
    new SysIdRoutine.Mechanism(
        (Measure<Voltage> volts) -> {
            m_armL.setVoltage(volts.in(Volts));
            m_armR.setVoltage(volts.in(Volts));
        },
        log -> {
            log.motor("arm-left")
                .voltage(
                    m_voltage.mut_replace(
                        m_armL.getAppliedOutput() * m_armL.getBusVoltage(), Volts)
                )
                .angularPosition(m_angle.mut_replace(m_encoderL.getPosition(), Radians))
                .angularVelocity(m_velocity.mut_replace(m_encoderL.getVelocity(), RadiansPerSecond));
            log.motor("arm-right")

                .voltage(
                    m_voltage.mut_replace(
                        m_armR.getAppliedOutput() * m_armR.getBusVoltage(), Volts)
                )
                .angularPosition(m_angle.mut_replace(m_encoderR.getPosition(), Radians))
                .angularVelocity(m_velocity.mut_replace(m_encoderR.getVelocity(), RadiansPerSecond));
            
        }, this));

    public Arm()
    {   
        m_armL.setInverted(false);
        m_armR.setInverted(true);

        m_encoderL.setPositionConversionFactor(2*Math.PI / Constants.ArmGearReduction);
        m_encoderR.setPositionConversionFactor(2*Math.PI / Constants.ArmGearReduction);
        reset();
    }

    // public Arm(double kP, double kI, double kD, double kS, double kG, double kV)
    // {   
    //     this();

    //     m_feedforward = new ArmFeedforward(kS, kG, kV);
    //     m_pidL = new ProfiledPIDController(kP, kI, kD, Constants.arm_constraints);
    //     m_pidR = new ProfiledPIDController(kP, kI, kD, Constants.arm_constraints);
    // }

    public void reset()
    {
        m_setpoint = Constants.arm_initial_position;
        m_encoderL.setPosition(Constants.arm_initial_position);
        m_encoderR.setPosition(Constants.arm_initial_position);

        m_pidL.reset(Constants.arm_initial_position);
        m_pidR.reset(Constants.arm_initial_position);
    }


    public void incrementAngleDegrees(double amount) {
        incrementAngle(amount * Math.PI / 180.0);
    }

    public void incrementAngle(double amount) {
        atAngle(m_setpoint + amount);
    }
    

    // in rad
    public void atAngle(double angle)
    {
        m_setpoint = angle;
        m_pidL.setGoal(m_setpoint);
        m_pidR.setGoal(m_setpoint);

        voltage(
            m_pidL.calculate(m_encoderL.getPosition()),
            m_feedforward.calculate(m_pidL.getSetpoint().position, m_pidL.getSetpoint().velocity),

            m_pidR.calculate(m_encoderR.getPosition()),
            m_feedforward.calculate(m_pidR.getSetpoint().position, m_pidR.getSetpoint().velocity)
        );
    }
    // in deg
    public void atAngleDegrees(double angle) 
    {
        atAngle(angle * Math.PI / 180.0);
    }

    // in rad
    public boolean hasClearedAngle(double angle) 
    {
        if(m_encoderL.getPosition() > angle
        && m_encoderR.getPosition() > angle) 
            return true;

        return false;
    }

    public boolean hasClearedAngleDegree(double angle)
    {
        return hasClearedAngle(angle * Math.PI/180);
    }
    // in rad
    public boolean isAtAngle(double angle) 
    {
        if(Math.abs(m_encoderL.getPosition() - angle) < Constants.ArmEpsilon
        && Math.abs(m_encoderR.getPosition() - angle) < Constants.ArmEpsilon) 
            return true;

        return false;
    }

    public boolean isAtAngleDegree(double angle)
    {
        return isAtAngle(angle * Math.PI/180);
    }

    public void voltage(double leftVoltsPID, double leftVoltsFFW, double rightVoltsPID, double rightVoltsFFW)
    {   
        Logger.recordOutput(getName() + "/leftPID", leftVoltsPID);
        Logger.recordOutput(getName() + "/rightPID", rightVoltsPID);
        Logger.recordOutput(getName() + "/leftFFW", leftVoltsFFW);
        Logger.recordOutput(getName() + "/rightFFW", rightVoltsFFW);

        m_armL.setVoltage(MathUtil.clamp(leftVoltsPID + leftVoltsFFW, -4, 4));
        m_armR.setVoltage(MathUtil.clamp(rightVoltsPID + rightVoltsFFW, -4, 4));
    }

    @Override
    public void periodic() 
    {
        Logger.recordOutput(getName() + "/positionL", m_encoderL.getPosition() * 180/Math.PI);
        Logger.recordOutput(getName() + "/positionR", m_encoderR.getPosition() * 180/Math.PI);
    }

    public Command sysIdQausistatic(SysIdRoutine.Direction direction)
    {
        return m_sysIdRoutine.quasistatic(direction);
    }

    public Command sysIdDynamic(SysIdRoutine.Direction direction)
    {
        return m_sysIdRoutine.dynamic(direction);
    }
}