using System;
using System.Collections.Generic;
using System.Globalization;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

/// <summary>Builds one normalized ST QP with exact constant-jerk integration and mode-specific stop 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)
    {
        return TryBuildCore(input, speedLimit, iterate, false, null, null, 0d, out problem, out failureReason);
    }

    internal bool TryBuildTrusted(LongitudinalPlanningInput input, PathSpeedLimit speedLimit,
        LongitudinalCandidate anchor, LongitudinalEnvelopeTrustRegion trustRegion, double strictTolerance,
        out QuadraticProgram problem, out string failureReason)
    {
        return TryBuildCore(input, speedLimit, anchor, false, trustRegion, anchor, strictTolerance,
            out problem, out failureReason);
    }

    /// <summary>Builds the bounded full-scope feasibility projection before objective optimization.</summary>
    public bool TryBuildInitialFeasibilityProjection(LongitudinalPlanningInput input, PathSpeedLimit speedLimit,
        out QuadraticProgram problem, out string failureReason)
    {
        return TryBuildInitialFeasibilityProjection(input, speedLimit,
            input == null ? null : CreateScheduleReferenceIterate(input), out problem, out failureReason);
    }

    public bool TryBuildInitialFeasibilityProjection(LongitudinalPlanningInput input, PathSpeedLimit speedLimit,
        LongitudinalCandidate linearizationIterate, out QuadraticProgram problem, out string failureReason)
    {
        problem = null;
        failureReason = string.Empty;
        if (input == null || input.PlanningScope != EmPlanningScope.FullDirectionSegment ||
            input.Mode != EmLongitudinalMode.ExactStopAtBoundary || linearizationIterate == null)
        {
            failureReason = "Initial feasibility projection is only defined for full-direction exact-stop planning.";
            return false;
        }
        return TryBuildCore(input, speedLimit, linearizationIterate, true, null, null, 0d,
            out problem, out failureReason);
    }

    private bool TryBuildCore(LongitudinalPlanningInput input, PathSpeedLimit speedLimit, LongitudinalCandidate iterate,
        bool useScheduleReferenceObjective, LongitudinalEnvelopeTrustRegion trustRegion,
        LongitudinalCandidate trustedAnchor, double strictTolerance, 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.PathUpperBoundS - input.PathUpperBoundS) > 1e-12d)
                throw new ArgumentException("The speed envelope upper bound must match actual lateral PathS.");

            IReadOnlyList<double> expectedTimes = input.KnotSchedule.KnotTimes;
            if (!HasMatchingTimes(iterate.KnotTimes, expectedTimes))
                throw new ArgumentException("The ST iterate time knots do not match the supplied knot schedule.");
            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 ((trustRegion == null) != (trustedAnchor == null))
                throw new ArgumentException("Trusted QP construction requires both a trust region and anchor.");
            if (trustRegion != null && (trustRegion.MinimumPathS.Count != layout.KnotCount ||
                trustRegion.MaximumPathS.Count != layout.KnotCount ||
                trustRegion.SpeedSlope.Count != layout.KnotCount ||
                trustRegion.SpeedIntercept.Count != layout.KnotCount))
            {
                throw new ArgumentException("The trust region 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 > directionMaximum + 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];
            if (useScheduleReferenceObjective)
                AddInitialFeasibilityObjective(input, layout, hessian, linearCost);
            else
                _objectiveBuilder.AddTerms(input, speedLimit, layout, iterate, hessian, linearCost);
            int stabilizationStart = GetStabilizationStart(input, expectedTimes, layout.KnotCount);
            int stationaryKnotCount = layout.KnotCount - stabilizationStart;
            int expectedRows = 10 * layout.KnotCount - 3 + 3 * stationaryKnotCount;
            var constraints = new SparseTripletBuilder(expectedRows, layout.VariableCount);
            var lower = new List<double>(expectedRows);
            var upper = new List<double>(expectedRows);
            int row = 0;
            AddVariableBounds(input, speedLimit, iterate, trustRegion, layout, maximumAcceleration,
                maximumDeceleration, maximumJerk, constraints, lower, upper, ref row);
            AddLowSpeedDecelerationReleaseEnvelope(layout, iterate, maximumJerk,
                constraints, lower, upper, ref row);
            AddMonotonicProgress(layout, constraints, lower, upper, ref row);
            AddExactDynamics(expectedTimes, layout, constraints, lower, upper, ref row);
            AddExactStart(input, layout, constraints, lower, upper, ref row);
            if (input.Mode == EmLongitudinalMode.ExactStopAtBoundary)
                AddExactStopTail(input, layout, stabilizationStart, constraints, lower, upper, ref row);
            if (row != expectedRows)
                throw new InvalidOperationException("ST constraint row accounting is inconsistent.");
            problem = new QuadraticProgram(hessian.Build(), linearCost, constraints.Build(), lower, upper);
            if (trustedAnchor != null)
            {
                LongitudinalQpAuditResult audit = LongitudinalQpFeasibilityAudit.Evaluate(problem, trustedAnchor,
                    strictTolerance, layout, stabilizationStart);
                if (!audit.IsFeasible)
                {
                    problem = null;
                    failureReason = "Planner invariant failure: strict anchor is outside trusted QP" +
                        ";row=" + audit.WorstRow + ";category=" + audit.Category +
                        ";residual=" + audit.MaximumResidual.ToString("R", CultureInfo.InvariantCulture) +
                        audit.Unit + ";tolerance=" + strictTolerance.ToString("R", CultureInfo.InvariantCulture);
                    return false;
                }
            }
            return true;
        }
        catch (ArgumentException exception)
        {
            failureReason = exception.Message;
            return false;
        }
    }

    private static LongitudinalCandidate CreateScheduleReferenceIterate(LongitudinalPlanningInput input)
    {
        int knotCount = input.KnotSchedule.KnotTimes.Count;
        return new LongitudinalCandidate(input.KnotSchedule.KnotTimes, input.KnotSchedule.ReferencePathS,
            input.KnotSchedule.ReferenceSpeedMetersPerSecond, new double[knotCount], new double[knotCount - 1]);
    }

    private static void AddInitialFeasibilityObjective(LongitudinalPlanningInput input, LongitudinalVariableLayout layout,
        SparseTripletBuilder hessian, IList<double> linearCost)
    {
        double progressScale = 1d;
        double speedScale = 1d;
        double accelerationScale = 1d;
        double jerkScale = 1d;
        for (int index = 0; index < layout.KnotCount; index++)
        {
            AddProjectionSquaredResidual(hessian, linearCost, layout.S(index), input.KnotSchedule.ReferencePathS[index],
                1d, progressScale);
            AddProjectionSquaredResidual(hessian, linearCost, layout.U(index),
                input.KnotSchedule.ReferenceSpeedMetersPerSecond[index], 10d, speedScale);
            AddProjectionSquaredResidual(hessian, linearCost, layout.A(index), 0d, 1e-3d, accelerationScale);
        }
        for (int index = 0; index < layout.KnotCount - 1; index++)
            AddProjectionSquaredResidual(hessian, linearCost, layout.J(index), 0d, 1e-3d, jerkScale);
    }

    private static void AddProjectionSquaredResidual(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));
        double coefficient = 2d * weight / (scale * scale);
        hessian.Add(variable, variable, coefficient);
        linearCost[variable] += -coefficient * reference;
    }

    private static void AddVariableBounds(LongitudinalPlanningInput input, PathSpeedLimit speedLimit,
        LongitudinalCandidate iterate, LongitudinalEnvelopeTrustRegion trustRegion,
        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.PathUpperBoundS)
                throw new ArgumentException("The ST iterate progress lies outside actual PathS bounds.");
            if (trustRegion == null)
                AddSingleVariableRow(constraints, lower, upper, layout.S(index), 0d, input.PathUpperBoundS, ref row);
            else
                AddSingleVariableRow(constraints, lower, upper, layout.S(index),
                    trustRegion.MinimumPathS[index], trustRegion.MaximumPathS[index], ref row);
            double maximumSpeed = index == 0
                ? input.DirectionMaximumSpeedMetersPerSecond
                : input.DirectionMaximumSpeedMetersPerSecond;
            AddSingleVariableRow(constraints, lower, upper, layout.U(index), 0d, maximumSpeed, ref row);
            if (index > 0)
            {
                if (trustRegion == null)
                {
                    AddLinearizedSpeedEnvelopeRow(speedLimit, iterate.S[index], layout.S(index), layout.U(index),
                        constraints, lower, upper, ref row);
                }
                else
                {
                    AddRow(constraints, lower, upper, row, new[]
                    {
                        new Coefficient(layout.U(index), 1d),
                        new Coefficient(layout.S(index), -trustRegion.SpeedSlope[index]),
                    }, -QuadraticProgram.MaximumFiniteBound, trustRegion.SpeedIntercept[index]);
                    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 AddLinearizedSpeedEnvelopeRow(PathSpeedLimit speedLimit, double pathS, int pathSVariable,
        int speedVariable, SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, ref int row)
    {
        int segment = FindSpeedEnvelopeSegment(speedLimit, pathS);
        double startS = speedLimit.PathS[segment];
        double endS = speedLimit.PathS[segment + 1];
        double startSpeed = speedLimit.MaximumSpeedMetersPerSecond[segment];
        double endSpeed = speedLimit.MaximumSpeedMetersPerSecond[segment + 1];
        double slope = (endSpeed - startSpeed) / (endS - startS);
        double intercept = startSpeed - slope * startS;
        AddRow(constraints, lower, upper, row, new[]
        {
            new Coefficient(speedVariable, 1d), new Coefficient(pathSVariable, -slope),
        }, -QuadraticProgram.MaximumFiniteBound, intercept);
        row++;
    }

    private static int FindSpeedEnvelopeSegment(PathSpeedLimit speedLimit, double pathS)
    {
        double clamped = Math.Max(speedLimit.PathS[0], Math.Min(speedLimit.PathUpperBoundS, pathS));
        for (int index = 0; index < speedLimit.PathS.Count - 1; index++)
        {
            if (clamped <= speedLimit.PathS[index + 1])
                return index;
        }
        return speedLimit.PathS.Count - 2;
    }

    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++;
        }
    }

    internal static void CalculateLowSpeedDecelerationReleaseTangent(
        double anchorAcceleration, double maximumJerk,
        out double accelerationCoefficient, out double lowerBound)
    {
        if (!IsFinite(anchorAcceleration) || !IsFinite(maximumJerk) || maximumJerk <= 0d)
            throw new ArgumentOutOfRangeException(nameof(anchorAcceleration));
        double a0 = Math.Min(0d, anchorAcceleration);
        accelerationCoefficient = -a0 / maximumJerk;
        lowerBound = -(a0 * a0) / (2d * maximumJerk);
    }

    private static void AddLowSpeedDecelerationReleaseEnvelope(LongitudinalVariableLayout layout,
        LongitudinalCandidate iterate, double maximumJerk, SparseTripletBuilder constraints,
        IList<double> lower, IList<double> upper, ref int row)
    {
        for (int index = 0; index < layout.KnotCount; index++)
        {
            CalculateLowSpeedDecelerationReleaseTangent(iterate.A[index], maximumJerk,
                out double accelerationCoefficient, out double lowerBound);
            AddRow(constraints, lower, upper, row, new[]
            {
                new Coefficient(layout.U(index), 1d),
                new Coefficient(layout.A(index), accelerationCoefficient),
            }, lowerBound, 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 AddExactStart(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);
    }

    private static void AddExactStopTail(LongitudinalPlanningInput input, LongitudinalVariableLayout layout,
        int stabilizationStart, SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, ref int row)
    {
        for (int index = stabilizationStart; index < layout.KnotCount; index++)
        {
            AddSingleVariableRow(constraints, lower, upper, layout.S(index), input.StopBoundaryPathS,
                input.StopBoundaryPathS, ref row);
            AddSingleVariableRow(constraints, lower, upper, layout.U(index), 0d, 0d, ref row);
            AddSingleVariableRow(constraints, lower, upper, layout.A(index), 0d, 0d, ref row);
        }
    }

    private static int GetStabilizationStart(LongitudinalPlanningInput input, IReadOnlyList<double> times,
        int knotCount)
    {
        if (input.Mode != EmLongitudinalMode.ExactStopAtBoundary)
            return knotCount;
        if (input.PlanningScope == EmPlanningScope.FullDirectionSegment)
        {
            if (input.KnotSchedule.TerminalHoldStartIndex < 1 ||
                input.KnotSchedule.TerminalHoldStartIndex >= knotCount)
            {
                throw new ArgumentException("Full-direction exact-stop schedules require an explicit terminal hold boundary.");
            }
            return input.KnotSchedule.TerminalHoldStartIndex;
        }
        return LongitudinalTerminalSchedule.GetStabilizationStartIndex(times,
            input.Configuration.Scheduling.OutputTimeStepSeconds);
    }

    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 static bool IsFinite(double value)
    {
        return !double.IsNaN(value) && !double.IsInfinity(value);
    }

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

        public int Variable { get; }

        public double Value { get; }
    }
}