using System;
using System.Collections.Generic;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

/// <summary>Assembles one lateral SQP QP with exact dynamics and finite hard bounds.</summary>
public sealed class LateralConstraintBuilder
{
    private const double Epsilon = 1e-12d;
    private readonly LateralObjectiveBuilder _objectiveBuilder;

    public LateralConstraintBuilder(LateralObjectiveBuilder objectiveBuilder)
    {
        _objectiveBuilder = objectiveBuilder ?? throw new ArgumentNullException(nameof(objectiveBuilder));
    }

    public bool TryBuild(LateralPlanningInput input, LateralCandidate linearization, out QuadraticProgram problem,
        out string failureReason)
    {
        problem = null;
        failureReason = string.Empty;
        if (input == null || linearization == null)
        {
            failureReason = "Lateral input and linearization are required.";
            return false;
        }
        if (!TryValidateCandidateStations(input, linearization, out failureReason))
            return false;

        try
        {
            var layout = new LateralVariableLayout(input.ReferenceStations.Count);
            var hessian = new SparseTripletBuilder(layout.VariableCount, layout.VariableCount, true);
            var linearCost = new double[layout.VariableCount];
            _objectiveBuilder.AddTerms(input, layout, linearization, hessian, linearCost);

            int terminalRows = input.TerminalType == EmTerminalType.RollingSafetyStop ? 0 : 2;
            var constraints = new SparseTripletBuilder(8 * layout.StationCount - 2 + terminalRows, layout.VariableCount);
            var lower = new List<double>();
            var upper = new List<double>();
            int row = 0;

            if (!TryAddLateralBounds(input, layout, linearization, constraints, lower, upper, ref row, out failureReason))
                return false;
            AddDerivativeBounds(input, layout, constraints, lower, upper, ref row);
            AddCurvatureBounds(input, layout, linearization, constraints, lower, upper, ref row);
            AddStartConstraints(input, layout, constraints, lower, upper, ref row);
            AddExactDynamics(input.ReferenceStations, layout, constraints, lower, upper, ref row);
            if (input.TerminalType != EmTerminalType.RollingSafetyStop)
                AddTerminalConstraints(layout, constraints, lower, upper, ref row);

            if (row != 8 * layout.StationCount - 2 + terminalRows)
                throw new InvalidOperationException("Lateral 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 bool TryValidateCandidateStations(LateralPlanningInput input, LateralCandidate candidate,
        out string failureReason)
    {
        failureReason = string.Empty;
        if (candidate.ReferenceStations.Count != input.ReferenceStations.Count)
        {
            failureReason = "Linearization station count does not match the lateral input.";
            return false;
        }
        for (int index = 0; index < input.ReferenceStations.Count; index++)
        {
            if (Math.Abs(candidate.ReferenceStations[index] - input.ReferenceStations[index]) > Epsilon)
            {
                failureReason = "Linearization stations do not match the lateral input.";
                return false;
            }
        }
        return true;
    }

    private static bool TryAddLateralBounds(LateralPlanningInput input, LateralVariableLayout layout,
        LateralCandidate linearization, SparseTripletBuilder constraints, IList<double> lower, IList<double> upper,
        ref int row, out string failureReason)
    {
        failureReason = string.Empty;
        double maximumOffset = RequireNonnegative(input.Configuration.Corridor.MaximumLateralOffsetMeters,
            "maximum lateral offset");
        double trustRegion = RequirePositive(input.Configuration.Lateral.MaximumLateralStepPerIterationMeters,
            "lateral trust region");
        double minimumDenominator = RequirePositive(input.Configuration.Frenet.MinimumFrenetDenominator,
            "minimum Frenet denominator");

        for (int station = 0; station < layout.StationCount; station++)
        {
            LateralInterval corridor = input.Corridor.Stations[station];
            double minimum = Math.Max(corridor.MinimumL, Math.Max(-maximumOffset, linearization.L[station] - trustRegion));
            double maximum = Math.Min(corridor.MaximumL, Math.Min(maximumOffset, linearization.L[station] + trustRegion));
            double referenceCurvature = ReferencePathInterpolator.Interpolate(input.ReferenceSegment,
                input.ReferenceStations[station]).GeometricCurvature;
            if (referenceCurvature > 0d)
                maximum = Math.Min(maximum, (1d - minimumDenominator) / referenceCurvature);
            else if (referenceCurvature < 0d)
                minimum = Math.Max(minimum, (1d - minimumDenominator) / referenceCurvature);

            if (!IsFinite(minimum) || !IsFinite(maximum) || minimum > maximum + Epsilon)
            {
                failureReason = "The lateral corridor, offset, trust-region, and Frenet denominator bounds do not intersect at station " + station + ".";
                return false;
            }
            AddSingleVariableRow(constraints, lower, upper, ref row, layout.L(station), minimum, maximum);
        }
        return true;
    }

    private static void AddDerivativeBounds(LateralPlanningInput input, LateralVariableLayout layout,
        SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, ref int row)
    {
        double slope = RequirePositive(input.Configuration.Lateral.MaximumLateralSlope, "maximum lateral slope");
        double second = RequirePositive(input.Configuration.Lateral.MaximumLateralSecondDerivativePerMeter,
            "maximum lateral second derivative");
        double third = RequirePositive(input.Configuration.Lateral.MaximumLateralThirdDerivativePerSquareMeter,
            "maximum lateral third derivative");
        for (int station = 0; station < layout.StationCount; station++)
        {
            AddSingleVariableRow(constraints, lower, upper, ref row, layout.DL(station), -slope, slope);
            AddSingleVariableRow(constraints, lower, upper, ref row, layout.DDL(station), -second, second);
        }
        for (int interval = 0; interval < layout.StationCount - 1; interval++)
            AddSingleVariableRow(constraints, lower, upper, ref row, layout.DDDL(interval), -third, third);
    }

    private static void AddCurvatureBounds(LateralPlanningInput input, LateralVariableLayout layout,
        LateralCandidate linearization, SparseTripletBuilder constraints, IList<double> lower, IList<double> upper,
        ref int row)
    {
        double maximumCurvature = LateralCurvatureLinearization.GetMaximumVehicleCurvature(input.Vehicle);
        IReadOnlyList<LateralCurvatureLinearization> affines = LateralCurvatureLinearization.Create(input, layout,
            linearization);
        for (int station = 0; station < affines.Count; station++)
        {
            LateralCurvatureLinearization affine = affines[station];
            AddRow(constraints, lower, upper, ref row, -maximumCurvature - affine.Constant,
                maximumCurvature - affine.Constant, affine.VariableIndices, affine.Gradient);
        }
    }

    private static void AddStartConstraints(LateralPlanningInput input, LateralVariableLayout layout,
        SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, ref int row)
    {
        double denominator = 1d - input.StartProjection.ReferencePoint.GeometricCurvature *
            input.StartProjection.LateralOffset;
        double startSlope = denominator * Math.Tan(input.StartProjection.HeadingError);
        if (!IsFinite(startSlope))
            throw new ArgumentException("The start lateral slope is non-finite.", nameof(input));
        AddSingleVariableRow(constraints, lower, upper, ref row, layout.L(0), input.StartProjection.LateralOffset,
            input.StartProjection.LateralOffset);
        AddSingleVariableRow(constraints, lower, upper, ref row, layout.DL(0), startSlope, startSlope);
    }

    private static void AddExactDynamics(IReadOnlyList<double> stations, LateralVariableLayout layout,
        SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, ref int row)
    {
        for (int interval = 0; interval < layout.StationCount - 1; interval++)
        {
            double ds = stations[interval + 1] - stations[interval];
            AddRow(constraints, lower, upper, ref row, 0d, 0d,
                new[] { layout.DDL(interval), layout.DDL(interval + 1), layout.DDDL(interval) },
                new[] { -1d, 1d, -ds });
            AddRow(constraints, lower, upper, ref row, 0d, 0d,
                new[] { layout.DL(interval), layout.DL(interval + 1), layout.DDL(interval), layout.DDDL(interval) },
                new[] { -1d, 1d, -ds, -0.5d * ds * ds });
            AddRow(constraints, lower, upper, ref row, 0d, 0d,
                new[] { layout.L(interval), layout.L(interval + 1), layout.DL(interval), layout.DDL(interval), layout.DDDL(interval) },
                new[] { -1d, 1d, -ds, -0.5d * ds * ds, -ds * ds * ds / 6d });
        }
    }

    private static void AddTerminalConstraints(LateralVariableLayout layout, SparseTripletBuilder constraints,
        IList<double> lower, IList<double> upper, ref int row)
    {
        AddSingleVariableRow(constraints, lower, upper, ref row, layout.L(layout.StationCount - 1), 0d, 0d);
        AddSingleVariableRow(constraints, lower, upper, ref row, layout.DL(layout.StationCount - 1), 0d, 0d);
    }

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

    private static void AddRow(SparseTripletBuilder constraints, IList<double> lower, IList<double> upper, ref int row,
        double minimum, double maximum, IReadOnlyList<int> variables, IReadOnlyList<double> coefficients)
    {
        if (!IsFinite(minimum) || !IsFinite(maximum) || minimum > maximum || variables.Count != coefficients.Count)
            throw new ArgumentException("Lateral constraint bounds are invalid.");
        for (int index = 0; index < variables.Count; index++)
            constraints.Add(row, variables[index], coefficients[index]);
        lower.Add(minimum);
        upper.Add(maximum);
        row++;
    }

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

    private static double RequireNonnegative(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);
    }
}