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

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

/// <summary>Runs bounded lateral SQP iterations and retains only independently validated candidates.</summary>
public sealed class SequentialConvexOptimizer
{
    private const double StationTolerance = 1e-12d;
    private readonly IQpSolver _qpSolver;
    private readonly LateralConstraintBuilder _constraintBuilder;
    private readonly LateralGeometryEvaluator _geometryEvaluator;
    private readonly LateralSolutionValidator _solutionValidator;

    public SequentialConvexOptimizer(IQpSolver qpSolver)
        : this(qpSolver, new LateralConstraintBuilder(new LateralObjectiveBuilder()), new LateralGeometryEvaluator(),
            new LateralSolutionValidator())
    {
    }

    internal SequentialConvexOptimizer(IQpSolver qpSolver, LateralConstraintBuilder constraintBuilder,
        LateralGeometryEvaluator geometryEvaluator, LateralSolutionValidator solutionValidator)
    {
        _qpSolver = qpSolver ?? throw new ArgumentNullException(nameof(qpSolver));
        _constraintBuilder = constraintBuilder ?? throw new ArgumentNullException(nameof(constraintBuilder));
        _geometryEvaluator = geometryEvaluator ?? throw new ArgumentNullException(nameof(geometryEvaluator));
        _solutionValidator = solutionValidator ?? throw new ArgumentNullException(nameof(solutionValidator));
    }

    public LateralPlanningResult Optimize(LateralPlanningInput input, CancellationToken cancellationToken)
    {
        if (input == null)
            return Failed(EmPlanningStatus.InvalidInput, "Lateral planning input is required.");

        if (!TryCreateSettings(input, out QpSolverSettings settings, out TimeSpan totalBudget, out double convergenceTolerance,
                out string configurationFailure))
        {
            return Failed(EmPlanningStatus.InvalidInput, configurationFailure);
        }

        LateralCandidate iterate = CreateInitialIterate(input);
        var warmStart = new double[new LateralVariableLayout(input.ReferenceStations.Count).VariableCount];
        LateralPath lastValidatedPath = null;
        double previousObjective = 0d;
        bool hasPreviousObjective = false;
        var stopwatch = Stopwatch.StartNew();

        for (int iteration = 0; iteration < input.Configuration.Solver.MaximumOuterIterations; iteration++)
        {
            if (cancellationToken.IsCancellationRequested)
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.Cancelled, "Lateral SQP was cancelled.");

            TimeSpan remainingBudget = totalBudget - stopwatch.Elapsed;
            if (remainingBudget <= TimeSpan.Zero)
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.SolverTimedOut,
                    "Lateral SQP exhausted its solve budget.");

            if (!_constraintBuilder.TryBuild(input, iterate, out QuadraticProgram problem, out string failureReason))
            {
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.LateralInfeasible,
                    "Lateral SQP constraints are infeasible: " + failureReason);
            }

            QpSolveResult solved = _qpSolver.Solve(problem,
                new QpSolverSettings(settings.MaximumIterations, settings.AbsoluteTolerance, settings.RelativeTolerance,
                    remainingBudget, settings.EnableWarmStart, settings.EnablePolishing, settings.EnableNativeVerboseOutput),
                warmStart, cancellationToken);

            if (cancellationToken.IsCancellationRequested)
                return Failed(EmPlanningStatus.Cancelled, "Lateral SQP was cancelled after the QP solve.");
            if (solved == null)
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.Failed, "The lateral QP solver returned no result.");

            if (solved.Status == QpSolveStatus.TimeLimit || solved.Status == QpSolveStatus.MaximumIterations)
            {
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.SolverTimedOut,
                    "The lateral QP solver timed out: " + solved.Diagnostic);
            }
            if (solved.Status == QpSolveStatus.Cancelled)
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.Cancelled,
                    "The lateral QP solver was cancelled: " + solved.Diagnostic);
            if (solved.Status == QpSolveStatus.PrimalInfeasible || solved.Status == QpSolveStatus.DualInfeasible)
            {
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.LateralInfeasible,
                    "The lateral QP solver reported infeasibility: " + solved.Diagnostic);
            }
            if (solved.Status == QpSolveStatus.SolverUnavailable)
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.SolverUnavailable,
                    "The lateral QP solver is unavailable: " + solved.Diagnostic);
            if (solved.Status != QpSolveStatus.Solved && solved.Status != QpSolveStatus.SolvedInaccurate)
            {
                return FallbackOrFailure(lastValidatedPath, EmPlanningStatus.Failed,
                    "The lateral QP solver failed: " + solved.Diagnostic);
            }
            if (solved.Status == QpSolveStatus.SolvedInaccurate && !HasStrictResiduals(solved, input.Configuration.Solver.StrictResidualTolerance))
            {
                continue;
            }
            if (!TryCreateCandidate(input.ReferenceStations, solved.Primal, out LateralCandidate candidate))
                continue;
            if (!_geometryEvaluator.TryEvaluate(input, candidate, out LateralPath evaluatedPath, out _))
                continue;

            // A geometrically evaluable QP solution that remains inside the static
            // corridor is the next SQP linearization point, even when the independent
            // validator rejects it.  Otherwise the next outer pass rebuilds the identical
            // convex subproblem and cannot correct that result.  An out-of-corridor vector
            // must not become an iterate: it can make the following trust region infeasible.
            LateralCandidate previousIterate = iterate;
            if (RespectsStaticCorridor(input, candidate))
            {
                iterate = candidate;
                warmStart = CopyValues(solved.Primal);
            }

            if (!_solutionValidator.TryValidate(input, candidate, evaluatedPath, out LateralPath validatedPath, out _))
                continue;

            double maximumLateralChange = MaximumLateralChange(previousIterate, candidate);
            double relativeObjectiveImprovement = hasPreviousObjective
                ? RelativeObjectiveImprovement(previousObjective, solved.Objective)
                : double.PositiveInfinity;
            lastValidatedPath = CopyPath(validatedPath);
            previousObjective = solved.Objective;
            hasPreviousObjective = true;

            if (maximumLateralChange <= convergenceTolerance && relativeObjectiveImprovement <= convergenceTolerance)
                return new LateralPlanningResult(EmPlanningStatus.Success, lastValidatedPath, string.Empty);
        }

        return lastValidatedPath == null
            ? Failed(EmPlanningStatus.LateralInfeasible, "No independently validated lateral candidate was found.")
            : new LateralPlanningResult(
                EmPlanningStatus.SuccessWithFallback,
                lastValidatedPath,
                "Lateral SQP reached its outer-iteration limit before strict convergence.");
    }

    private static bool TryCreateSettings(LateralPlanningInput input, out QpSolverSettings settings, out TimeSpan totalBudget,
        out double convergenceTolerance, out string failureReason)
    {
        settings = null;
        totalBudget = TimeSpan.Zero;
        convergenceTolerance = 0d;
        failureReason = string.Empty;
        SolverConfiguration solver = input.Configuration.Solver;
        SchedulingConfiguration scheduling = input.Configuration.Scheduling;
        if (solver == null || scheduling == null || solver.MaximumOuterIterations <= 0 ||
            !IsPositiveFinite(solver.AbsoluteTolerance) || !IsPositiveFinite(solver.RelativeTolerance) ||
            !IsPositiveFinite(solver.StrictResidualTolerance) || !IsPositiveFinite(scheduling.SolverTimeoutSeconds))
        {
            failureReason = "The lateral SQP 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;
            return true;
        }
        catch (ArgumentException exception)
        {
            failureReason = exception.Message;
            return false;
        }
    }

    private static LateralCandidate CreateInitialIterate(LateralPlanningInput input)
    {
        int stationCount = input.ReferenceStations.Count;
        var l = new double[stationCount];
        var dl = new double[stationCount];
        var ddl = new double[stationCount];
        var dddl = new double[stationCount - 1];
        bool coversAllStations = input.PreviousTrajectorySeed.Count >= 2 &&
            input.PreviousTrajectorySeed[0].ReferenceS <= input.ReferenceStations[0] + StationTolerance &&
            input.PreviousTrajectorySeed[input.PreviousTrajectorySeed.Count - 1].ReferenceS >=
            input.ReferenceStations[stationCount - 1] - StationTolerance;

        for (int index = 0; index < stationCount; index++)
        {
            LateralInterval corridor = input.Corridor.Stations[index];
            l[index] = coversAllStations
                ? InterpolateSeedL(input.PreviousTrajectorySeed, input.ReferenceStations[index])
                : Clamp(0d, corridor.MinimumL, corridor.MaximumL);
        }

        double startDenominator = 1d - input.StartProjection.ReferencePoint.GeometricCurvature *
            input.StartProjection.LateralOffset;
        l[0] = input.StartProjection.LateralOffset;
        dl[0] = startDenominator * Math.Tan(input.StartProjection.HeadingError);
        return new LateralCandidate(input.ReferenceStations, l, dl, ddl, dddl);
    }

    private static bool TryCreateCandidate(IReadOnlyList<double> stations, IReadOnlyList<double> primal,
        out LateralCandidate candidate)
    {
        candidate = null;
        if (primal == null)
            return false;
        try
        {
            var layout = new LateralVariableLayout(stations.Count);
            if (primal.Count != layout.VariableCount)
                return false;
            var l = new double[layout.StationCount];
            var dl = new double[layout.StationCount];
            var ddl = new double[layout.StationCount];
            var dddl = new double[layout.StationCount - 1];
            for (int index = 0; index < layout.StationCount; index++)
            {
                l[index] = primal[layout.L(index)];
                dl[index] = primal[layout.DL(index)];
                ddl[index] = primal[layout.DDL(index)];
            }
            for (int index = 0; index < dddl.Length; index++)
                dddl[index] = primal[layout.DDDL(index)];
            candidate = new LateralCandidate(stations, l, dl, ddl, dddl);
            return true;
        }
        catch (ArgumentException)
        {
            return false;
        }
    }

    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 double MaximumLateralChange(LateralCandidate previous, LateralCandidate current)
    {
        double maximum = 0d;
        for (int index = 0; index < previous.L.Count; index++)
            maximum = Math.Max(maximum, Math.Abs(current.L[index] - previous.L[index]));
        return maximum;
    }

    private static bool RespectsStaticCorridor(LateralPlanningInput input, LateralCandidate candidate)
    {
        for (int index = 0; index < candidate.L.Count; index++)
        {
            LateralInterval interval = input.Corridor.Stations[index];
            if (candidate.L[index] < interval.MinimumL || candidate.L[index] > interval.MaximumL)
                return false;
        }
        return true;
    }

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

    private static LateralPlanningResult FallbackOrFailure(LateralPath path, EmPlanningStatus failureStatus, string failureReason)
    {
        return failureStatus == EmPlanningStatus.Cancelled || path == null
            ? Failed(failureStatus, failureReason)
            : new LateralPlanningResult(EmPlanningStatus.SuccessWithFallback, path, failureReason);
    }

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

    private static LateralPath CopyPath(LateralPath source)
    {
        var points = new List<LateralPathPoint>(source.Points.Count);
        for (int index = 0; index < source.Points.Count; index++)
        {
            LateralPathPoint point = source.Points[index];
            points.Add(new LateralPathPoint(point.ReferenceS, point.PathS, point.L, point.DL, point.DDL, point.DDDL,
                point.X, point.Y, point.VehicleYaw, point.GeometricCurvature, point.VehicleCurvature,
                point.VehicleCurvatureDerivative));
        }
        return new LateralPath(points, true);
    }

    private static double[] CopyValues(IReadOnlyList<double> source)
    {
        var copy = new double[source.Count];
        for (int index = 0; index < source.Count; index++)
            copy[index] = source[index];
        return copy;
    }

    private static double InterpolateSeedL(IReadOnlyList<FrenetProjection> seed, double referenceS)
    {
        if (referenceS <= seed[0].ReferenceS)
            return seed[0].LateralOffset;
        for (int index = 1; index < seed.Count; index++)
        {
            if (referenceS <= seed[index].ReferenceS)
            {
                FrenetProjection lower = seed[index - 1];
                FrenetProjection upper = seed[index];
                double span = upper.ReferenceS - lower.ReferenceS;
                return span <= StationTolerance ? upper.LateralOffset : lower.LateralOffset +
                    (upper.LateralOffset - lower.LateralOffset) * (referenceS - lower.ReferenceS) / span;
            }
        }
        return seed[seed.Count - 1].LateralOffset;
    }

    private static double Clamp(double value, double minimum, double maximum)
    {
        return Math.Max(minimum, Math.Min(maximum, value));
    }

    private static bool IsPositiveFinite(double value)
    {
        return !double.IsNaN(value) && !double.IsInfinity(value) && value > 0d;
    }
}
