using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;

namespace MultiWheelC.TrajectoryPlanning.EMPlanner;

internal sealed class TrajectorySampleSchedule
{
    private const double ZeroTolerance = 1e-12d;

    public TrajectorySampleSchedule(LongitudinalCandidate candidate, double outputTimeStepSeconds, double holdDurationSeconds,
        EmLongitudinalMode mode, bool resampleMotion)
    {
        if (candidate == null)
            throw new ArgumentNullException(nameof(candidate));
        if (!IsFinite(outputTimeStepSeconds) || outputTimeStepSeconds <= 0d)
            throw new ArgumentOutOfRangeException(nameof(outputTimeStepSeconds));
        if (!IsFinite(holdDurationSeconds) || holdDurationSeconds < 0d)
            throw new ArgumentOutOfRangeException(nameof(holdDurationSeconds));
        if (!Enum.IsDefined(typeof(EmLongitudinalMode), mode))
            throw new ArgumentOutOfRangeException(nameof(mode));

        var samples = new List<TrajectorySample>(candidate.KnotTimes.Count + 4);
        double previousPathS = double.NegativeInfinity;
        if (resampleMotion)
        {
            double finalTime = candidate.KnotTimes[candidate.KnotTimes.Count - 1];
            int sourceInterval = 0;
            for (double sampleTime = 0d; sampleTime < finalTime - ZeroTolerance;
                sampleTime += outputTimeStepSeconds)
            {
                AddSample(Interpolate(candidate, sampleTime, ref sourceInterval), samples, ref previousPathS, candidate);
            }
            AddSample(Interpolate(candidate, finalTime, ref sourceInterval), samples, ref previousPathS, candidate);
        }
        else
        {
            for (int index = 0; index < candidate.KnotTimes.Count; index++)
            {
                AddSample(new TrajectorySample(candidate.KnotTimes[index], candidate.S[index],
                    Math.Max(0d, candidate.U[index]), candidate.A[index], index < candidate.J.Count ? candidate.J[index] : 0d,
                    false), samples, ref previousPathS, candidate);
            }
        }

        if (mode != EmLongitudinalMode.ExactStopAtBoundary)
        {
            Samples = new ReadOnlyCollection<TrajectorySample>(samples);
            TerminalAnchorSampleIndex = -1;
            return;
        }

        int sourceStabilizationStart = resampleMotion
            ? FindTerminalStationaryTailStart(candidate)
            : LongitudinalTerminalSchedule.GetStabilizationStartIndex(candidate.KnotTimes, outputTimeStepSeconds);
        double stabilizationStartTime = candidate.KnotTimes[sourceStabilizationStart];
        int stabilizationStart = 0;
        while (stabilizationStart < samples.Count - 1 &&
            samples[stabilizationStart].TimeFromStart < stabilizationStartTime - ZeroTolerance)
        {
            stabilizationStart++;
        }
        double stopPathS = samples[stabilizationStart].PathS;
        for (int index = stabilizationStart; index < samples.Count; index++)
        {
            if (Math.Abs(samples[index].PathS - stopPathS) > ZeroTolerance ||
                Math.Abs(samples[index].ProgressSpeed) > ZeroTolerance || Math.Abs(samples[index].Acceleration) > ZeroTolerance ||
                Math.Abs(samples[index].Jerk) > ZeroTolerance)
            {
                throw new ArgumentException("An exact stop requires a stationary S/U/A/J tail.", nameof(candidate));
            }
        }

        TerminalAnchorSampleIndex = stabilizationStart;
        TrajectorySample terminal = samples[samples.Count - 1];
        double holdElapsed = 0d;
        while (holdElapsed < holdDurationSeconds - ZeroTolerance)
        {
            holdElapsed = Math.Min(holdDurationSeconds, holdElapsed + outputTimeStepSeconds);
            samples.Add(new TrajectorySample(terminal.TimeFromStart + holdElapsed, terminal.PathS, 0d, 0d, 0d, true));
        }

        Samples = new ReadOnlyCollection<TrajectorySample>(samples);
    }

    public IReadOnlyList<TrajectorySample> Samples { get; }
    public int TerminalAnchorSampleIndex { get; }

    internal static bool ExceedsMaximumSampleCount(LongitudinalCandidate candidate, double outputTimeStepSeconds,
        double holdDurationSeconds, EmLongitudinalMode mode, bool resampleMotion, int maximumSampleCount,
        out int requiredSampleCount)
    {
        requiredSampleCount = 0;
        if (candidate == null)
            throw new ArgumentNullException(nameof(candidate));
        if (maximumSampleCount < 1)
            throw new ArgumentOutOfRangeException(nameof(maximumSampleCount));

        if (resampleMotion)
        {
            double finalTime = candidate.KnotTimes[candidate.KnotTimes.Count - 1];
            for (double sampleTime = 0d; sampleTime < finalTime - ZeroTolerance;
                sampleTime += outputTimeStepSeconds)
            {
                if (IncrementAndExceedsMaximum(ref requiredSampleCount, maximumSampleCount))
                    return true;
            }
            if (IncrementAndExceedsMaximum(ref requiredSampleCount, maximumSampleCount))
                return true;
        }
        else
        {
            for (int index = 0; index < candidate.KnotTimes.Count; index++)
            {
                if (IncrementAndExceedsMaximum(ref requiredSampleCount, maximumSampleCount))
                    return true;
            }
        }

        if (mode != EmLongitudinalMode.ExactStopAtBoundary)
            return false;

        double holdElapsed = 0d;
        while (holdElapsed < holdDurationSeconds - ZeroTolerance)
        {
            holdElapsed = Math.Min(holdDurationSeconds, holdElapsed + outputTimeStepSeconds);
            if (IncrementAndExceedsMaximum(ref requiredSampleCount, maximumSampleCount))
                return true;
        }
        return false;
    }

    private static void AddSample(TrajectorySample sample, ICollection<TrajectorySample> samples,
        ref double previousPathS, LongitudinalCandidate candidate)
    {
        if (sample.PathS < previousPathS)
            throw new ArgumentException("Trajectory PathS cannot decrease.", nameof(candidate));
        if (sample.ProgressSpeed < -ZeroTolerance)
            throw new ArgumentException("Longitudinal progress speed cannot be negative.", nameof(candidate));
        samples.Add(sample);
        previousPathS = sample.PathS;
    }

    private static bool IncrementAndExceedsMaximum(ref int sampleCount, int maximumSampleCount)
    {
        checked { sampleCount++; }
        return sampleCount > maximumSampleCount;
    }

    private static TrajectorySample Interpolate(LongitudinalCandidate candidate, double sampleTime, ref int sourceInterval)
    {
        int lastKnot = candidate.KnotTimes.Count - 1;
        if (sampleTime >= candidate.KnotTimes[lastKnot] - ZeroTolerance)
        {
            return new TrajectorySample(candidate.KnotTimes[lastKnot], candidate.S[lastKnot],
                Math.Max(0d, candidate.U[lastKnot]), candidate.A[lastKnot], 0d, false);
        }
        while (sourceInterval < lastKnot - 1 &&
            sampleTime >= candidate.KnotTimes[sourceInterval + 1] - ZeroTolerance)
        {
            sourceInterval++;
        }
        if (Math.Abs(sampleTime - candidate.KnotTimes[sourceInterval]) <= ZeroTolerance)
        {
            return new TrajectorySample(candidate.KnotTimes[sourceInterval], candidate.S[sourceInterval],
                Math.Max(0d, candidate.U[sourceInterval]), candidate.A[sourceInterval], candidate.J[sourceInterval], false);
        }
        double dt = sampleTime - candidate.KnotTimes[sourceInterval];
        double jerk = candidate.J[sourceInterval];
        double acceleration = candidate.A[sourceInterval] + jerk * dt;
        double speed = candidate.U[sourceInterval] + candidate.A[sourceInterval] * dt + 0.5d * jerk * dt * dt;
        double pathS = candidate.S[sourceInterval] + candidate.U[sourceInterval] * dt +
            0.5d * candidate.A[sourceInterval] * dt * dt + jerk * dt * dt * dt / 6d;
        return new TrajectorySample(sampleTime, pathS, Math.Max(0d, speed), acceleration, jerk, false);
    }

    private static int FindTerminalStationaryTailStart(LongitudinalCandidate candidate)
    {
        int start = candidate.KnotTimes.Count - 1;
        double terminalPathS = candidate.S[start];
        while (start > 0 && Math.Abs(candidate.S[start - 1] - terminalPathS) <= ZeroTolerance &&
            Math.Abs(candidate.U[start - 1]) <= ZeroTolerance && Math.Abs(candidate.A[start - 1]) <= ZeroTolerance &&
            Math.Abs(candidate.J[start - 1]) <= ZeroTolerance)
        {
            start--;
        }
        return start;
    }

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

internal sealed class TrajectorySample
{
    public TrajectorySample(double timeFromStart, double pathS, double progressSpeed, double acceleration, double jerk,
        bool isHoldSample)
    {
        TimeFromStart = timeFromStart;
        PathS = pathS;
        ProgressSpeed = progressSpeed;
        Acceleration = acceleration;
        Jerk = jerk;
        IsHoldSample = isHoldSample;
    }

    public double TimeFromStart { get; }
    public double PathS { get; }
    public double ProgressSpeed { get; }
    public double Acceleration { get; }
    public double Jerk { get; }
    public bool IsHoldSample { get; }
}