import wpilib
import wpilib.drive
import wpimath.controller
import math
import ctre

def clamp(x, low, high):
    if x < low:
        return low
    elif x > high:
        return high
    else:
        return x

class MyRobot(wpilib.TimedRobot):
    def robotInit(self):
        self.leftFront = ctre.WPI_VictorSPX(0)
        self.leftRear = ctre.WPI_VictorSPX(1)
        self.rightFront = ctre.WPI_VictorSPX(2)
        self.rightRear = ctre.WPI_VictorSPX(3)

        self.left = wpilib.MotorControllerGroup(self.leftFront, self.leftRear)
        self.right = wpilib.MotorControllerGroup(self.rightFront, self.rightRear)
        self.right.setInverted(True)

        self.drivetrain = wpilib.drive.DifferentialDrive(self.left, self.right)

        self.primary = wpilib.XboxController(0)

        self.PULSES_PER_REVOLUTION = 512
        self.WHEEL_DIAMETER = 6
        self.DISTANCE_PER_PULSE = math.pi * self.WHEEL_DIAMETER / self.PULSES_PER_REVOLUTION

        self.leftDrivetrainEncoder = wpilib.Encoder(0, 1)
        self.leftDrivetrainEncoder.setDistancePerPulse(self.DISTANCE_PER_PULSE)
        self.rightDrivetrainEncoder = wpilib.Encoder(2, 3)
        self.rightDrivetrainEncoder.setDistancePerPulse(self.DISTANCE_PER_PULSE)

        self.gyro = wpilib.ADXRS450_Gyro()
        self.gyro.calibrate()

        self.ENCODER_PID_P = .1
        self.ENCODER_PID_I = 0
        self.ENCODER_PID_D = .01
        self.ENCODER_PID_TOLERANCE = 1
        self.AUTO_ENCODER_SETPOINT = 60

        self.GYRO_PID_P = .1
        self.GYRO_PID_I = 0
        self.GYRO_PID_D = .01
        self.GYRO_PID_TOLERANCE = 1
        self.AUTO_GYRO_SETPOINT = 90

        self.leftDrivetrainController = wpimath.controller.PIDController(
            self.ENCODER_PID_P,
            self.ENCODER_PID_I,
            self.ENCODER_PID_D
        )
        self.leftDrivetrainController.setTolerance(self.ENCODER_PID_TOLERANCE)
        self.leftDrivetrainController.setSetpoint(self.AUTO_ENCODER_SETPOINT)

        self.rightDrivetrainController = wpimath.controller.PIDController(
            self.ENCODER_PID_P,
            self.ENCODER_PID_I,
            self.ENCODER_PID_D
        )
        self.rightDrivetrainController.setTolerance(self.ENCODER_PID_TOLERANCE)
        self.rightDrivetrainController.setSetpoint(self.AUTO_ENCODER_SETPOINT)

        self.gyroDrivetrainController = wpimath.controller.PIDController(
            self.GYRO_PID_P,
            self.GYRO_PID_I, 
            self.GYRO_PID_D
        )
        self.gyroDrivetrainController.setTolerance(self.GYRO_PID_TOLERANCE)
        self.gyroDrivetrainController.setSetpoint(self.AUTO_GYRO_SETPOINT)


    def autonomousInit(self):
        self.leftDrivetrainEncoder.reset()
        self.rightDrivetrainEncoder.reset()

        self.autoState = "Encoder Step"

    def autonomousPeriodic(self):
        if self.autoState == "Encoder Step":
            if not self.leftDrivetrainController.atSetpoint() and not self.rightDrivetrainController.atSetpoint():
                self.drivetrain.tankDrive(
                    clamp(
                        self.leftDrivetrainController.calculate(self.leftDrivetrainEncoder.getDistance()),
                        -.5, 
                        .5
                    ),
                    clamp(
                        self.rightDrivetrainController.calculate(self.rightDrivetrainEncoder.getDistance()),
                        -.5,
                        .5
                    )
                )
            else:
                self.drivetrain.tankDrive(0, 0)
                self.gyro.reset()
                self.autoState = "Gyro Step"
        elif self.autoState == "Gyro Step":
            if not self.gyroDrivetrainController.atSetpoint():
                speedValue = clamp(
                    self.gyroDrivetrainController.calculate(self.gyro.getAngle()), 
                    -.5, 
                    .5
                )

                self.drivetrain(speedValue, -speedValue)
            else:
                self.drivetrain.tankDrive(0, 0)
                self.autoState = "End Step"
        else:
            self.drivetrain.tankDrive(0, 0)

    def teleopInit(self):
        pass

    def teleopPeriodic(self):
        self.drivetrain.arcadeDrive(self.primary.getLeftY(), self.primary.getLeftX())

if __name__ == "__main__":
    wpilib.run(MyRobot)