Swerve.h

#pragma once
 
#include <deque>
#include <memory>
#include <string>
#include <vector>
 
#include <Eigen/Core>
#include <choreo/Choreo.h>
#include <ctre/phoenix6/swerve/impl/SwerveDrivePoseEstimator.hpp>
#include <frc/DriverStation.h>
#include <frc/EigenCore.h>
#include <frc/Notifier.h>
#include <frc/controller/ProfiledPIDController.h>
#include <frc/estimator/SwerveDrivePoseEstimator.h>
#include <frc/geometry/Pose2d.h>
#include <frc/kinematics/ChassisSpeeds.h>
#include <frc/kinematics/SwerveDriveKinematics.h>
#include <frc/trajectory/TrapezoidProfile.h>
#include <frc2/command/InstantCommand.h>
#include <networktables/StructArrayTopic.h>
#include <networktables/StructTopic.h>
#include <units/acceleration.h>
#include <units/angle.h>
#include <units/angular_velocity.h>
#include <units/dimensionless.h>
#include <units/length.h>
#include <units/time.h>
#include <units/velocity.h>
#include <wpi/interpolating_map.h>
 
#include "valkyrie/BaseSubsystem.h"
#include "valkyrie/controllers/BaseController.h"
#include "valkyrie/controllers/PIDF.h"
#include "valkyrie/drivetrain/swerve/SwerveModule.h"
#include "valkyrie/gyro/Gyro.h"
#include "valkyrie/util/CoordinateFrames.h"
#include "valkyrie/util/LockedAccess.h"
 
namespace valor {
namespace swerve {
 
class Swerve : public valor::BaseSubsystem {
   public:
    Swerve(std::string name, std::vector<std::pair<std::unique_ptr<BaseController>, std::unique_ptr<BaseController>>> modules,
           std::shared_ptr<gyro::Gyro> gyro, units::meter_t module_radius, units::meter_t _wheelDiameter, std::vector<int> moduleCoordsX,
           std::vector<int> moduleCoordsY, units::meter_t driveBaseRadius);
 
    void AnalyzeDashboard() override;
    void LoggablePeriodic() override;
    void Reset() override;
 
    void SimulationPeriodic() override;
 
    void ResetGyro();
 
    void ApplyRequest(requests::ApplyChassisSpeeds);
    void ApplyRequest(requests::ApplyStates);

    void ResetOdometry(FrameWrapper<FieldFrame, frc::Pose2d> pose);
    void ResetEncoders();

    std::vector<frc::SwerveModulePosition> GetModulePositions();

    std::vector<frc::SwerveModuleState> GetModuleStates();
 
    double GetSkiddingRatio();
    bool IsRobotSkidding();
 
    constexpr std::vector<std::unique_ptr<SwerveModule>>& GetModules() { return swerveModules; }

    void UpdateAngularAcceleration();
    units::angular_acceleration::radians_per_second_squared_t GetSmoothedAngularAcceleration();
    double RotationLerping(double);
    void FollowTrajectory(const choreo::SwerveSample& sample);
 
    constexpr units::meter_t GetDriveBaseRadius() const { return driveBaseRadius; };
 
    constexpr const auto& GetState() const { return state; }
 
    constexpr auto& GetState() { return state; }
 
   protected:
    FrameWrapper<FieldFrame, frc::ChassisSpeeds> commandedSpeeds;
 
    units::meters_per_second_t maxDriveSpeed;
    units::radians_per_second_t maxRotationSpeed;
 
    std::shared_ptr<gyro::Gyro> gyro;
    std::unique_ptr<ctre::phoenix6::swerve::impl::SwerveDriveKinematics> kinematics;
    util::Locked<std::unique_ptr<ctre::phoenix6::swerve::impl::SwerveDrivePoseEstimator>> rawEstimator;
    util::Locked<std::unique_ptr<ctre::phoenix6::swerve::impl::SwerveDrivePoseEstimator>> calcEstimator;
 
    bool toast;
 
    void EnableCarpetGrain(double grainMultiplier, bool roughTowardsRed);
 
    FrameWrapper<FieldFrame, frc::ChassisSpeeds> VectorLockedSpeeds(FrameWrapper<FieldFrame, frc::ChassisSpeeds> speeds,
                                                                    FrameWrapper<FieldFrame, frc::Rotation2d> heading);
 
    std::vector<std::unique_ptr<SwerveModule>> swerveModules;
 
    void StartOdomThread(units::hertz_t freq);
    void UpdateOdometry();
 
    void SetRotationalLerpMap(wpi::interpolating_map<double, double> _interpolationMap);
 
    static FrameWrapper<FieldFrame, frc::ChassisSpeeds> ReflectDriveSpeeds(FrameWrapper<FieldFrame, frc::ChassisSpeeds> unreflectedSpeeds) {
        if (frc::DriverStation::GetAlliance().value_or(frc::DriverStation::Alliance::kBlue) == frc::DriverStation::Alliance::kRed) {
            return frc::ChassisSpeeds{-unreflectedSpeeds->vx, -unreflectedSpeeds->vy, unreflectedSpeeds->omega};
        }
        return unreflectedSpeeds;
    }
 
    frc::Timer resetOdom;
 
   private:
    std::deque<units::angular_acceleration::radians_per_second_squared_t> yawRateBuffer;  // TODO Replace with a moving average filter
 
    static constexpr size_t ACCEL_BUFFER_SIZE = 10;  // need to adjust // TODO: Replace with an actual moving average filter
    units::angular_velocity::radians_per_second_t lastYawRate = 0_rad_per_s;  // TODO: Replace with a moving average filter
 
    units::angular_acceleration::radians_per_second_squared_t angularAcceleration = 0_rad_per_s_sq;
 
    bool useCarpetGrain;
    double carpetGrainMultiplier;
    bool roughTowardsRed;
    void CalculateCarpetPose();
 
    struct {
        FrameWrapper<FieldFrame, frc::Pose2d> rawPose;
        FrameWrapper<FieldFrame, frc::Pose2d> calculatedPose;
        FrameWrapper<FieldFrame, frc::ChassisSpeeds> fieldSpeeds;
        FrameWrapper<RobotFrame, frc::ChassisSpeeds> robotSpeeds;
 
        std::vector<frc::SwerveModuleState> moduleStates;
 
        frc::Rotation3d gyroRotation;
        gyro::Gyro::AngularVelocity3d gyroVelocity;
    } state;
 
    units::meter_t driveBaseRadius;
 
    frc::Notifier odometryThread{[this] { UpdateOdometry(); }};
 
    frc::PIDController xChoreoController{10.0, 0.0, 0.0};
    frc::PIDController yChoreoController{10.0, 0.0, 0.0};
    frc::PIDController headingChoreoController{7.5, 0.0, 0.0};
 
    wpi::interpolating_map<double, double> rotationalLerpMap;
 
    // SwerveSetpointGenerator setpointGenerator;
    // SwerveSetpoint previousSetpoint;
};
 
}  // namespace swerve
}  // namespace valor