using System;
using System.Collections.Generic;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

/// <summary>Builds one normalized ST QP with exact constant-jerk integration and hard terminal conditions.</summary>
public sealed class LongitudinalConstraintBuilder
{
    private readonly LongitudinalObjectiveBuilder _objectiveBuilder;

    public LongitudinalConstraintBuilder(LongitudinalObjectiveBuilder objectiveBuilder)
    {
        _objectiveBuilder = objectiveBuilder ?? throw new ArgumentNullException(nameof(objectiveBuilder));
    }

    public bool TryBuild(LongitudinalPlanningInput input, PathSpeedLimit speedLimit, LongitudinalCandidate iterate,
        out QuadraticProgram problem, out string failureReason)
    {
        problem = null;
        failureReason = string.Empty;
        try
        {
            if (input == null || speedLimit == null || iterate == null)
                throw new ArgumentException("ST input, speed envelope, and iterate are required.");
            if (Math.Abs(speedLimit.TerminalPathS - input.TerminalPathS) > 1e-12d)
                throw new ArgumentException("The speed envelope terminal must match actual lateral PathS.");

            IReadOnlyList<double> expectedTimes = LongitudinalCandidate.CreateKnotTimes(
                input.Configuration.Scheduling.TimeHorizonSeconds, input.Configuration.Scheduling.OutputTimeStepSeconds);
            if (!HasMatchingTimes(iterate.KnotTimes, expectedTimes))
                throw new ArgumentException("The ST iterate time knots do not match the configured horizon.");
            var layout = new LongitudinalVariableLayout(expectedTimes.Count);
            if (iterate.S.Count != layout.KnotCount || iterate.U.Count != layout.KnotCount ||
                iterate.A.Count != layout.KnotCount || iterate.J.Count != layout.KnotCount - 1)
            {
                throw new ArgumentException("The ST iterate does not match the configured knot layout.");
            }
            if (!PathSpeedLimitBuilder.TryGetLimits(input, out double directionMaximum, out double maximumAcceleration,
                    out double maximumDeceleration, out double maximumJerk, out _, out _, out failureReason))
            {
                return false;
            }
            if (input.InitialProgressSpeedMetersPerSecond > speedLimit.MaximumSpeedAt(0d) + 1e-12d ||
                input.InitialAccelerationMetersPerSecondSquared < -maximumDeceleration - 1e-12d ||
                input.InitialAccelerationMetersPerSecondSquared > maximumAcceleration + 1e-12d)
            {
                failureReason = "The initial ST state violates hard bounds.";
                return false;
            }

            var hessian = new SparseTripletBuilder(layout.VariableCount, layout.VariableCount, true);
            var linearCost = new double[layout.VariableCount];
            _objectiveBuilder.AddTerms(input, speedLimit, layout, iterate, hessian, linearCost);
            var constraints = new SparseTripletBuilder(8 * layout.KnotCount, layout.VariableCount);
            var lower = new List<double>(8 * layout.KnotCount);
            var upper = new List<double>(8 * layout.KnotCount);
            int row = 0;
            AddVariableBounds(input, speedLimit, iterate, layout, maximumAcceleration, maximumDeceleration, maximumJerk,
                constraints, lower, upper, ref row);
            AddMonotonicProgress(layout, constraints, lower, upper, ref row);
            AddExactDynamics(expectedTimes, layout, constraints, lower, upper, ref row);
            AddExactStartAndTerminal(input, layout, constraints, lower, upper, ref row);
            if (row != 8 * layout.KnotCount)
                throw new InvalidOperationException("ST constraint row accounting is inconsistent.");
            problem = new QuadraticProgram(hessian.Build(), linearCost, constraints.Build(), lower, upper);
            return true;
        }
        catch (ArgumentException exception)
        {
            failureReason = exception.Message;
            return false;
        }
    }

    private static void AddVariableBounds(LongitudinalPlanningInput input, PathSpeedLimit speedLimit,
        LongitudinalCandidate iterate, LongitudinalVariableLayout layout, double maximumAcceleration,
        double maximumDeceleration, double maximumJerk, SparseTripletBuilder constraints, IList<double> lower,
        IList<double> upper, ref int row)
    {
        for (int index = 0; index < layout.KnotCount; index++)
        {
            if (iterate.S[index] < 0d || iterate.S[index] > input.TerminalPathS)
                throw new ArgumentException("The ST iterate progress lies outside actual PathS bounds.");
            AddSingleVariableRow(constraints, lower, upper, layout.S(index), 0d, input.TerminalPathS, ref row);
            AddSingleVariableRow(constraints, lower, upper, layout.U(index), 0d,
                Math.Min(input.DirectionMaximumSpeedMetersPerSecond, speedLimit.MaximumSpeedAt(iterate.S[index])), ref row);
            AddSingleVariableRow(constraints, lower, upper, layout.A(index), -maximumDeceleration, maximumAcceleration,
                ref row);
        }
        for (int index = 0; index < layout.KnotCount - 1; index++)
            AddSingleVariableRow(constraints, lower, upper, layout.J(index), -maximumJerk, maximumJerk, ref row);
    }

    private static void AddMonotonicProgress(LongitudinalVariableLayout layout, SparseTripletBuilder constraints,
        IList<double> lower, IList<double> upper, ref int row)
    {
        for (int index = 0; index < layout.KnotCount - 1; index++)
        {
            AddRow(constraints, lower, upper, row, new[]
            {
                new Coefficient(layout.S(index + 1), 1d), new Coefficient(layout.S(index), -1d),
            }, 0d, QuadraticProgram.MaximumFiniteBound);
            row++;
        }
    }

    private static void AddExactDynamics(IReadOnlyList<double> times, LongitudinalVariableLayout layout,
        SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, ref int row)
    {
        for (int index = 0; index < layout.KnotCount - 1; index++)
        {
            double dt = times[index + 1] - times[index];
            AddRow(constraints, lower, upper, row, new[]
            {
                new Coefficient(layout.A(index + 1), 1d), new Coefficient(layout.A(index), -1d),
                new Coefficient(layout.J(index), -dt),
            }, 0d, 0d);
            row++;
            AddRow(constraints, lower, upper, row, new[]
            {
                new Coefficient(layout.U(index + 1), 1d), new Coefficient(layout.U(index), -1d),
                new Coefficient(layout.A(index), -dt), new Coefficient(layout.J(index), -0.5d * dt * dt),
            }, 0d, 0d);
            row++;
            AddRow(constraints, lower, upper, row, new[]
            {
                new Coefficient(layout.S(index + 1), 1d), new Coefficient(layout.S(index), -1d),
                new Coefficient(layout.U(index), -dt), new Coefficient(layout.A(index), -0.5d * dt * dt),
                new Coefficient(layout.J(index), -dt * dt * dt / 6d),
            }, 0d, 0d);
            row++;
        }
    }

    private static void AddExactStartAndTerminal(LongitudinalPlanningInput input, LongitudinalVariableLayout layout,
        SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, ref int row)
    {
        AddSingleVariableRow(constraints, lower, upper, layout.S(0), 0d, 0d, ref row);
        AddSingleVariableRow(constraints, lower, upper, layout.U(0), input.InitialProgressSpeedMetersPerSecond,
            input.InitialProgressSpeedMetersPerSecond, ref row);
        AddSingleVariableRow(constraints, lower, upper, layout.A(0), input.InitialAccelerationMetersPerSecondSquared,
            input.InitialAccelerationMetersPerSecondSquared, ref row);
        AddSingleVariableRow(constraints, lower, upper, layout.S(layout.KnotCount - 1), input.TerminalPathS,
            input.TerminalPathS, ref row);
        AddSingleVariableRow(constraints, lower, upper, layout.U(layout.KnotCount - 1), 0d, 0d, ref row);
    }

    private static void AddSingleVariableRow(SparseTripletBuilder constraints, IList<double> lower, IList<double> upper,
        int variable, double minimum, double maximum, ref int row)
    {
        AddRow(constraints, lower, upper, row, new[] { new Coefficient(variable, 1d) }, minimum, maximum);
        row++;
    }

    private static void AddRow(SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, int row,
        IReadOnlyList<Coefficient> coefficients, double minimum, double maximum)
    {
        for (int index = 0; index < coefficients.Count; index++)
            constraints.Add(row, coefficients[index].Variable, coefficients[index].Value);
        lower.Add(minimum);
        upper.Add(maximum);
    }

    private static bool HasMatchingTimes(IReadOnlyList<double> actual, IReadOnlyList<double> expected)
    {
        if (actual.Count != expected.Count)
            return false;
        for (int index = 0; index < expected.Count; index++)
        {
            if (Math.Abs(actual[index] - expected[index]) > 1e-12d)
                return false;
        }
        return true;
    }

    private readonly struct Coefficient
    {
        public Coefficient(int variable, double value)
        {
            Variable = variable;
            Value = value;
        }

        public int Variable { get; }

        public double Value { get; }
    }
}