using System;
using System.Collections.Generic;
using EMPlannerVerificationHost;
using MultiWheelC.TrajectoryPlanning.CoarsePath;
using MultiWheelC.TrajectoryPlanning.PathSmoothing;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

internal static class LongitudinalModelChecks
{
    public static void Run()
    {
        VerifiesFinitePathSIndexedSpeedEnvelope();
        VerifiesStoppingEnvelopeIsRefinedOnActualPathS();
        VerifiesStoppingEnvelopeUsesDiscreteTimeTailStations();
        VerifiesStoppingPrecheckBeforeQpAssembly();
        VerifiesReferenceHorizonSelectionKeepsTheCurrentSegmentBoundary();
        VerifiesTimeKnotLayoutDynamicsObjectiveAndHardConstraints();
    }

    private static void VerifiesFinitePathSIndexedSpeedEnvelope()
    {
        LateralPath directionPath = CreatePath(new[]
        {
            new PathFixture(0d, 0d, 0d, 0d),
            new PathFixture(1d, 1d, 0d, 0d),
        });
        var directionInput = new LongitudinalPlanningInput(directionPath, TravelDirection.Forward, 0d, 0d,
            EmTerminalType.Goal, EmPlannerConfiguration.CreateDefault(), Array.Empty<double>(), Array.Empty<double>());
        EmPlanningStatus directionStatus = new PathSpeedLimitBuilder().Build(directionInput,
            out PathSpeedLimit directionEnvelope, out string directionFailureReason);
        Verification.Equal(EmPlanningStatus.Success, directionStatus, "default direction speed limit status: " +
            directionFailureReason);
        Verification.NearlyEqual(0.20d, directionEnvelope.DirectionMaximumSpeedMetersPerSecond,
            "default direction speed limit");

        EmPlannerConfiguration configuration = EmPlannerConfiguration.CreateDefault();
        configuration.Longitudinal.MaximumForwardSpeedMetersPerSecond = 1d;
        configuration.Longitudinal.MaximumReverseSpeedMetersPerSecond = 1d;
        LateralPath path = CreatePath(new[]
        {
            new PathFixture(10d, 0d, 0d, 0d),
            new PathFixture(20d, 2d, 2d, 4d),
            new PathFixture(20.5d, 4d, 20d, 0d),
            new PathFixture(21d, 5d, 0d, 0d),
        });
        var input = new LongitudinalPlanningInput(path, TravelDirection.Forward, 0d, 0d,
            EmTerminalType.Goal, configuration, Array.Empty<double>(), Array.Empty<double>());

        EmPlanningStatus status = new PathSpeedLimitBuilder().Build(input, out PathSpeedLimit envelope,
            out string failureReason);
        Verification.Equal(EmPlanningStatus.Success, status, "speed envelope status: " + failureReason);
        Verification.NearlyEqual(1d, envelope.DirectionMaximumSpeedMetersPerSecond, "overridden direction speed limit");
        Verification.NearlyEqual(Math.Sqrt(0.20d / 2d), envelope.LateralAccelerationLimitAt(2d),
            "curvature lateral-acceleration limit");
        Verification.NearlyEqual(0.50d / 4d, envelope.CurvatureRateLimitAt(2d), "curvature-rate limit");
        Verification.True(double.IsFinite(envelope.LateralAccelerationLimitAt(0d)) &&
            double.IsFinite(envelope.CurvatureRateLimitAt(0d)), "zero curvature limits stay finite");
        Verification.NearlyEqual(Math.Sqrt(2d * 0.30d * (5d - 4d)), envelope.StoppingLimitAt(4d),
            "stopping speed limit");
        Verification.NearlyEqual(Math.Sqrt(0.20d / 20d), envelope.MaximumSpeedAt(4d),
            "combined limit chooses the finite minimum");
        double interpolationQueryPathS = 1.013d;
        int upperStation = 1;
        while (envelope.PathS[upperStation] < interpolationQueryPathS)
            upperStation++;
        double lowerPathS = envelope.PathS[upperStation - 1];
        double upperPathS = envelope.PathS[upperStation];
        double fraction = (interpolationQueryPathS - lowerPathS) / (upperPathS - lowerPathS);
        double expectedInterpolatedSpeed = envelope.MaximumSpeedMetersPerSecond[upperStation - 1] +
            (envelope.MaximumSpeedMetersPerSecond[upperStation] - envelope.MaximumSpeedMetersPerSecond[upperStation - 1]) *
            fraction;
        Verification.NearlyEqual(expectedInterpolatedSpeed, envelope.MaximumSpeedAt(interpolationQueryPathS),
            "speed envelope interpolates by PathS rather than ReferenceS");
        Verification.NearlyEqual(0d, envelope.MaximumSpeedAt(5d), "terminal speed is exactly zero");
    }

    private static void VerifiesStoppingEnvelopeIsRefinedOnActualPathS()
    {
        EmPlannerConfiguration configuration = EmPlannerConfiguration.CreateDefault();
        configuration.Longitudinal.MaximumForwardSpeedMetersPerSecond = 1d;
        configuration.Longitudinal.MaximumReverseSpeedMetersPerSecond = 1d;
        LateralPath path = CreatePath(new[]
        {
            new PathFixture(0d, 0d, 0d, 0d),
            new PathFixture(100d, 2d, 0d, 0d),
        });
        var input = new LongitudinalPlanningInput(path, TravelDirection.Forward, 0d, 0d,
            EmTerminalType.Goal, configuration, Array.Empty<double>(), Array.Empty<double>());

        EmPlanningStatus status = new PathSpeedLimitBuilder().Build(input, out PathSpeedLimit envelope,
            out string failureReason);
        Verification.Equal(EmPlanningStatus.Success, status, "refined stopping envelope status: " + failureReason);
        Verification.True(envelope.PathS.Count > path.Points.Count, "stopping envelope inserts actual-PathS refinement stations");
        Verification.NearlyEqual(Math.Sqrt(2d * 0.30d * (2d - 1.5d)), envelope.MaximumSpeedAt(1.5d),
            "refined stopping envelope avoids a sparse terminal chord");
    }

    private static void VerifiesStoppingEnvelopeUsesDiscreteTimeTailStations()
    {
        EmPlannerConfiguration configuration = EmPlannerConfiguration.CreateDefault();
        configuration.Longitudinal.MaximumForwardSpeedMetersPerSecond = 1d;
        configuration.Longitudinal.MaximumReverseSpeedMetersPerSecond = 1d;
        LateralPath path = CreatePath(new[]
        {
            new PathFixture(0d, 0d, 0d, 0d),
            new PathFixture(100d, 2d, 0d, 0d),
        });
        var input = new LongitudinalPlanningInput(path, TravelDirection.Forward, 0d, 0d,
            EmTerminalType.Goal, configuration, Array.Empty<double>(), Array.Empty<double>());

        EmPlanningStatus status = new PathSpeedLimitBuilder().Build(input, out PathSpeedLimit envelope,
            out string failureReason);
        Verification.Equal(EmPlanningStatus.Success, status, "discrete stopping-tail status: " + failureReason);
        double timeStep = configuration.Scheduling.OutputTimeStepSeconds;
        double deceleration = configuration.Longitudinal.MaximumDecelerationMetersPerSecondSquared;
        double firstTailDistance = 0.5d * deceleration * timeStep * timeStep;
        double firstTailPathS = input.TerminalPathS - firstTailDistance;
        Verification.NearlyEqual(deceleration * timeStep, envelope.StoppingLimitAt(firstTailPathS),
            "first stopping-tail station matches one discrete deceleration time step");
        double jerk = configuration.Longitudinal.MaximumJerkMetersPerSecondCubed;
        double firstJerkTailDistance = jerk * timeStep * timeStep * timeStep / 6d;
        double firstJerkTailPathS = input.TerminalPathS - firstJerkTailDistance;
        Verification.NearlyEqual(Math.Sqrt(2d * deceleration * firstJerkTailDistance),
            envelope.StoppingLimitAt(firstJerkTailPathS),
            "first stopping-tail station matches one discrete jerk-release time step");
    }

    private static void VerifiesStoppingPrecheckBeforeQpAssembly()
    {
        EmPlannerConfiguration configuration = EmPlannerConfiguration.CreateDefault();
        LateralPath shortPath = CreatePath(new[]
        {
            new PathFixture(0d, 0d, 0d, 0d),
            new PathFixture(100d, 0.01d, 0d, 0d),
        });
        var input = new LongitudinalPlanningInput(shortPath, TravelDirection.Forward, 0.20d, 0.20d,
            EmTerminalType.RollingSafetyStop, configuration, Array.Empty<double>(), Array.Empty<double>());

        EmPlanningStatus status = new PathSpeedLimitBuilder().Build(input, out PathSpeedLimit envelope,
            out string failureReason);
        Verification.Equal(EmPlanningStatus.StoppingDistanceInsufficient, status,
            "jerk/deceleration stopping precheck status");
        Verification.True(envelope == null, "stopping-distance failure does not create a speed envelope");
        Verification.True(failureReason.Length != 0, "stopping-distance failure explains the rejection");
    }

    private static void VerifiesReferenceHorizonSelectionKeepsTheCurrentSegmentBoundary()
    {
        EmPlannerConfiguration configuration = EmPlannerConfiguration.CreateDefault();
        DirectionSegmentView shortGoal = CreateSegment(0.25d, EmBoundaryType.Goal);
        EmPlanningStatus status = new PlanningHorizonSelector().Select(shortGoal, 0d, 0d, 0d, configuration,
            out PlanningHorizonSelection goalSelection, out string goalFailure);

        Verification.Equal(EmPlanningStatus.Success, status, "goal horizon status: " + goalFailure);
        Verification.Equal(EmTerminalType.Goal, goalSelection.TerminalType, "goal terminal type");
        Verification.NearlyEqual(0.25d, goalSelection.TerminalReferenceS, "goal terminal reference S");

        DirectionSegmentView longSegment = CreateSegment(4d, EmBoundaryType.GearSwitchApproach);
        status = new PlanningHorizonSelector().Select(longSegment, 0d, 0d, 0d, configuration,
            out PlanningHorizonSelection rollingSelection, out string rollingFailure);
        Verification.Equal(EmPlanningStatus.Success, status, "rolling horizon status: " + rollingFailure);
        Verification.Equal(EmTerminalType.RollingSafetyStop, rollingSelection.TerminalType, "rolling terminal type");
        Verification.True(rollingSelection.TerminalReferenceS >= 0d && rollingSelection.TerminalReferenceS < 4d,
            "rolling horizon remains within the current segment");
    }

    private static void VerifiesTimeKnotLayoutDynamicsObjectiveAndHardConstraints()
    {
        var layout = new LongitudinalVariableLayout(5);
        Verification.Equal(19, layout.VariableCount, "ST variable count");
        for (int index = 0; index < 5; index++)
        {
            Verification.Equal(index, layout.S(index), "s index " + index);
            Verification.Equal(5 + index, layout.U(index), "u index " + index);
            Verification.Equal(10 + index, layout.A(index), "a index " + index);
        }
        for (int index = 0; index < 4; index++)
            Verification.Equal(15 + index, layout.J(index), "j index " + index);

        double[] times = { 0d, 0.05d, 0.10d, 0.15d, 0.20d };
        double[] jerk = { 0.30d, -0.10d, 0.20d, -0.20d };
        LongitudinalCandidate integrated = LongitudinalCandidate.Integrate(times, 0d, 0.10d, 0.02d, jerk);
        for (int index = 0; index < jerk.Length; index++)
        {
            double dt = times[index + 1] - times[index];
            Verification.NearlyEqual(integrated.A[index] + dt * integrated.J[index], integrated.A[index + 1],
                "exact ST acceleration dynamics " + index);
            Verification.NearlyEqual(integrated.U[index] + dt * integrated.A[index] + 0.5d * dt * dt * integrated.J[index],
                integrated.U[index + 1], "exact ST speed dynamics " + index);
            Verification.NearlyEqual(integrated.S[index] + dt * integrated.U[index] +
                0.5d * dt * dt * integrated.A[index] + dt * dt * dt * integrated.J[index] / 6d,
                integrated.S[index + 1], "exact ST progress dynamics " + index);
        }
        Verification.True(integrated.SatisfiesExactDiscreteDynamics(1e-12d), "integrated ST candidate validates dynamics");

        EmPlannerConfiguration configuration = CreateUnitScaleConfiguration();
        LateralPath path = CreatePath(new[]
        {
            new PathFixture(0d, 0d, 0d, 0d),
            new PathFixture(1d, 1d, 0d, 0d),
            new PathFixture(2d, 2d, 0d, 0d),
        });
        var input = new LongitudinalPlanningInput(path, TravelDirection.Forward, 0.10d, 0.02d,
            EmTerminalType.Goal, configuration, new[] { 0d, 0.10d, 0.20d, 0.30d, 0.40d },
            new[] { 0.20d, 0.20d, 0.20d, 0.20d, 0.20d });
        EmPlanningStatus speedStatus = new PathSpeedLimitBuilder().Build(input, out PathSpeedLimit envelope,
            out string speedFailure);
        Verification.Equal(EmPlanningStatus.Success, speedStatus, "unit-scale speed envelope: " + speedFailure);

        Verification.True(new LongitudinalConstraintBuilder(new LongitudinalObjectiveBuilder()).TryBuild(input, envelope,
            integrated, out QuadraticProgram problem, out string failureReason), "ST QP builds: " + failureReason);
        Verification.NearlyEqual(30d, MatrixValue(problem.UpperTriangularP, layout.U(0), layout.U(0)),
            "normalized speed and previous-U P coefficient");
        Verification.NearlyEqual(2d, MatrixValue(problem.UpperTriangularP, layout.A(0), layout.A(0)),
            "normalized acceleration P coefficient");
        Verification.NearlyEqual(20d, MatrixValue(problem.UpperTriangularP, layout.J(0), layout.J(0)),
            "normalized jerk P coefficient");
        Verification.NearlyEqual(2.5d, MatrixValue(problem.UpperTriangularP, layout.S(0), layout.S(0)),
            "normalized previous-S P coefficient");
        Verification.NearlyEqual(0d, MatrixValue(problem.UpperTriangularP, layout.S(4), layout.S(4)),
            "fixed terminal S has no progress-reward coefficient");

        FindSingleVariableBounds(problem, layout.S(0), out double sLower, out double sUpper);
        Verification.NearlyEqual(0d, sLower, "S lower bound");
        Verification.NearlyEqual(2d, sUpper, "S upper bound");
        FindSingleVariableBounds(problem, layout.U(1), out double uLower, out double uUpper);
        Verification.NearlyEqual(0d, uLower, "U nonnegative bound");
        Verification.NearlyEqual(envelope.MaximumSpeedAt(integrated.S[1]), uUpper,
            "U upper bound samples envelope at current S iterate");
        FindSingleVariableBounds(problem, layout.A(1), out double aLower, out double aUpper);
        Verification.NearlyEqual(-1d, aLower, "deceleration lower bound");
        Verification.NearlyEqual(1d, aUpper, "acceleration upper bound");
        FindSingleVariableBounds(problem, layout.J(1), out double jLower, out double jUpper);
        Verification.NearlyEqual(-1d, jLower, "jerk lower bound");
        Verification.NearlyEqual(1d, jUpper, "jerk upper bound");

        Verification.Equal(1, CountExactEqualityRows(problem, new Dictionary<int, double> { { layout.S(0), 1d } }, 0d),
            "exact initial S");
        Verification.Equal(1, CountExactEqualityRows(problem, new Dictionary<int, double> { { layout.U(0), 1d } }, 0.10d),
            "exact initial U");
        Verification.Equal(1, CountExactEqualityRows(problem, new Dictionary<int, double> { { layout.A(0), 1d } }, 0.02d),
            "exact initial A");
        Verification.Equal(1, CountExactEqualityRows(problem, new Dictionary<int, double> { { layout.S(4), 1d } }, 2d),
            "exact terminal S");
        Verification.Equal(1, CountExactEqualityRows(problem, new Dictionary<int, double> { { layout.U(4), 1d } }, 0d),
            "exact terminal U");
        Verification.Equal(1, CountBoundedRow(problem, new Dictionary<int, double>
        {
            { layout.S(1), 1d }, { layout.S(0), -1d },
        }, 0d, QuadraticProgram.MaximumFiniteBound), "monotonic S hard constraint");
        Verification.Equal(1, CountExactEqualityRows(problem, new Dictionary<int, double>
        {
            { layout.A(1), 1d }, { layout.A(0), -1d }, { layout.J(0), -0.05d },
        }, 0d), "exact ST acceleration equation");
        Verification.Equal(1, CountExactEqualityRows(problem, new Dictionary<int, double>
        {
            { layout.U(1), 1d }, { layout.U(0), -1d }, { layout.A(0), -0.05d }, { layout.J(0), -0.00125d },
        }, 0d), "exact ST speed equation");
        Verification.Equal(1, CountExactEqualityRows(problem, new Dictionary<int, double>
        {
            { layout.S(1), 1d }, { layout.S(0), -1d }, { layout.U(0), -0.05d }, { layout.A(0), -0.00125d },
            { layout.J(0), -0.000020833333333333333d },
        }, 0d), "exact ST progress equation");
    }

    private static LateralPath CreatePath(IReadOnlyList<PathFixture> fixtures)
    {
        var points = new List<LateralPathPoint>(fixtures.Count);
        for (int index = 0; index < fixtures.Count; index++)
        {
            PathFixture fixture = fixtures[index];
            points.Add(new LateralPathPoint(fixture.ReferenceS, fixture.PathS, 0d, 0d, 0d, 0d, fixture.PathS, 0d,
                0d, fixture.Curvature, fixture.Curvature, fixture.CurvatureDerivative));
        }
        return new LateralPath(points, true);
    }

    private static DirectionSegmentView CreateSegment(double length, EmBoundaryType endBoundaryType)
    {
        var points = new List<SmoothedPathPoint>
        {
            Point(0d, 0d),
            Point(length, length),
        };
        return new DirectionSegmentView(0, TravelDirection.Forward, points,
            new ReferenceBoundary(0, 0d, EmBoundaryType.None, 0d),
            new ReferenceBoundary(0, length, endBoundaryType, length), 0d);
    }

    private static SmoothedPathPoint Point(double x, double s)
    {
        return new SmoothedPathPoint(x, 0d, 0d, 0d, s, TravelDirection.Forward, 0d, 0d, 0d, 1d, false,
            SmoothedPathPointSource.Anchor);
    }

    private static EmPlannerConfiguration CreateUnitScaleConfiguration()
    {
        EmPlannerConfiguration configuration = EmPlannerConfiguration.CreateDefault();
        configuration.Scheduling.TimeHorizonSeconds = 0.20d;
        configuration.Scheduling.OutputTimeStepSeconds = 0.05d;
        configuration.Longitudinal.MaximumForwardSpeedMetersPerSecond = 1d;
        configuration.Longitudinal.MaximumReverseSpeedMetersPerSecond = 1d;
        configuration.Longitudinal.MaximumAccelerationMetersPerSecondSquared = 1d;
        configuration.Longitudinal.MaximumDecelerationMetersPerSecondSquared = 1d;
        configuration.Longitudinal.MaximumJerkMetersPerSecondCubed = 1d;
        configuration.Longitudinal.MaximumLateralAccelerationMetersPerSecondSquared = 1d;
        configuration.Longitudinal.MaximumCurvatureRatePerMeterPerSecond = 1d;
        return configuration;
    }

    private static double MatrixValue(SparseCscMatrix matrix, int row, int column)
    {
        for (int index = matrix.ColumnPointers[column]; index < matrix.ColumnPointers[column + 1]; index++)
        {
            if (matrix.RowIndices[index] == row)
                return matrix.Values[index];
        }
        return 0d;
    }

    private static void FindSingleVariableBounds(QuadraticProgram problem, int variable, out double lower, out double upper)
    {
        for (int row = 0; row < problem.ConstraintCount; row++)
        {
            if (RowMatches(problem.ConstraintMatrix, row, new Dictionary<int, double> { { variable, 1d } }))
            {
                lower = problem.LowerBounds[row];
                upper = problem.UpperBounds[row];
                return;
            }
        }
        throw new InvalidOperationException("No single-variable bounds were found for variable " + variable + ".");
    }

    private static int CountExactEqualityRows(QuadraticProgram problem, IReadOnlyDictionary<int, double> expected,
        double bound)
    {
        return CountBoundedRow(problem, expected, bound, bound);
    }

    private static int CountBoundedRow(QuadraticProgram problem, IReadOnlyDictionary<int, double> expected,
        double lower, double upper)
    {
        int count = 0;
        for (int row = 0; row < problem.ConstraintCount; row++)
        {
            if (Math.Abs(problem.LowerBounds[row] - lower) <= 1e-12d &&
                Math.Abs(problem.UpperBounds[row] - upper) <= 1e-12d && RowMatches(problem.ConstraintMatrix, row, expected))
            {
                count++;
            }
        }
        return count;
    }

    private static bool RowMatches(SparseCscMatrix matrix, int targetRow, IReadOnlyDictionary<int, double> expected)
    {
        var actual = new Dictionary<int, double>();
        for (int column = 0; column < matrix.ColumnCount; column++)
        {
            for (int index = matrix.ColumnPointers[column]; index < matrix.ColumnPointers[column + 1]; index++)
            {
                if (matrix.RowIndices[index] == targetRow)
                    actual[column] = matrix.Values[index];
            }
        }
        if (actual.Count != expected.Count)
            return false;
        foreach (KeyValuePair<int, double> pair in expected)
        {
            if (!actual.TryGetValue(pair.Key, out double actualValue) || Math.Abs(actualValue - pair.Value) > 1e-12d)
                return false;
        }
        return true;
    }

    private sealed class PathFixture
    {
        public PathFixture(double referenceS, double pathS, double curvature, double curvatureDerivative)
        {
            ReferenceS = referenceS;
            PathS = pathS;
            Curvature = curvature;
            CurvatureDerivative = curvatureDerivative;
        }

        public double ReferenceS { get; }
        public double PathS { get; }
        public double Curvature { get; }
        public double CurvatureDerivative { get; }
    }
}