using System;
using System.Collections.Generic;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

/// <summary>Normalized 0.5*x'Px+q'x ST objective terms with no terminal-progress reward.</summary>
public sealed class LongitudinalObjectiveBuilder
{
    public void AddTerms(LongitudinalPlanningInput input, PathSpeedLimit speedLimit, LongitudinalVariableLayout layout,
        LongitudinalCandidate iterate, SparseTripletBuilder hessian, IList<double> linearCost)
    {
        if (input == null || speedLimit == null || layout == null || iterate == null || hessian == null || linearCost == null ||
            linearCost.Count != layout.VariableCount || iterate.KnotTimes.Count != layout.KnotCount)
        {
            throw new ArgumentException("ST objective inputs do not match the variable layout.");
        }

        LongitudinalConfiguration configuration = input.Configuration.Longitudinal;
        LongitudinalWeights weights = configuration.Weights ?? throw new ArgumentException("Longitudinal weights are required.");
        double speedScale = RequirePositive(input.DirectionMaximumSpeedMetersPerSecond, nameof(speedScale));
        double accelerationScale = RequirePositive(Math.Max(configuration.MaximumAccelerationMetersPerSecondSquared,
            configuration.MaximumDecelerationMetersPerSecondSquared), nameof(accelerationScale));
        double jerkScale = RequirePositive(configuration.MaximumJerkMetersPerSecondCubed, nameof(jerkScale));
        double progressScale = input.TerminalPathS > 0d ? input.TerminalPathS : 1d;

        for (int index = 0; index < layout.KnotCount; index++)
        {
            AddSquaredResidual(hessian, linearCost, layout.U(index), speedLimit.MaximumSpeedAt(iterate.S[index]),
                weights.ReferenceSpeed, speedScale);
            AddSquaredResidual(hessian, linearCost, layout.A(index), 0d, weights.Acceleration, accelerationScale);
            if (index < layout.KnotCount - 1 && index < input.PreviousPathS.Count)
            {
                AddSquaredResidual(hessian, linearCost, layout.S(index), input.PreviousPathS[index],
                    weights.PreviousTrajectory, progressScale);
                AddSquaredResidual(hessian, linearCost, layout.U(index), input.PreviousProgressSpeedMetersPerSecond[index],
                    weights.PreviousTrajectory, speedScale);
            }
        }
        for (int index = 0; index < layout.KnotCount - 1; index++)
            AddSquaredResidual(hessian, linearCost, layout.J(index), 0d, weights.Jerk, jerkScale);
        AddSquaredResidual(hessian, linearCost, layout.A(layout.KnotCount - 1), 0d, weights.TerminalAcceleration,
            accelerationScale);
    }

    private static void AddSquaredResidual(SparseTripletBuilder hessian, IList<double> linearCost, int variable,
        double reference, double weight, double scale)
    {
        if (!IsFinite(reference) || !IsFinite(weight) || weight < 0d || !IsFinite(scale) || scale <= 0d)
            throw new ArgumentOutOfRangeException(nameof(reference));
        if (weight == 0d)
            return;
        double coefficient = 2d * weight / (scale * scale);
        hessian.Add(variable, variable, coefficient);
        linearCost[variable] += -coefficient * reference;
    }

    private static double RequirePositive(double value, string name)
    {
        if (!IsFinite(value) || value <= 0d)
            throw new ArgumentOutOfRangeException(name);
        return value;
    }

    private static bool IsFinite(double value)
    {
        return !double.IsNaN(value) && !double.IsInfinity(value);
    }
}