001/************************* PROJECT RON *************************/
002/* Copyright (c) 2026 StuyPulse Robotics. All rights reserved. */
003/* Use of this source code is governed by an MIT-style license */
004/* that can be found in the repository LICENSE file.           */
005/***************************************************************/
006package com.stuypulse.robot.util;
007
008import edu.wpi.first.math.geometry.Translation2d;
009import edu.wpi.first.wpilibj.Timer;
010
011public class TranslationMotionProfile {
012
013    // Default number of times to apply filter (helps accuracy)
014    private static final int kDefaultSteps = 64;
015
016    // Limits for each of the derivatives
017    private final Number velocityLimit;
018
019    private final Number accelerationLimit;
020
021    // The last output / velocity
022    private Translation2d output;
023
024    private Translation2d velocity;
025
026    // Number of times to apply filter (helps accuracy)
027    private final int steps;
028
029    private double previousTime;
030
031    public TranslationMotionProfile(
032            Number velLimit,
033            Number accelLimit,
034            Translation2d startingTranslation,
035            Translation2d startingVelocity,
036            int steps) {
037        velocityLimit = velLimit;
038        accelerationLimit = accelLimit;
039        output = startingTranslation;
040        velocity = startingVelocity;
041        this.steps = steps;
042    }
043
044    public TranslationMotionProfile(
045            Number velLimit, Number accelLimit, Translation2d startingTranslation, Translation2d startingVelocity) {
046        this(velLimit, accelLimit, startingTranslation, startingVelocity, kDefaultSteps);
047    }
048
049    public TranslationMotionProfile(Number velLimit, Number accelLimit) {
050        this(velLimit, accelLimit, Translation2d.kZero, Translation2d.kZero, kDefaultSteps);
051    }
052
053    private static Translation2d clamp(Translation2d translation, double maxMagnitude) {
054        if (maxMagnitude <= 0.0) return Translation2d.kZero;
055
056        final double magnitude = translation.getNorm();
057
058        if (maxMagnitude <= magnitude) {
059            return translation.times(maxMagnitude / magnitude);
060        }
061
062        return translation;
063    }
064
065    public static Translation2d normalize(Translation2d translation) {
066        final double magnitude = translation.getNorm();
067        if (magnitude <= 1e-9) {
068            return new Translation2d(1, 0);
069        }
070        return translation.div(magnitude);
071    }
072
073    public Translation2d get(Translation2d target) {
074        double dt = resetTimer() / steps;
075        for (int i = 0; i < steps; ++i) {
076            // if there is a accel limit, limit the amount the velocity can change
077            if (0 < accelerationLimit.doubleValue()) {
078                // amount of windup in system (how long it would take to slow down)
079                double windup = velocity.getNorm() / accelerationLimit.doubleValue();
080                // If the windup is too small, just use normal algorithm to limit acceleration
081                if (windup < dt) {
082                    // Calculate acceleration needed to reach target
083                    Translation2d accel = target.minus(output).div(dt).minus(velocity);
084                    // Try to reach it while abiding by accel limit
085                    velocity = velocity.plus(clamp(accel, dt * accelerationLimit.doubleValue()));
086                } else {
087                    // the position it would end up if it attempted to come to a full stop
088                    Translation2d windA = // windup caused by acceleration
089                            velocity.times(0.5 * (dt + windup));
090                    // where the robot will end up
091                    Translation2d future = output.plus(windA);
092                    // Calculate acceleration needed to come to stop at target throughout windup
093                    Translation2d accel = target.minus(future).div(windup);
094                    // Try to reach it while abiding by accel limit
095                    velocity = velocity.plus(clamp(accel, dt * accelerationLimit.doubleValue()));
096                }
097            } else {
098                // make the velocity the difference between target and current
099                velocity = target.minus(output).div(dt);
100            }
101            // if there is an velocity limit, limit the velocity
102            if (0 < velocityLimit.doubleValue()) {
103                velocity = clamp(velocity, velocityLimit.doubleValue());
104            }
105            Translation2d error = target.minus(output);
106            Translation2d unitError = normalize(error);
107            double parallelMag = velocity.dot(unitError);
108            Translation2d accelParallel = unitError.times(parallelMag);
109            Translation2d accelPerpendicular = velocity.minus(accelParallel);
110            double damping = Math.pow(0.5, dt);
111            velocity = accelParallel.plus(accelPerpendicular.times(damping));
112            // adjust output by calculated velocity
113            output = output.plus(velocity.times(dt));
114        }
115        return output;
116    }
117
118    private double resetTimer() {
119        double now = Timer.getFPGATimestamp();
120        double elapsed = now - previousTime;
121        previousTime = now;
122        return elapsed;
123    }
124}