FlyShotHost/src/Flyshot.ControllerClientCompat/MoveJointTrajectoryGenerator.cs

using Flyshot.Core.Domain;

namespace Flyshot.ControllerClientCompat;

internal static class MoveJointTrajectoryGenerator
{
    private const double BaseMoveJointDurationSeconds = 0.320;
    private const double VelocityShapeCoefficient = 2.0759961613199973;
    private const double AccelerationShapeCoefficient = 7.986313199999984;
    private const double JerkShapeCoefficient = 36.12609273600853;
    private const int MaxMoveJointSampleCount = 1_000_000;

    private static readonly double[] CapturedMvpointAlpha =
    [
        0.000000000000,
        0.000012196163,
        0.000106156906,
        0.000764380061,
        0.002550804028,
        0.006029689194,
        0.011765134027,
        0.020321400844,
        0.032262426551,
        0.048152469303,
        0.068555498563,
        0.093895155669,
        0.124210027377,
        0.159174512929,
        0.198230386318,
        0.240813559900,
        0.286359937276,
        0.334305411725,
        0.384085546646,
        0.435136609163,
        0.486894129077,
        0.538794033110,
        0.590272360135,
        0.640764719629,
        0.689707151220,
        0.736535405849,
        0.780685354316,
        0.821592775628,
        0.858693734065,
        0.891423926949,
        0.919286047395,
        0.942156722091,
        0.960255163676,
        0.974119666692,
        0.984314536393,
        0.991403790959,
        0.995951593494,
        0.998522142663,
        0.999679443354,
        0.999987892657,
        1.000000000000
    ];

    public static TrajectoryResult CreateResult(
        RobotProfile robot,
        IReadOnlyList<double> startJoints,
        IReadOnlyList<double> targetJoints,
        double speedRatio)
    {
        ArgumentNullException.ThrowIfNull(robot);
        ArgumentNullException.ThrowIfNull(startJoints);
        ArgumentNullException.ThrowIfNull(targetJoints);

        if (speedRatio <= 0.0 || double.IsNaN(speedRatio) || double.IsInfinity(speedRatio))
        {
            throw new InvalidOperationException("Speed ratio must be greater than zero for MoveJoint execution.");
        }

        if (startJoints.Count != robot.DegreesOfFreedom || targetJoints.Count != robot.DegreesOfFreedom)
        {
            throw new InvalidOperationException($"MoveJoint expects {robot.DegreesOfFreedom} joints.");
        }

        var requestedDurationSeconds = ResolveDurationSeconds(robot, startJoints, targetJoints);
        var samplePeriodSeconds = robot.ServoPeriod.TotalSeconds * speedRatio;
        var durationSeconds = AlignDurationToServoStep(requestedDurationSeconds, samplePeriodSeconds);
        var denseJointTrajectory = GenerateDenseTrajectory(startJoints, targetJoints, durationSeconds, samplePeriodSeconds);

        return new TrajectoryResult(
            programName: "move-joint",
            method: PlanningMethod.Doubles,
            isValid: true,
            duration: TimeSpan.FromSeconds(durationSeconds),
            shotEvents: Array.Empty<ShotEvent>(),
            triggerTimeline: Array.Empty<TrajectoryDoEvent>(),
            artifacts: Array.Empty<TrajectoryArtifact>(),
            failureReason: null,
            usedCache: false,
            originalWaypointCount: 2,
            plannedWaypointCount: denseJointTrajectory.Count,
            denseJointTrajectory: denseJointTrajectory);
    }

    internal static double ResolveDurationSeconds(RobotProfile robot, IReadOnlyList<double> startJoints, IReadOnlyList<double> targetJoints)
    {
        var duration = BaseMoveJointDurationSeconds;

        for (var index = 0; index < robot.DegreesOfFreedom; index++)
        {
            var distance = Math.Abs(targetJoints[index] - startJoints[index]);
            if (distance <= 0.0)
            {
                continue;
            }

            var limit = robot.JointLimits[index];
            var velocityDuration = distance * VelocityShapeCoefficient / limit.VelocityLimit;
            var accelerationDuration = Math.Sqrt(distance * AccelerationShapeCoefficient / limit.AccelerationLimit);
            var jerkDuration = Math.Cbrt(distance * JerkShapeCoefficient / limit.JerkLimit);

            duration = Math.Max(duration, Math.Max(velocityDuration, Math.Max(accelerationDuration, jerkDuration)));
        }

        return duration;
    }

    internal static double AlignDurationToServoStep(double durationSeconds, double samplePeriodSeconds)
    {
        if (samplePeriodSeconds <= 0.0 || double.IsNaN(samplePeriodSeconds) || double.IsInfinity(samplePeriodSeconds))
        {
            throw new InvalidOperationException("Speed ratio must be greater than zero for MoveJoint execution.");
        }

        var intervals = ResolveSampleIntervalCount(durationSeconds, samplePeriodSeconds);

        return Math.Max(1, intervals) * samplePeriodSeconds;
    }

    private static long ResolveSampleIntervalCount(double durationSeconds, double samplePeriodSeconds)
    {
        var rawIntervals = durationSeconds / samplePeriodSeconds;
        if (double.IsNaN(rawIntervals) || double.IsInfinity(rawIntervals))
        {
            throw new InvalidOperationException("MoveJoint sample count is not representable.");
        }

        var intervals = (long)Math.Ceiling(rawIntervals - 1e-9);
        if (intervals < 1 || intervals + 1 > MaxMoveJointSampleCount)
        {
            throw new InvalidOperationException($"MoveJoint sample count must be between 2 and {MaxMoveJointSampleCount}.");
        }

        return intervals;
    }

    internal static IReadOnlyList<IReadOnlyList<double>> GenerateDenseTrajectory(
        IReadOnlyList<double> startJoints,
        IReadOnlyList<double> targetJoints,
        double durationSeconds,
        double samplePeriodSeconds)
    {
        var sampleCount = ResolveSampleIntervalCount(durationSeconds, samplePeriodSeconds) + 1;
        var rows = new List<IReadOnlyList<double>>(checked((int)sampleCount));

        for (var index = 0L; index < sampleCount; index++)
        {
            var time = Math.Min(index * samplePeriodSeconds, durationSeconds);
            rows.Add(CreateRow(time, durationSeconds, startJoints, targetJoints));
        }

        return rows;
    }

    private static IReadOnlyList<double> CreateRow(double timeSeconds, double durationSeconds, IReadOnlyList<double> startJoints, IReadOnlyList<double> targetJoints)
    {
        var u = durationSeconds <= 0.0 ? 1.0 : Math.Clamp(timeSeconds / durationSeconds, 0.0, 1.0);
        var alpha = InterpolateCapturedAlpha(u);
        var row = new double[startJoints.Count + 1];
        row[0] = Math.Round(timeSeconds, 9);

        for (var index = 0; index < startJoints.Count; index++)
        {
            row[index + 1] = startJoints[index] + ((targetJoints[index] - startJoints[index]) * alpha);
        }

        return row;
    }

    internal static double InterpolateCapturedAlpha(double normalizedTime)
    {
        var clamped = Math.Clamp(normalizedTime, 0.0, 1.0);
        var scaledIndex = clamped * (CapturedMvpointAlpha.Length - 1);
        var lower = (int)Math.Floor(scaledIndex);
        var upper = Math.Min(lower + 1, CapturedMvpointAlpha.Length - 1);
        var fraction = scaledIndex - lower;

        return CapturedMvpointAlpha[lower]
            + ((CapturedMvpointAlpha[upper] - CapturedMvpointAlpha[lower]) * fraction);
    }
}