using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

/// <summary>Bounded ST envelope iteration retaining only independently validated physical candidates.</summary>
public sealed class SequentialLongitudinalOptimizer
{
    private const int MaximumEnvelopeIterations = 5;
    private const double OrdinaryEnvelopeProbeLookaheadSteps = 1d;
    private const double OrdinaryTerminalProbeFraction = 0.5d;
    private readonly IQpSolver _qpSolver;
    private readonly PathSpeedLimitBuilder _speedLimitBuilder;
    private readonly LongitudinalConstraintBuilder _constraintBuilder;
    private readonly LongitudinalSolutionValidator _solutionValidator;

    public SequentialLongitudinalOptimizer(IQpSolver qpSolver)
        : this(qpSolver, new PathSpeedLimitBuilder(), new LongitudinalConstraintBuilder(new LongitudinalObjectiveBuilder()),
            new LongitudinalSolutionValidator())
    {
    }

    internal SequentialLongitudinalOptimizer(IQpSolver qpSolver, PathSpeedLimitBuilder speedLimitBuilder,
        LongitudinalConstraintBuilder constraintBuilder, LongitudinalSolutionValidator solutionValidator)
    {
        _qpSolver = qpSolver ?? throw new ArgumentNullException(nameof(qpSolver));
        _speedLimitBuilder = speedLimitBuilder ?? throw new ArgumentNullException(nameof(speedLimitBuilder));
        _constraintBuilder = constraintBuilder ?? throw new ArgumentNullException(nameof(constraintBuilder));
        _solutionValidator = solutionValidator ?? throw new ArgumentNullException(nameof(solutionValidator));
    }

    public LongitudinalPlanningResult Optimize(LongitudinalPlanningInput input, CancellationToken cancellationToken)
    {
        if (input == null)
            return Failed(EmPlanningStatus.InvalidInput, "Longitudinal planning input is required.");
        if (cancellationToken.IsCancellationRequested)
            return Failed(EmPlanningStatus.Cancelled, "Longitudinal optimization was cancelled.");
        if (!TryCreateSettings(input, out QpSolverSettings settings, out TimeSpan totalBudget, out double convergenceTolerance,
                out int iterationLimit, out string configurationFailure))
        {
            return Failed(EmPlanningStatus.InvalidInput, configurationFailure);
        }

        EmPlanningStatus speedStatus = _speedLimitBuilder.Build(input, out PathSpeedLimit speedLimit, out string speedFailure);
        if (speedStatus != EmPlanningStatus.Success)
            return Failed(speedStatus, speedFailure);

        LongitudinalCandidate iterate = CreateInitialIterate(input);
        double[] warmStart = ToPrimal(iterate);
        bool hasDynamicsConsistentInitialWarmStart = iterate.SatisfiesExactDiscreteDynamics(1e-12d);
        LongitudinalCandidate lastStrictCandidate;
        if (!_solutionValidator.TryValidate(input, speedLimit, iterate, out lastStrictCandidate, out _))
            lastStrictCandidate = null;
        string lastCandidateRejection = string.Empty;
        bool hasPreviousObjective = false;
        double previousObjective = 0d;
        var stopwatch = Stopwatch.StartNew();

        for (int iteration = 0; iteration < iterationLimit; iteration++)
        {
            if (cancellationToken.IsCancellationRequested)
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.Cancelled, "Longitudinal optimization was cancelled.");
            TimeSpan remainingBudget = totalBudget - stopwatch.Elapsed;
            if (remainingBudget <= TimeSpan.Zero)
            {
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.SolverTimedOut,
                    "Longitudinal optimization exhausted its solve budget.");
            }
            if (!_constraintBuilder.TryBuild(input, speedLimit, iterate, out QuadraticProgram problem, out string buildFailure))
            {
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.LongitudinalInfeasible,
                    "Longitudinal constraints are infeasible: " + buildFailure);
            }

            QpSolveResult solved = _qpSolver.Solve(problem,
                new QpSolverSettings(settings.MaximumIterations, settings.AbsoluteTolerance, settings.RelativeTolerance,
                    remainingBudget, settings.EnableWarmStart && (iteration > 0 || hasDynamicsConsistentInitialWarmStart),
                    settings.EnablePolishing,
                    settings.EnableNativeVerboseOutput),
                warmStart, cancellationToken);
            if (solved == null)
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.Failed, "The longitudinal QP solver returned no result.");
            if (solved.Status == QpSolveStatus.TimeLimit || solved.Status == QpSolveStatus.MaximumIterations)
            {
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.SolverTimedOut,
                    "The longitudinal QP solver timed out (status=" + solved.NativeStatus +
                    ", iterations=" + solved.Iterations + ", primal=" + solved.PrimalResidual +
                    ", dual=" + solved.DualResidual + "): " + solved.Diagnostic);
            }
            if (solved.Status == QpSolveStatus.Cancelled)
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.Cancelled,
                    "The longitudinal QP solver was cancelled: " + solved.Diagnostic);
            if (solved.Status == QpSolveStatus.PrimalInfeasible || solved.Status == QpSolveStatus.DualInfeasible)
            {
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.LongitudinalInfeasible,
                    "The longitudinal QP solver reported infeasibility: " + solved.Diagnostic);
            }
            if (solved.Status == QpSolveStatus.SolverUnavailable)
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.SolverUnavailable,
                    "The longitudinal QP solver is unavailable: " + solved.Diagnostic);
            if (solved.Status != QpSolveStatus.Solved && solved.Status != QpSolveStatus.SolvedInaccurate)
            {
                return FallbackOrFailure(lastStrictCandidate, EmPlanningStatus.Failed,
                    "The longitudinal QP solver failed: " + solved.Diagnostic);
            }
            if (solved.Status == QpSolveStatus.SolvedInaccurate && !HasStrictResiduals(solved, convergenceTolerance))
            {
                lastCandidateRejection = "SolvedInaccurate residuals exceed the strict acceptance tolerance" +
                    " (primal=" + solved.PrimalResidual + ", dual=" + solved.DualResidual + ").";
                if (TryCreateCandidate(iterate.KnotTimes, solved.Primal, out LongitudinalCandidate inaccurateCandidate))
                    warmStart = ToPrimal(inaccurateCandidate);
                continue;
            }
            if (!TryCreateCandidate(iterate.KnotTimes, solved.Primal, out LongitudinalCandidate candidate))
            {
                lastCandidateRejection = "The solver primal does not match the ST variable layout.";
                continue;
            }
            if (!_solutionValidator.TryValidate(input, speedLimit, candidate, out LongitudinalCandidate validated,
                    out string validationFailure))
            {
                string rejection = validationFailure + CreateEnvelopeDiagnostic(speedLimit, iterate, candidate,
                    iteration + 1);
                lastCandidateRejection = string.IsNullOrEmpty(lastCandidateRejection)
                    ? rejection
                    : lastCandidateRejection + " | " + rejection;
                if (TryCreateEnvelopeIterate(input, iterate, candidate, out LongitudinalCandidate nextIterate))
                {
                    iterate = nextIterate;
                    warmStart = ToPrimal(candidate);
                }
                continue;
            }

            double maximumChange = MaximumProgressOrSpeedChange(iterate, validated);
            double relativeObjectiveImprovement = hasPreviousObjective
                ? RelativeObjectiveImprovement(previousObjective, solved.Objective)
                : double.PositiveInfinity;
            lastStrictCandidate = CopyCandidate(validated);
            iterate = validated;
            warmStart = ToPrimal(validated);
            previousObjective = solved.Objective;
            hasPreviousObjective = true;
            if (maximumChange <= convergenceTolerance && relativeObjectiveImprovement <= convergenceTolerance)
                return new LongitudinalPlanningResult(EmPlanningStatus.Success, lastStrictCandidate, string.Empty);
        }

        return lastStrictCandidate == null
            ? Failed(EmPlanningStatus.LongitudinalInfeasible, "No strictly validated longitudinal candidate was found. " +
                lastCandidateRejection)
            : new LongitudinalPlanningResult(EmPlanningStatus.Success, lastStrictCandidate, string.Empty);
    }

    private static bool TryCreateSettings(LongitudinalPlanningInput input, out QpSolverSettings settings,
        out TimeSpan totalBudget, out double convergenceTolerance, out int iterationLimit, out string failureReason)
    {
        settings = null;
        totalBudget = TimeSpan.Zero;
        convergenceTolerance = 0d;
        iterationLimit = 0;
        failureReason = string.Empty;
        if (input.Configuration == null || input.Configuration.Solver == null || input.Configuration.Scheduling == null)
        {
            failureReason = "Longitudinal solver configuration is required.";
            return false;
        }
        SolverConfiguration solver = input.Configuration.Solver;
        SchedulingConfiguration scheduling = input.Configuration.Scheduling;
        if (solver.MaximumOuterIterations <= 0 || solver.MaximumOsqpIterations <= 0 ||
            !IsPositiveFinite(solver.AbsoluteTolerance) || !IsPositiveFinite(solver.RelativeTolerance) ||
            !IsPositiveFinite(solver.StrictResidualTolerance) || !IsPositiveFinite(scheduling.SolverTimeoutSeconds))
        {
            failureReason = "Longitudinal solver configuration is invalid.";
            return false;
        }
        try
        {
            totalBudget = TimeSpan.FromSeconds(scheduling.SolverTimeoutSeconds);
            settings = new QpSolverSettings(solver.MaximumOsqpIterations, solver.AbsoluteTolerance, solver.RelativeTolerance,
                totalBudget, solver.WarmStart, solver.Polish, solver.NativeVerbose);
            convergenceTolerance = solver.StrictResidualTolerance;
            iterationLimit = Math.Min(MaximumEnvelopeIterations, solver.MaximumOuterIterations);
            return true;
        }
        catch (ArgumentException exception)
        {
            failureReason = exception.Message;
            return false;
        }
    }

    private static LongitudinalCandidate CreateInitialIterate(LongitudinalPlanningInput input)
    {
        IReadOnlyList<double> times = LongitudinalCandidate.CreateKnotTimes(input.Configuration.Scheduling.TimeHorizonSeconds,
            input.Configuration.Scheduling.OutputTimeStepSeconds);
        if (TryCreateCruiseThenBrakeSeed(input, times, out LongitudinalCandidate brakingSeed))
            return brakingSeed;
        int knotCount = times.Count;
        var s = new double[knotCount];
        var u = new double[knotCount];
        var a = new double[knotCount];
        var j = new double[knotCount - 1];
        double horizon = times[knotCount - 1];
        double requestedSpeed = Math.Min(input.DirectionMaximumSpeedMetersPerSecond,
            Math.Max(0d, input.TerminalPathS / horizon));
        if (!JerkLimitedStoppingMath.TryCalculate(requestedSpeed, 0d,
                input.Configuration.Longitudinal.MaximumDecelerationMetersPerSecondSquared,
                input.Configuration.Longitudinal.MaximumJerkMetersPerSecondCubed,
                out JerkLimitedStoppingProfile terminalStop, out _))
        {
            throw new ArgumentOutOfRangeException(nameof(input));
        }
        double cruiseDistance = Math.Max(0d, input.TerminalPathS - terminalStop.DistanceMeters);
        for (int index = 0; index < knotCount; index++)
        {
            double fraction = (double)index / (knotCount - 1);
            s[index] = Math.Min(input.TerminalPathS, cruiseDistance * fraction + terminalStop.DistanceMeters * fraction * fraction);
            u[index] = index == 0 ? input.InitialProgressSpeedMetersPerSecond : requestedSpeed;
            a[index] = index == 0 ? input.InitialAccelerationMetersPerSecondSquared : 0d;
        }
        s[0] = 0d;
        s[knotCount - 1] = input.TerminalPathS;
        u[knotCount - 1] = 0d;
        a[knotCount - 1] = 0d;
        return new LongitudinalCandidate(times, s, u, a, j);
    }

    private static bool TryCreateCruiseThenBrakeSeed(LongitudinalPlanningInput input, IReadOnlyList<double> times,
        out LongitudinalCandidate candidate)
    {
        candidate = null;
        double initialSpeed = input.InitialProgressSpeedMetersPerSecond;
        double initialAcceleration = input.InitialAccelerationMetersPerSecondSquared;
        if (initialSpeed <= 0d || Math.Abs(initialAcceleration) > 1e-12d)
            return false;
        double timeStep = times[1] - times[0];
        for (int index = 1; index < times.Count - 1; index++)
        {
            if (Math.Abs((times[index + 1] - times[index]) - timeStep) > 1e-12d)
                return false;
        }
        LongitudinalConfiguration configuration = input.Configuration.Longitudinal;
        int intervalCount = times.Count - 1;
        int maximumRampIntervals = Math.Min(intervalCount / 2, checked((int)Math.Floor(
            configuration.MaximumDecelerationMetersPerSecondSquared /
            (configuration.MaximumJerkMetersPerSecondCubed * timeStep))));
        for (int rampIntervals = maximumRampIntervals; rampIntervals >= 1; rampIntervals--)
        {
            for (int plateauIntervals = 0; 2 * rampIntervals + plateauIntervals <= intervalCount; plateauIntervals++)
            {
                double jerkMagnitude = initialSpeed / (rampIntervals * (rampIntervals + plateauIntervals) *
                    timeStep * timeStep);
                double peakDeceleration = jerkMagnitude * rampIntervals * timeStep;
                if (jerkMagnitude > configuration.MaximumJerkMetersPerSecondCubed + 1e-12d ||
                    peakDeceleration > configuration.MaximumDecelerationMetersPerSecondSquared + 1e-12d)
                {
                    continue;
                }
                int brakingIntervals = 2 * rampIntervals + plateauIntervals;
                double brakingDistance = 0.5d * initialSpeed * brakingIntervals * timeStep;
                if (brakingDistance > input.TerminalPathS + 1e-12d)
                    continue;
                int maximumCruiseIntervals = intervalCount - brakingIntervals;
                int cruiseIntervals = Math.Min(maximumCruiseIntervals, Math.Max(0, checked((int)Math.Floor(
                    (input.TerminalPathS - brakingDistance) / (initialSpeed * timeStep) + 1e-12d))));
                var jerk = new double[intervalCount];
                int cursor = cruiseIntervals;
                for (int index = 0; index < rampIntervals; index++)
                    jerk[cursor++] = -jerkMagnitude;
                cursor += plateauIntervals;
                for (int index = 0; index < rampIntervals; index++)
                    jerk[cursor++] = jerkMagnitude;
                LongitudinalCandidate integrated = LongitudinalCandidate.Integrate(times, 0d, initialSpeed, 0d, jerk);
                if (integrated.S[integrated.S.Count - 1] <= input.TerminalPathS + 1e-12d)
                {
                    candidate = integrated;
                    return true;
                }
            }
        }
        return false;
    }

    private static bool TryCreateCandidate(IReadOnlyList<double> times, IReadOnlyList<double> primal,
        out LongitudinalCandidate candidate)
    {
        candidate = null;
        if (primal == null)
            return false;
        try
        {
            var layout = new LongitudinalVariableLayout(times.Count);
            if (primal.Count != layout.VariableCount)
                return false;
            var s = new double[layout.KnotCount];
            var u = new double[layout.KnotCount];
            var a = new double[layout.KnotCount];
            var j = new double[layout.KnotCount - 1];
            for (int index = 0; index < layout.KnotCount; index++)
            {
                s[index] = primal[layout.S(index)];
                u[index] = primal[layout.U(index)];
                a[index] = primal[layout.A(index)];
            }
            for (int index = 0; index < layout.KnotCount - 1; index++)
                j[index] = primal[layout.J(index)];
            candidate = new LongitudinalCandidate(times, s, u, a, j);
            return true;
        }
        catch (ArgumentException)
        {
            return false;
        }
    }

    private static double[] ToPrimal(LongitudinalCandidate candidate)
    {
        var layout = new LongitudinalVariableLayout(candidate.KnotTimes.Count);
        var primal = new double[layout.VariableCount];
        for (int index = 0; index < layout.KnotCount; index++)
        {
            primal[layout.S(index)] = candidate.S[index];
            primal[layout.U(index)] = candidate.U[index];
            primal[layout.A(index)] = candidate.A[index];
        }
        for (int index = 0; index < layout.KnotCount - 1; index++)
            primal[layout.J(index)] = candidate.J[index];
        return primal;
    }

    private static bool TryCreateEnvelopeIterate(LongitudinalPlanningInput input, LongitudinalCandidate previous,
        LongitudinalCandidate candidate, out LongitudinalCandidate nextIterate)
    {
        nextIterate = null;
        if (candidate.S.Count != previous.S.Count)
            return false;
        var candidateProgressSamples = new double[candidate.S.Count];
        double priorProgress = double.NegativeInfinity;
        double priorPreviousProgress = double.NegativeInfinity;
        for (int index = 0; index < candidate.S.Count; index++)
        {
            double candidateProgress = candidate.S[index];
            double previousProgress = previous.S[index];
            if (!IsFinite(previousProgress) || previousProgress < 0d || previousProgress > input.TerminalPathS ||
                previousProgress < priorPreviousProgress)
            {
                return false;
            }
            if (!IsFinite(candidateProgress))
            {
                candidateProgress = previousProgress;
            }
            candidateProgress = Math.Max(0d, Math.Min(input.TerminalPathS, candidateProgress));
            if (index == candidate.S.Count - 1)
                candidateProgress = input.TerminalPathS;
            candidateProgress = Math.Max(priorProgress, candidateProgress);
            candidateProgressSamples[index] = candidateProgress;
            priorProgress = candidateProgress;
            priorPreviousProgress = previousProgress;
        }
        var progress = new double[candidate.S.Count];
        double previousNextProgress = 0d;
        for (int index = 0; index < progress.Length; index++)
        {
            double candidateProgress = candidateProgressSamples[index];
            if (index == 0 || index == progress.Length - 1 || candidateProgress >= input.TerminalPathS)
            {
                progress[index] = candidateProgress;
            }
            else
            {
                double timeStep = candidate.KnotTimes[index + 1] - candidate.KnotTimes[index];
                double iterationAdvance = Math.Max(0d, candidateProgress - previous.S[index]);
                double candidateSpeed = IsFinite(candidate.U[index]) ? Math.Max(0d, candidate.U[index]) : 0d;
                double lookaheadAdvance = IsFinite(candidate.U[index])
                    ? OrdinaryEnvelopeProbeLookaheadSteps * candidateSpeed * timeStep
                    : 0d;
                double terminalLimitedAdvance = OrdinaryTerminalProbeFraction *
                    (input.TerminalPathS - candidateProgress);
                double advance = Math.Min(Math.Max(iterationAdvance, lookaheadAdvance), terminalLimitedAdvance);
                progress[index] = candidateProgress + advance;
            }
            progress[index] = Math.Max(previousNextProgress, progress[index]);
            previousNextProgress = progress[index];
        }
        nextIterate = new LongitudinalCandidate(candidate.KnotTimes, progress, previous.U, previous.A, previous.J);
        return true;
    }

    private static bool HasStrictResiduals(QpSolveResult result, double tolerance)
    {
        return IsPositiveFinite(tolerance) && result.PrimalResidual >= 0d && result.DualResidual >= 0d &&
            result.PrimalResidual <= tolerance && result.DualResidual <= tolerance;
    }

    private static string CreateEnvelopeDiagnostic(PathSpeedLimit speedLimit, LongitudinalCandidate iterate,
        LongitudinalCandidate candidate, int iteration)
    {
        int worstIndex = -1;
        double worstExcess = double.NegativeInfinity;
        for (int index = 0; index < candidate.S.Count; index++)
        {
            if (!IsFinite(candidate.S[index]) || !IsFinite(candidate.U[index]))
                continue;
            double candidateProgress = Math.Max(0d, Math.Min(speedLimit.TerminalPathS, candidate.S[index]));
            double limit = speedLimit.MaximumSpeedAt(candidateProgress);
            double excess = candidate.U[index] - limit;
            if (excess > worstExcess)
            {
                worstExcess = excess;
                worstIndex = index;
            }
        }
        if (worstIndex < 0)
            return "";
        return " Envelope iteration " + iteration + " used PathS=" + iterate.S[worstIndex] +
            " and produced PathS=" + candidate.S[worstIndex] + " at its largest speed-envelope excess.";
    }

    private static double MaximumProgressOrSpeedChange(LongitudinalCandidate previous, LongitudinalCandidate current)
    {
        double maximum = 0d;
        for (int index = 0; index < previous.S.Count; index++)
        {
            maximum = Math.Max(maximum, Math.Abs(current.S[index] - previous.S[index]));
            maximum = Math.Max(maximum, Math.Abs(current.U[index] - previous.U[index]));
        }
        return maximum;
    }

    private static double RelativeObjectiveImprovement(double previous, double current)
    {
        return Math.Abs(previous - current) / Math.Max(1d, Math.Abs(previous));
    }

    private static LongitudinalPlanningResult FallbackOrFailure(LongitudinalCandidate candidate,
        EmPlanningStatus failureStatus, string failureReason)
    {
        return candidate == null
            ? Failed(failureStatus, failureReason)
            : new LongitudinalPlanningResult(EmPlanningStatus.SuccessWithFallback, candidate, failureReason);
    }

    private static LongitudinalPlanningResult Failed(EmPlanningStatus status, string reason)
    {
        return new LongitudinalPlanningResult(status, null, reason);
    }

    private static LongitudinalCandidate CopyCandidate(LongitudinalCandidate source)
    {
        return new LongitudinalCandidate(source.KnotTimes, source.S, source.U, source.A, source.J);
    }

    private static bool IsPositiveFinite(double value)
    {
        return IsFinite(value) && value > 0d;
    }

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