TrajectoryFollower.h

#pragma once
 
#include <algorithm>
#include <cmath>
#include <functional>
#include <optional>
 
#include <choreo/trajectory/SwerveSample.h>
#include <choreo/trajectory/Trajectory.h>
#include <frc/Timer.h>
#include <frc/apriltag/AprilTagFieldLayout.h>
#include <frc/controller/ProfiledPIDController.h>
#include <frc/geometry/Pose2d.h>
#include <frc/geometry/Rotation2d.h>
#include <frc/kinematics/ChassisSpeeds.h>
#include <frc/kinematics/SwerveModuleState.h>
#include <frc2/command/Command.h>
#include <frc2/command/CommandHelper.h>
#include <frc2/command/CommandPtr.h>
#include <frc2/command/Commands.h>
#include <frc2/command/FunctionalCommand.h>
#include <frc2/command/SequentialCommandGroup.h>
#include <units/acceleration.h>
#include <units/time.h>
#include <valkyrie/Loggable.h>
#include <valkyrie/WorldAlign.h>
#include <valkyrie/controllers/BaseController.h>
#include <valkyrie/controllers/PIDController.h>
#include <valkyrie/controllers/PIDF.h>
#include <valkyrie/drivetrain/swerve/Swerve.h>
 
namespace valor {

struct AccelerationState {
    units::meters_per_second_squared_t ax;
    units::meters_per_second_squared_t ay;
    units::radians_per_second_squared_t alpha;
};

struct AutoSwerveState {
    units::second_t timestamp;
    frc::Pose2d pose;
    frc::ChassisSpeeds speeds;
    AccelerationState accel;
};

class TrajectoryFollower {
   public:
    TrajectoryFollower() = delete;
    TrajectoryFollower(TrajectoryFollower&) = delete;
    TrajectoryFollower(TrajectoryFollower&&) = delete;

    enum class TrajectoryEndCondition { TIME, POSITION };
 
    enum class FieldSymmetry { MIRROR, ROTATE };

    static void Register(swerve::Swerve& s) { swerve = &s; }

    static frc2::CommandPtr FollowPath(
        std::optional<choreo::Trajectory<choreo::SwerveSample>> traj, TrajectoryEndCondition end, bool isMirrored, bool resetOdom,
        units::meter_t tolerance = 0.25_m, units::second_t timeWithoutCorrection = 0_s,
        std::function<bool()> interruption = [] { return false; });

    static void SetPID(WorldAlign::TranslationalPIDF_t x, WorldAlign::TranslationalPIDF_t y, WorldAlign::RotationalPIDF_t rot);
 
   private:
    class FollowCommand : public frc2::CommandHelper<frc2::Command, FollowCommand> {
       public:
        FollowCommand(choreo::Trajectory<choreo::SwerveSample> traj, TrajectoryEndCondition end, bool isMirrored, units::meter_t tolerance,
                      units::second_t timeWithoutCorrection, std::function<bool()> interruption)
              : traj{traj},
                end{end},
                isMirrored{isMirrored},
                tolerance{tolerance},
                timeWithoutCorrection{timeWithoutCorrection},
                interruption{interruption} {}
 
        void Initialize() override;
 
        void Execute() override;
 
        bool IsFinished() override;
 
       private:
        choreo::Trajectory<choreo::SwerveSample> traj;
        TrajectoryEndCondition end;
        bool isMirrored;
        units::meter_t tolerance;
        units::second_t timeWithoutCorrection;
        std::function<bool()> interruption;
 
        frc::Timer timer;
        PIDController<units::meter, units::meters_per_second> xController{{}};
        PIDController<units::meter, units::meters_per_second> yController{{}};
        PIDController<units::radian, units::radians_per_second> rotController{{}};
    };

    static AutoSwerveState ConvertChoreoSampleIntoSwerveState(choreo::SwerveSample sample);
    static AutoSwerveState Interop(units::second_t time, std::vector<AutoSwerveState> traj);
    static AutoSwerveState FlipSwerveState(AutoSwerveState state, FieldSymmetry sym);
    static AutoSwerveState MirrorSwerveState(AutoSwerveState state);
 
    static bool IsAllianceFlipped();
 
    static WorldAlign::TranslationalPIDF_t xPIDF, yPIDF;
    static WorldAlign::RotationalPIDF_t rotPIDF;
 
    static swerve::Swerve* swerve;
 
    static Loggable loggable;
};
}  // namespace valor