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

namespace MultiWheelC.TrajectoryPlanning.CoarsePath.Search;

/// <summary>
/// 基于不可变栅格地图的目标反向八邻域 Dijkstra 启发式。
/// 距离单位为 m；对角移动仅在两个对应正交邻格都未占据时允许，以避免从障碍夹角穿越。
/// </summary>
public sealed class GridDijkstraHeuristic
{
    private readonly PlanningGridMap _map;
    private readonly double[] _costs;

    /// <summary>
    /// 从目标栅格预计算所有可达自由格到目标的二维最短距离。
    /// 参数：map 必须是已就绪的不可变地图；goalRow、goalCol 为地图内且未占据的目标格索引。
    /// 失败：地图为空、未就绪或目标格无效时抛出异常。
    /// </summary>
    public GridDijkstraHeuristic(PlanningGridMap map, int goalRow, int goalCol)
    {
        if (!TryCreate(map, goalRow, goalCol, PlanningOperationBudget.Unlimited(CancellationToken.None),
            out GridDijkstraHeuristic heuristic, out _))
            throw new InvalidOperationException("Unbounded Dijkstra construction unexpectedly stopped.");
        _map = heuristic._map;
        _costs = heuristic._costs;
    }

    private GridDijkstraHeuristic(PlanningGridMap map, double[] costs)
    {
        _map = map;
        _costs = costs;
    }

    /// <summary>使用共享预算创建完整二维启发式；停止时不返回部分成本数组。</summary>
    internal static bool TryCreate(PlanningGridMap map, int goalRow, int goalCol, PlanningOperationBudget budget,
        out GridDijkstraHeuristic heuristic, out PlanningOperationStopReason stopReason)
    {
        if (map == null) throw new ArgumentNullException(nameof(map));
        if (!map.PlanningReady) throw new ArgumentException("The planning map must be ready.", nameof(map));
        if (goalRow < 0 || goalRow >= map.Rows || goalCol < 0 || goalCol >= map.Cols)
            throw new ArgumentOutOfRangeException(nameof(goalRow));
        if (map.IsOccupied(goalRow, goalCol))
            throw new ArgumentException("The goal grid cell must be free.", nameof(goalRow));
        if (budget == null) throw new ArgumentNullException(nameof(budget));

        heuristic = null;
        stopReason = budget.GetStopReason();
        if (stopReason != PlanningOperationStopReason.None) return false;
        var costs = new double[checked(map.Rows * map.Cols)];
        int workItemCount = 0;
        for (int index = 0; index < costs.Length; index++)
        {
            stopReason = budget.CheckEvery(ref workItemCount);
            if (stopReason != PlanningOperationStopReason.None) return false;
            costs[index] = double.PositiveInfinity;
        }
        if (!TryBuild(map, costs, goalRow, goalCol, budget, ref workItemCount, out stopReason)) return false;
        heuristic = new GridDijkstraHeuristic(map, costs);
        stopReason = PlanningOperationStopReason.None;
        return true;
    }

    /// <summary>
    /// 查询指定栅格到构造时目标格的二维最短距离。
    /// 参数：row、col 为从零开始的栅格索引。
    /// 返回：单位 m 的有限最短距离；自由格不可达或索引越界时返回正无穷。
    /// </summary>
    public double GetCost(int row, int col)
    {
        return row < 0 || row >= _map.Rows || col < 0 || col >= _map.Cols
            ? double.PositiveInfinity
            : _costs[row * _map.Cols + col];
    }

    private static bool TryBuild(PlanningGridMap map, double[] costs, int goalRow, int goalCol,
        PlanningOperationBudget budget, ref int workItemCount, out PlanningOperationStopReason stopReason)
    {
        var openList = new BinaryMinHeap<int>();
        int goalIndex = goalRow * map.Cols + goalCol;
        costs[goalIndex] = 0d;
        openList.Push(goalIndex, 0d, 0d, 0d);

        while (openList.Count > 0)
        {
            stopReason = budget.CheckEvery(ref workItemCount);
            if (stopReason != PlanningOperationStopReason.None) return false;
            int currentIndex = openList.Pop();
            int currentRow = currentIndex / map.Cols;
            int currentCol = currentIndex % map.Cols;
            double currentCost = costs[currentIndex];

            for (int rowOffset = -1; rowOffset <= 1; rowOffset++)
            for (int colOffset = -1; colOffset <= 1; colOffset++)
            {
                stopReason = budget.CheckEvery(ref workItemCount);
                if (stopReason != PlanningOperationStopReason.None) return false;
                if (rowOffset == 0 && colOffset == 0) continue;

                int nextRow = currentRow + rowOffset;
                int nextCol = currentCol + colOffset;
                if (nextRow < 0 || nextRow >= map.Rows || nextCol < 0 || nextCol >= map.Cols || map.IsOccupied(nextRow, nextCol))
                    continue;

                bool isDiagonal = rowOffset != 0 && colOffset != 0;
                if (isDiagonal && (map.IsOccupied(currentRow + rowOffset, currentCol) || map.IsOccupied(currentRow, currentCol + colOffset)))
                    continue;

                double stepCost = isDiagonal ? Math.Sqrt(2d) * map.ResolutionMeters : map.ResolutionMeters;
                double candidateCost = currentCost + stepCost;
                int nextIndex = nextRow * map.Cols + nextCol;
                if (candidateCost >= costs[nextIndex]) continue;

                costs[nextIndex] = candidateCost;
                openList.Push(nextIndex, candidateCost, 0d, 0d);
            }
        }
        stopReason = PlanningOperationStopReason.None;
        return true;
    }
}
