using System;
using System.Threading;
using MultiWheelC.TrajectoryPlanning.Utils;

namespace MultiWheelC.TrajectoryPlanning.Mapping;

/// <summary>针对行主序二值栅格计算精确欧氏距离平方的内部算法。</summary>
internal static class EuclideanDistanceTransform
{
    /// <summary>
    /// 计算每个栅格到最近障碍栅格的距离平方。
    ///
    /// 参数：occupied 为行主序占据数组，非零表示障碍；rows、cols 为数组尺寸。
    /// 返回：行主序距离平方数组，单位为栅格边长的平方；不含任何 mm 或 m 换算。
    /// </summary>
    public static double[] ComputeSquaredDistances(byte[] occupied, int rows, int cols)
    {
        if (!TryComputeSquaredDistances(occupied, rows, cols, PlanningOperationBudget.Unlimited(CancellationToken.None),
            out double[] squared, out _))
            throw new InvalidOperationException("Unbounded Euclidean distance transform unexpectedly stopped.");
        return squared;
    }

    /// <summary>使用共享预算计算距离平方；停止时不返回部分数组。</summary>
    internal static bool TryComputeSquaredDistances(byte[] occupied, int rows, int cols, PlanningOperationBudget budget,
        out double[] squaredDistances, out PlanningOperationStopReason stopReason)
    {
        if (occupied == null) throw new ArgumentNullException(nameof(occupied));
        if (budget == null) throw new ArgumentNullException(nameof(budget));
        squaredDistances = null;
        stopReason = budget.GetStopReason();
        if (stopReason != PlanningOperationStopReason.None) return false;
        double noObstacleDistanceSquared = (double)rows * rows + (double)cols * cols + 1d;
        var intermediate = new double[occupied.Length];
        var result = new double[occupied.Length];
        var input = new double[Math.Max(rows, cols)];
        var output = new double[Math.Max(rows, cols)];
        int workItemCount = 0;
        for (int row = 0; row < rows; row++)
        {
            int offset = row * cols;
            for (int col = 0; col < cols; col++)
            {
                stopReason = budget.CheckEvery(ref workItemCount);
                if (stopReason != PlanningOperationStopReason.None) return false;
                input[col] = occupied[offset + col] == 0 ? noObstacleDistanceSquared : 0d;
            }
            if (!TryTransform1D(input, cols, output, budget, ref workItemCount, out stopReason)) return false;
            for (int col = 0; col < cols; col++)
            {
                stopReason = budget.CheckEvery(ref workItemCount);
                if (stopReason != PlanningOperationStopReason.None) return false;
                intermediate[offset + col] = output[col];
            }
        }
        for (int col = 0; col < cols; col++)
        {
            for (int row = 0; row < rows; row++)
            {
                stopReason = budget.CheckEvery(ref workItemCount);
                if (stopReason != PlanningOperationStopReason.None) return false;
                input[row] = intermediate[row * cols + col];
            }
            if (!TryTransform1D(input, rows, output, budget, ref workItemCount, out stopReason)) return false;
            for (int row = 0; row < rows; row++)
            {
                stopReason = budget.CheckEvery(ref workItemCount);
                if (stopReason != PlanningOperationStopReason.None) return false;
                result[row * cols + col] = output[row];
            }
        }
        squaredDistances = result;
        stopReason = PlanningOperationStopReason.None;
        return true;
    }

    private static bool TryTransform1D(double[] f, int length, double[] result, PlanningOperationBudget budget,
        ref int workItemCount, out PlanningOperationStopReason stopReason)
    {
        var locations = new int[length];
        var boundaries = new double[length + 1];
        int k = 0;
        locations[0] = 0;
        boundaries[0] = double.NegativeInfinity;
        boundaries[1] = double.PositiveInfinity;
        for (int q = 1; q < length; q++)
        {
            stopReason = budget.CheckEvery(ref workItemCount);
            if (stopReason != PlanningOperationStopReason.None) return false;
            double intersection;
            do
            {
                int p = locations[k];
                intersection = ((f[q] + (double)q * q) - (f[p] + (double)p * p)) / (2d * (q - p));
                if (intersection <= boundaries[k])
                {
                    k--;
                    stopReason = budget.CheckEvery(ref workItemCount);
                    if (stopReason != PlanningOperationStopReason.None) return false;
                }
            } while (k >= 0 && intersection <= boundaries[k]);
            if (k < 0)
            {
                k = 0; locations[0] = q; boundaries[0] = double.NegativeInfinity; boundaries[1] = double.PositiveInfinity;
            }
            else
            {
                k++; locations[k] = q; boundaries[k] = intersection; boundaries[k + 1] = double.PositiveInfinity;
            }
        }
        k = 0;
        for (int q = 0; q < length; q++)
        {
            stopReason = budget.CheckEvery(ref workItemCount);
            if (stopReason != PlanningOperationStopReason.None) return false;
            while (boundaries[k + 1] < q) k++;
            double delta = q - locations[k];
            result[q] = delta * delta + f[locations[k]];
        }
        stopReason = PlanningOperationStopReason.None;
        return true;
    }
}