robot.py

#!/usr/bin/env python3
import enum
import math

import rev
import ctre
import magicbot
import wpilib

from automations.alignment import (
    HatchIntakeAligner,
    HatchDepositAligner,
    CargoDepositAligner,
)
from automations.cargo import CargoManager
from automations.climb import ClimbAutomation
from automations.hatch import HatchAutomation
from components.cargo import CargoManipulator
from components.hatch import Hatch
from components.vision import Vision
from components.climb import Climber
from pyswervedrive.chassis import SwerveChassis
from pyswervedrive.module import SwerveModule
from utilities.functions import constrain_angle, rescale_js
from utilities.navx import NavX


class FieldAngle(enum.Enum):
    CARGO_FRONT = 0
    CARGO_RIGHT = math.pi / 2
    CARGO_LEFT = -math.pi / 2
    LOADING_STATION = math.pi

    @classmethod
    def closest(cls, robot_heading: float) -> "FieldAngle":
        return min(cls, key=lambda a: abs(constrain_angle(robot_heading - a.value)))


class Robot(magicbot.MagicRobot):
    # Declare magicbot components here using variable annotations.
    # NOTE: ORDER IS IMPORTANT.
    # Any components that actuate objects should be declared after
    # any higher-level components (automations) that depend on them.

    # Automations
    cargo: CargoManager
    cargo_deposit: CargoDepositAligner
    climb_automation: ClimbAutomation
    hatch_deposit: HatchDepositAligner
    hatch_intake: HatchIntakeAligner
    hatch_automation: HatchAutomation

    # Actuators
    cargo_component: CargoManipulator
    chassis: SwerveChassis
    hatch: Hatch

    climber: Climber

    vision: Vision

    offset_rotation_rate = 20

    def createObjects(self):
        """Create motors and stuff here."""

        # a + + b - + c - - d + -
        x_dist = 0.2625
        y_dist = 0.2665
        self.module_a = SwerveModule(  # front left module
            "a",
            steer_talon=ctre.TalonSRX(3),
            drive_talon=ctre.TalonSRX(4),
            x_pos=x_dist,
            y_pos=y_dist,
        )
        self.module_b = SwerveModule(  # front right module
            "b",
            steer_talon=ctre.TalonSRX(7),
            drive_talon=ctre.TalonSRX(8),
            x_pos=-x_dist,
            y_pos=y_dist,
        )
        self.module_c = SwerveModule(  # bottom left module
            "c",
            steer_talon=ctre.TalonSRX(1),
            drive_talon=ctre.TalonSRX(6),
            x_pos=-x_dist,
            y_pos=-y_dist,
        )
        self.module_d = SwerveModule(  # bottom right module
            "d",
            steer_talon=ctre.TalonSRX(23),
            drive_talon=ctre.TalonSRX(24),
            x_pos=x_dist,
            y_pos=-y_dist,
        )
        self.imu = NavX()

        wpilib.SmartDashboard.putData("Gyro", self.imu.ahrs)

        # hatch objects
        self.hatch_fingers = wpilib.DoubleSolenoid(7, 6)
        self.hatch_punchers = wpilib.Solenoid(0)
        self.hatch_enable_piston = wpilib.DoubleSolenoid(3, 2)

        self.hatch_left_limit_switch = wpilib.DigitalInput(8)
        self.hatch_right_limit_switch = wpilib.DigitalInput(9)

        self.climber_front_motor = rev.CANSparkMax(10, rev.MotorType.kBrushless)
        self.climber_back_motor = rev.CANSparkMax(11, rev.MotorType.kBrushless)
        self.climber_front_podium_switch = wpilib.DigitalInput(4)
        self.climber_back_podium_switch = wpilib.DigitalInput(5)
        self.climber_drive_motor = ctre.TalonSRX(20)
        self.climber_pistons = wpilib.DoubleSolenoid(forwardChannel=4, reverseChannel=5)

        # cargo related objects
        self.intake_motor = ctre.VictorSPX(9)
        self.intake_switch = wpilib.DigitalInput(0)
        self.arm_motor = rev.CANSparkMax(2, rev.MotorType.kBrushless)

        # boilerplate setup for the joystick
        self.joystick = wpilib.Joystick(0)
        self.gamepad = wpilib.XboxController(1)

        self.spin_rate = 1.5

    def autonomous(self):
        self.imu.resetHeading()
        self.chassis.set_heading_sp(0)
        self.hatch.enable_hatch = True
        self.hatch_intake.alignment_speed = 0.5
        self.hatch_deposit.alignment_speed = 0.5
        self.chassis.derate_drive_modules(6)
        super().autonomous()

    def disabledPeriodic(self):
        self.chassis.set_inputs(0, 0, 0)
        self.vision.execute()  # Keep the time offset calcs running

    def teleopInit(self):
        """Initialise driver control."""
        self.chassis.set_inputs(0, 0, 0)
        self.hatch_intake.alignment_speed = 0.5
        self.hatch_deposit.alignment_speed = 0.5
        self.chassis.derate_drive_modules(9)

    def teleopPeriodic(self):
        """Allow the drivers to control the robot."""
        throttle = max(0.1, (1 - self.joystick.getThrottle()) / 2)  # min 10%

        # this is where the joystick inputs get converted to numbers that are sent
        # to the chassis component. we rescale them using the rescale_js function,
        # in order to make their response exponential, and to set a dead zone -
        # which just means if it is under a certain value a 0 will be sent
        # TODO: Tune these constants for whatever robot they are on
        joystick_vx = -rescale_js(
            self.joystick.getY(), deadzone=0.1, exponential=1.5, rate=4 * throttle
        )
        joystick_vy = -rescale_js(
            self.joystick.getX(), deadzone=0.1, exponential=1.5, rate=4 * throttle
        )
        joystick_vz = -rescale_js(
            self.joystick.getZ(), deadzone=0.2, exponential=20.0, rate=self.spin_rate
        )
        joystick_hat = self.joystick.getPOV()

        # Allow big stick movements from the driver to break out of an automation
        if abs(joystick_vx) > 0.5 or abs(joystick_vy) > 0.5:
            self.hatch_intake.done()
            self.hatch_deposit.done()
            self.cargo_deposit.done()

        if not self.chassis.automation_running:
            if joystick_vx or joystick_vy or joystick_vz:
                self.chassis.set_inputs(
                    joystick_vx,
                    joystick_vy,
                    joystick_vz,
                    field_oriented=not self.joystick.getRawButton(6),
                )
            else:
                self.chassis.set_inputs(0, 0, 0)

            if joystick_hat != -1:
                if self.cargo_component.has_cargo or self.cargo.is_executing:
                    constrained_angle = -constrain_angle(
                        math.radians(joystick_hat) + math.pi
                    )
                else:
                    constrained_angle = -constrain_angle(math.radians(joystick_hat))
                self.chassis.set_heading_sp(constrained_angle)

        # Starts Hatch Alignment and Cargo State Machines
        if (
            self.joystick.getTrigger()
            or self.gamepad.getTriggerAxis(self.gamepad.Hand.kLeft) > 0.5
            or self.gamepad.getTriggerAxis(self.gamepad.Hand.kRight) > 0.5
        ):
            angle = FieldAngle.closest(self.imu.getAngle())
            self.logger.info("closest field angle: %s", angle)
            if self.cargo_component.has_cargo:
                self.cargo_deposit.engage()
            else:
                if angle is FieldAngle.LOADING_STATION:
                    self.hatch_intake.engage()
                else:
                    self.hatch_deposit.engage()
            self.chassis.set_heading_sp(angle.value)

        # Hatch Manual Outake/Intake
        if self.joystick.getRawButtonPressed(5) or self.gamepad.getBumperPressed(6):
            angle = FieldAngle.closest(self.imu.getAngle())
            self.logger.info("closest field angle: %s", angle)
            if angle is not FieldAngle.LOADING_STATION:
                self.hatch_automation.outake()
            else:
                self.hatch_automation.grab()

        if self.gamepad.getXButtonPressed():
            self.hatch.retract_fingers()
            self.hatch.retract()

        # Stops Cargo Intake Motor
        if self.gamepad.getBButtonPressed():
            self.cargo.outake_cargo_ship(force=True)

        # Toggles the Heading Hold
        if self.joystick.getRawButtonPressed(8):
            if self.chassis.hold_heading:
                self.chassis.heading_hold_off()
            else:
                self.chassis.heading_hold_on()

        # Resets the IMU's Heading
        if self.joystick.getRawButtonPressed(7):
            self.imu.resetHeading()
            self.chassis.set_heading_sp(0)

        # Start Button starts Climb State Machine
        if self.gamepad.getStartButtonPressed() and self.gamepad.getRawButtonPressed(5):
            self.climb_automation.start_climb_lv3()

        # Back Button Ends Climb State Machine
        if self.gamepad.getBackButtonPressed():
            if self.gamepad.getRawButtonPressed(5):
                self.climb_automation.abort()
            else:
                self.climb_automation.done()

        # Cargo Floor Intake
        if self.gamepad.getAButtonPressed():
            self.cargo.intake_floor(force=True)

        # Cargo Loading Station Intake
        if self.gamepad.getYButtonPressed():
            self.cargo.intake_loading(force=True)
            self.chassis.set_heading_sp(
                FieldAngle.CARGO_FRONT.value
            )  # Reversed side of robot

        if self.gamepad.getPOV() != -1:
            speed = 0.65
            azimuth = math.radians(-self.gamepad.getPOV())
            if self.cargo_component.has_cargo:
                azimuth += math.pi
            self.chassis.set_inputs(
                speed * math.cos(azimuth),
                speed * math.sin(azimuth),
                0,
                field_oriented=False,
            )

    def robotPeriodic(self):
        # super().robotPeriodic()
        wpilib.SmartDashboard.updateValues()

    def testPeriodic(self):
        self.vision.execute()  # Keep the time offset calcs running

        joystick_vx = -rescale_js(
            self.joystick.getY(), deadzone=0.1, exponential=1.5, rate=0.5
        )
        for button, module in zip((5, 3, 4, 6), self.chassis.modules):
            if self.joystick.getRawButton(button):
                module.store_steer_offsets()
                module.steer_motor.set(ctre.ControlMode.PercentOutput, joystick_vx)
                if self.joystick.getTriggerPressed():
                    module.steer_motor.set(
                        ctre.ControlMode.Position,
                        module.steer_motor.getSelectedSensorPosition(0)
                        + self.offset_rotation_rate,
                    )
                if self.joystick.getRawButtonPressed(2):
                    module.steer_motor.set(
                        ctre.ControlMode.Position,
                        module.steer_motor.getSelectedSensorPosition(0)
                        - self.offset_rotation_rate,
                    )

        if self.joystick.getRawButtonPressed(8):
            for module in self.chassis.modules:
                module.drive_motor.set(ctre.ControlMode.PercentOutput, 0.3)

        if self.joystick.getRawButtonPressed(12):
            for module in self.chassis.modules:
                module.steer_motor.set(
                    ctre.ControlMode.Position, module.steer_enc_offset
                )

        if self.gamepad.getStartButton():
            self.climber.retract_all()
            self.climber.execute()

        if self.gamepad.getPOV() != -1:
            speed = 0.1
            azimuth = math.radians(-self.gamepad.getPOV())
            for module in self.chassis.modules:
                module.set_velocity(
                    speed * math.cos(azimuth),
                    speed * math.sin(azimuth),
                    absolute_rotation=True,
                )

        if self.gamepad.getTriggerAxis(self.gamepad.Hand.kLeft) > 0.5:
            self.hatch_enable_piston.set(wpilib.DoubleSolenoid.Value.kReverse)


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