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(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);
        CurvatureAffine[] curvature = CreateCurvatureAffines(input, layout, linearization);
        for (int station = 0; station < curvature.Length; station++)
        {
            AddSquaredResidual(hessian, linearCost, curvature[station].Indices, 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 CurvatureAffine[] CreateCurvatureAffines(LateralPlanningInput input, LateralVariableLayout layout,
        LateralCandidate linearization)
    {
        var affines = new CurvatureAffine[layout.StationCount];
        double directionSign = input.ReferenceSegment.Direction == TravelDirection.Forward ? 1d : -1d;
        for (int station = 0; station < layout.StationCount; station++)
        {
            FrenetReferencePoint reference = ReferencePathInterpolator.Interpolate(input.ReferenceSegment,
                input.ReferenceStations[station]);
            double l = linearization.L[station];
            double dl = linearization.DL[station];
            double ddl = linearization.DDL[station];
            double referenceCurvature = reference.GeometricCurvature;
            double referenceCurvatureDerivative = directionSign * reference.VehicleCurvatureDerivative;
            double a = 1d - referenceCurvature * l;
            double denominatorSquared = a * a + dl * dl;
            if (!IsFinite(denominatorSquared) || denominatorSquared <= 0d)
                throw new ArgumentException("Curvature linearization denominator is invalid.", nameof(linearization));

            double denominatorPow3Over2 = denominatorSquared * Math.Sqrt(denominatorSquared);
            double denominatorPow5Over2 = denominatorPow3Over2 * denominatorSquared;
            double numerator = a * a * referenceCurvature + a * ddl +
                referenceCurvatureDerivative * l * dl + 2d * referenceCurvature * dl * dl;
            double geometricCurvature = numerator / denominatorPow3Over2;
            double dNumeratorDLateral = -2d * a * referenceCurvature * referenceCurvature -
                referenceCurvature * ddl + referenceCurvatureDerivative * dl;
            double dNumeratorDSlope = referenceCurvatureDerivative * l + 4d * referenceCurvature * dl;
            double dDenominatorSquaredDLateral = -2d * a * referenceCurvature;
            double dDenominatorSquaredDSlope = 2d * dl;
            double dGeometricDLateral = dNumeratorDLateral / denominatorPow3Over2 -
                1.5d * numerator * dDenominatorSquaredDLateral / denominatorPow5Over2;
            double dGeometricDSlope = dNumeratorDSlope / denominatorPow3Over2 -
                1.5d * numerator * dDenominatorSquaredDSlope / denominatorPow5Over2;
            double dGeometricDSecondDerivative = a / denominatorPow3Over2;
            double vehicleCurvature = directionSign * geometricCurvature;
            double[] gradient =
            {
                directionSign * dGeometricDLateral,
                directionSign * dGeometricDSlope,
                directionSign * dGeometricDSecondDerivative,
            };
            double constant = vehicleCurvature - gradient[0] * l - gradient[1] * dl - gradient[2] * ddl;
            if (!IsFinite(vehicleCurvature) || !IsFinite(constant) || !IsFinite(gradient[0]) ||
                !IsFinite(gradient[1]) || !IsFinite(gradient[2]))
            {
                throw new ArgumentException("Curvature linearization is non-finite.", nameof(linearization));
            }
            affines[station] = new CurvatureAffine(new[] { layout.L(station), layout.DL(station), layout.DDL(station) },
                gradient, constant);
        }
        return affines;
    }

    private static void AddCurvatureVariationTerms(IReadOnlyList<double> stations, CurvatureAffine[] curvature,
        SparseTripletBuilder hessian, IList<double> linearCost, double weight, double scale)
    {
        for (int station = 0; station < curvature.Length; station++)
        {
            int lower = station == 0 ? 0 : station - 1;
            int upper = station == curvature.Length - 1 ? curvature.Length - 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));
            CurvatureAffine left = curvature[lower];
            CurvatureAffine right = curvature[upper];
            var indices = new int[left.Indices.Length + right.Indices.Length];
            var gradient = new double[indices.Length];
            for (int index = 0; index < left.Indices.Length; index++)
            {
                indices[index] = left.Indices[index];
                gradient[index] = -left.Gradient[index] / ds;
                indices[left.Indices.Length + index] = right.Indices[index];
                gradient[left.Indices.Length + 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 GetMaximumVehicleCurvature(VehicleParameters vehicle)
    {
        if (vehicle == null)
            throw new ArgumentNullException(nameof(vehicle));
        if (vehicle.MaximumCurvaturePerMeter.HasValue)
            return vehicle.MaximumCurvaturePerMeter.Value;
        if (vehicle.MinimumTurningRadiusMeters.HasValue && vehicle.MinimumTurningRadiusMeters.Value > 0d)
            return 1d / vehicle.MinimumTurningRadiusMeters.Value;
        throw new ArgumentException("Vehicle maximum curvature is required.", nameof(vehicle));
    }

    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);
    }

    private sealed class CurvatureAffine
    {
        public CurvatureAffine(int[] indices, double[] gradient, double constant)
        {
            Indices = indices;
            Gradient = gradient;
            Constant = constant;
        }

        public int[] Indices { get; }
        public double[] Gradient { get; }
        public double Constant { get; }
    }
}