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

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

/// <summary>Builds normalized squared-residual costs in OSQP's 0.5*x'P*x + q'x convention.</summary>
public sealed class LateralObjectiveBuilder
{
    public void AddTerms(LateralPlanningInput input, LateralVariableLayout layout, LateralCandidate linearization,
        SparseTripletBuilder hessian, IList<double> linearCost)
    {
        if (input == null)
            throw new ArgumentNullException(nameof(input));
        if (layout == null)
            throw new ArgumentNullException(nameof(layout));
        if (linearization == null)
            throw new ArgumentNullException(nameof(linearization));
        if (hessian == null)
            throw new ArgumentNullException(nameof(hessian));
        if (linearCost == null || linearCost.Count != layout.VariableCount)
            throw new ArgumentException("Linear cost must match the lateral layout.", nameof(linearCost));

        LateralConfiguration lateral = input.Configuration.Lateral ?? throw new ArgumentException("Missing lateral configuration.");
        LateralWeights weights = lateral.Weights ?? throw new ArgumentException("Missing lateral weights.");
        double lateralScale = RequirePositive(input.Configuration.Corridor.MaximumLateralOffsetMeters, "lateral scale");
        double slopeScale = RequirePositive(lateral.MaximumLateralSlope, "slope scale");
        double secondDerivativeScale = RequirePositive(lateral.MaximumLateralSecondDerivativePerMeter, "second-derivative scale");
        double thirdDerivativeScale = RequirePositive(lateral.MaximumLateralThirdDerivativePerSquareMeter, "third-derivative scale");
        double curvatureScale = RequirePositive(LateralCurvatureLinearization.GetMaximumVehicleCurvature(input.Vehicle),
            "curvature scale");
        double curvatureVariationScale = GetCurvatureVariationScale(input);

        for (int station = 0; station < layout.StationCount; station++)
        {
            AddSquaredResidual(hessian, linearCost, new[] { layout.L(station) }, new[] { 1d }, 0d,
                weights.ReferenceOffset, lateralScale);
            AddSquaredResidual(hessian, linearCost, new[] { layout.DL(station) }, new[] { 1d }, 0d,
                weights.HeadingDeviation, slopeScale);
            AddSquaredResidual(hessian, linearCost, new[] { layout.DDL(station) }, new[] { 1d }, 0d,
                weights.SecondDerivative, secondDerivativeScale);
        }
        for (int interval = 0; interval < layout.StationCount - 1; interval++)
        {
            AddSquaredResidual(hessian, linearCost, new[] { layout.DDDL(interval) }, new[] { 1d }, 0d,
                weights.ThirdDerivative, thirdDerivativeScale);
        }

        AddPreviousTrajectoryTerms(input, layout, hessian, linearCost, weights.PreviousTrajectory, lateralScale);
        IReadOnlyList<LateralCurvatureLinearization> curvature = LateralCurvatureLinearization.Create(input, layout,
            linearization);
        for (int station = 0; station < curvature.Count; station++)
        {
            AddSquaredResidual(hessian, linearCost, curvature[station].VariableIndices, curvature[station].Gradient,
                curvature[station].Constant, weights.Curvature, curvatureScale);
        }
        AddCurvatureVariationTerms(input.ReferenceStations, curvature, hessian, linearCost, weights.CurvatureVariation,
            curvatureVariationScale);
        if (input.TerminalType == EmTerminalType.RollingSafetyStop)
        {
            AddSquaredResidual(hessian, linearCost, new[] { layout.L(layout.StationCount - 1) }, new[] { 1d }, 0d,
                weights.RollingTerminal, lateralScale);
        }
    }

    private static void AddPreviousTrajectoryTerms(LateralPlanningInput input, LateralVariableLayout layout,
        SparseTripletBuilder hessian, IList<double> linearCost, double weight, double lateralScale)
    {
        if (input.PreviousTrajectorySeed.Count == 0)
            return;

        for (int station = 0; station < layout.StationCount; station++)
        {
            double previousL = InterpolatePreviousL(input.PreviousTrajectorySeed, input.ReferenceStations[station]);
            AddSquaredResidual(hessian, linearCost, new[] { layout.L(station) }, new[] { 1d }, -previousL,
                weight, lateralScale);
        }
    }

    private static void AddCurvatureVariationTerms(IReadOnlyList<double> stations,
        IReadOnlyList<LateralCurvatureLinearization> curvature,
        SparseTripletBuilder hessian, IList<double> linearCost, double weight, double scale)
    {
        for (int station = 0; station < curvature.Count; station++)
        {
            int lower = station == 0 ? 0 : station - 1;
            int upper = station == curvature.Count - 1 ? curvature.Count - 1 : station + 1;
            double ds = stations[upper] - stations[lower];
            if (!IsFinite(ds) || ds <= 0d)
                throw new ArgumentException("Curvature variation requires strictly increasing stations.", nameof(stations));
            LateralCurvatureLinearization left = curvature[lower];
            LateralCurvatureLinearization right = curvature[upper];
            var indices = new int[left.VariableIndices.Count + right.VariableIndices.Count];
            var gradient = new double[indices.Length];
            for (int index = 0; index < left.VariableIndices.Count; index++)
            {
                indices[index] = left.VariableIndices[index];
                gradient[index] = -left.Gradient[index] / ds;
                indices[left.VariableIndices.Count + index] = right.VariableIndices[index];
                gradient[left.VariableIndices.Count + index] = right.Gradient[index] / ds;
            }
            AddSquaredResidual(hessian, linearCost, indices, gradient, (right.Constant - left.Constant) / ds, weight, scale);
        }
    }

    private static void AddSquaredResidual(SparseTripletBuilder hessian, IList<double> linearCost,
        IReadOnlyList<int> indices, IReadOnlyList<double> gradient, double constant, double weight, double scale)
    {
        if (indices.Count != gradient.Count || indices.Count == 0 || !IsFinite(constant))
            throw new ArgumentException("Affine residual is invalid.");
        if (!IsFinite(weight) || weight < 0d)
            throw new ArgumentOutOfRangeException(nameof(weight));
        double coefficient = 2d * weight / (scale * scale);
        for (int left = 0; left < indices.Count; left++)
        {
            if (!IsFinite(gradient[left]) || indices[left] < 0 || indices[left] >= linearCost.Count)
                throw new ArgumentOutOfRangeException(nameof(gradient));
            linearCost[indices[left]] += coefficient * constant * gradient[left];
            for (int right = left; right < indices.Count; right++)
            {
                if (!IsFinite(gradient[right]) || indices[right] < 0 || indices[right] >= linearCost.Count)
                    throw new ArgumentOutOfRangeException(nameof(gradient));
                int row = Math.Min(indices[left], indices[right]);
                int column = Math.Max(indices[left], indices[right]);
                hessian.Add(row, column, coefficient * gradient[left] * gradient[right]);
            }
        }
    }

    private static double InterpolatePreviousL(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;
                if (span <= 0d)
                    return upper.LateralOffset;
                return lower.LateralOffset + (upper.LateralOffset - lower.LateralOffset) *
                    (referenceS - lower.ReferenceS) / span;
            }
        }
        return seed[seed.Count - 1].LateralOffset;
    }

    private static double GetCurvatureVariationScale(LateralPlanningInput input)
    {
        double maximum = 0d;
        for (int station = 0; station < input.ReferenceStations.Count; station++)
        {
            FrenetReferencePoint reference = ReferencePathInterpolator.Interpolate(input.ReferenceSegment,
                input.ReferenceStations[station]);
            maximum = Math.Max(maximum, Math.Abs(reference.VehicleCurvatureDerivative));
        }
        return Math.Max(1d, maximum);
    }

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

}