-- Sourced from: https://github.com/tdepke2/OpenComputersProjects/blob/master/tminer/dig_test.lua
-- Demo of A* algorithm for pathfinding in a 3D scalar grid.

-- Priority queue using a binary heap.
local dstructs = {}
dstructs.PriorityQueue = {}

local function heapify(arr, comp, length, i)
  local bestIndex = i
  local best = arr[i]

  if i * 2 <= length and comp(best, arr[i * 2]) then
    bestIndex = i * 2
    best = arr[bestIndex]
  end
  if i * 2 + 1 <= length and comp(best, arr[i * 2 + 1]) then
    bestIndex = i * 2 + 1
    best = arr[bestIndex]
  end
  if bestIndex ~= i then
    arr[bestIndex] = arr[i]
    arr[i] = best
    return heapify(arr, comp, length, bestIndex)
  end
end

--- dstructs.PriorityQueue:new([arr: table[, comp: function]]): table
-- Construction from an array takes O(n) time.
-- The comp is less-than by default (max heap).
function dstructs.PriorityQueue:new(arr, comp)
  self.__index = self
  self = setmetatable({}, self)

  self.comp = comp or function(a, b)
    return a < b
  end
  if arr == nil then
    self.arr = {}
    self.length = 0
  else
    self.arr = arr
    self.length = #arr
    for i = math.floor(self.length / 2), 1, -1 do
      heapify(self.arr, self.comp, self.length, i)
    end
  end

  return self
end

--- dstructs.PriorityQueue:empty(): boolean
function dstructs.PriorityQueue:empty()
  return self.length == 0
end

--- dstructs.PriorityQueue:size(): number
function dstructs.PriorityQueue:size()
  return self.length
end

--- dstructs.PriorityQueue:top(): any
function dstructs.PriorityQueue:top()
  return self.arr[1]
end

--- dstructs.PriorityQueue:push(val: any)
-- Uses heapify-up. Time complexity is O(log n)
function dstructs.PriorityQueue:push(val)
  self.length = self.length + 1
  self.arr[self.length] = val

  local i = self.length
  local parent = math.floor(i / 2)
  while i > 1 and self.comp(self.arr[parent], self.arr[i]) do
    self.arr[i], self.arr[parent] = self.arr[parent], self.arr[i]
    i = parent
    parent = math.floor(i / 2)
  end
end

--- dstructs.PriorityQueue:pop(): any
-- Uses heapify-down. Time complexity is O(log n)
function dstructs.PriorityQueue:pop()
  local topElement = self.arr[1]
  self.arr[1] = self.arr[self.length]
  self.arr[self.length] = nil
  self.length = self.length - 1

  heapify(self.arr, self.comp, self.length, 1)
  return topElement
end


local reprioritizeCounter = 0
local reprioritizeCounterFalse = 0
local reprioritizeCounterAvgIndex = 0


--- dstructs.PriorityQueue:updateKey(oldVal: any, newVal: any[, compResult: boolean]): boolean
-- Time complexity is O(n)
function dstructs.PriorityQueue:updateKey(oldVal, newVal, compResult)
  reprioritizeCounter = reprioritizeCounter + 1
  compResult = compResult or self.comp(oldVal, newVal)
  local i = 1
  while i <= self.length and self.arr[i] ~= oldVal do
    i = i + 1
  end
  reprioritizeCounterAvgIndex = reprioritizeCounterAvgIndex + i
  if i > self.length then
    --print("updateKey() false")
    reprioritizeCounterFalse = reprioritizeCounterFalse + 1
    return false
  end

  self.arr[i] = newVal
  if compResult then
    --print("updateKey() heap-up")
    local parent = math.floor(i / 2)
    while i > 1 and self.comp(self.arr[parent], self.arr[i]) do
      self.arr[i], self.arr[parent] = self.arr[parent], self.arr[i]
      i = parent
      parent = math.floor(i / 2)
    end
  else
    --print("updateKey() heap-down")
    heapify(self.arr, self.comp, self.length, i)
  end
  return true
end




local checkNeighborCounter = 0

local findPathHelper

-- findPath(grid: table, start: table, stop: table[, turnCost: number]): number,
--   table, number
--
-- Computes the lowest-cost optimal path between two points in a 3-dimensional
-- scalar field. The scalar field (grid) is a 1-indexed table sequence of
-- numbers where the indices increase in the positive x, positive z, then
-- positive y direction (same as the geolyzer scan format). The size of the grid
-- is defined by grid.xSize, grid.ySize, and grid.zSize where coordinates go
-- from 0 to grid.<n>Size - 1. The start and stop values are 3D coordinates in
-- the grid, both are tables with x, y, and z in indices 1, 2, and 3
-- respectively. Scalar values in the grid represent the cost to pass through
-- that point, the special value -1 indicates a barrier that cannot be passed
-- through.
--
-- The turnCost is an optional argument that specifies the unit path cost when
-- moving in a new direction. This is used to bias straight paths in the
-- solution when there are multiple optimal paths. By setting this to 1, the
-- bias is disabled (and a path through an area with uniform scalar values may
-- take seemingly random turns). Setting this above 1 can make straight paths
-- take priority over optimal ones, which has the potential to reduce the path
-- length in the solution.
--
-- Returns the computed path cost, table sequence of indices in the grid along
-- the path (in reverse order), and the length of this sequence. If no solution
-- exists (the goal is blocked behind barriers), then math.huge is returned for
-- the path cost.
--
-- This algorithm uses a modified version of A* search to add bias for
-- straight-line paths. See below for references:
-- https://en.wikipedia.org/wiki/A*_search_algorithm
-- https://www.redblobgames.com/pathfinding/a-star/implementation.html
local function findPath(grid, start, stop, turnCost)
  -- The heuristic used here calculates the Manhattan distance between the target and stop node.
  local function pathCostHeuristicSum(pathCost, x, y, z)
    return pathCost + math.abs(x - stop[1]) + math.abs(y - stop[2]) + math.abs(z - stop[3])
  end

  return findPathHelper(grid, start, stop, turnCost, pathCostHeuristicSum)
end


-- findPathWeighted(grid: table, start: table, stop: table[, turnCost: number,
--   suboptimalityBound: number]): number, table, number
--
-- Variation of findPath(). This uses Weighted A* to find a suboptimal path
-- between the start and stop points. The trade-off is reduced exploration of
-- nodes in the grid which therefore increases speed of the calculation. If the
-- suboptimalityBound is provided, it should be a value greater than 1. This
-- value is multiplied with the heuristic result to make it non-admissible (it
-- overestimates the distance to the goal). As the suboptimalityBound approaches
-- infinity, the search turns into a greedy best-first search.
--
-- See here for details and more variations:
-- http://theory.stanford.edu/~amitp/GameProgramming/Variations.html
local function findPathWeighted(grid, start, stop, turnCost, suboptimalityBound)
  suboptimalityBound = suboptimalityBound or 1.5
  local function pathCostHeuristicSum(pathCost, x, y, z)
    return pathCost + (math.abs(x - stop[1]) + math.abs(y - stop[2]) + math.abs(z - stop[3])) * suboptimalityBound
  end

  return findPathHelper(grid, start, stop, turnCost, pathCostHeuristicSum)
end


-- findPathPWXD(grid: table, start: table, stop: table[, turnCost: number,
--   suboptimalityBound: number]): number, table, number
--
-- Another variation of findPath(). This uses Weighted A* (piecewise Convex
-- Downward) to find a suboptimal path between the start and stop points. This
-- often gives better speed over Weighted A* because the scaling of the path
-- cost is applied as we get closer to the goal. If the suboptimalityBound is
-- provided, it should be a value greater than 1.
--
-- See here for demonstration and implementation details:
-- https://www.movingai.com/SAS/SUB/
-- https://webdocs.cs.ualberta.ca/~nathanst/papers/chen2021general.pdf
local function findPathPWXD(grid, start, stop, turnCost, suboptimalityBound)
  suboptimalityBound = suboptimalityBound or 1.5
  local function pathCostHeuristicSum(pathCost, x, y, z)
    local h = math.abs(x - stop[1]) + math.abs(y - stop[2]) + math.abs(z - stop[3])
    if h > pathCost then
      return pathCost + h
    else
      return (pathCost + (2 * suboptimalityBound - 1) * h) / suboptimalityBound
    end
  end

  return findPathHelper(grid, start, stop, turnCost, pathCostHeuristicSum)
end


-- Helper function for findPath() and variations.
findPathHelper = function(grid, start, stop, turnCost, pathCostHeuristicSum)
  turnCost = turnCost or 1.0001

  -- Find indices of the start and stop positions in the grid.
  local startNode = (start[2] * grid.zSize + start[3]) * grid.xSize + start[1] + 1
  local stopNode = (stop[2] * grid.zSize + stop[3]) * grid.xSize + stop[1] + 1

  -- For node n, cameFrom[n] is the node immediately preceding it on the cheapest path to get to n.
  local cameFrom = {}

  -- For node n, pathCosts[n] is the cost of the cheapest path from start to n.
  local pathCosts = setmetatable({[startNode] = 0}, {
    __index = function()
      return math.huge
    end
  })

  -- For node n, nodePriorities[n] is the current best guess of the total path cost to get to the goal when going through n.
  local nodePriorities = setmetatable({[startNode] = 0}, {
    __index = function()
      return math.huge
    end
  })

  -- Nodes that have been discovered and need to be expanded. The priority queue keeps the node with the smallest estimated total path cost at the front.
  -- A node that is removed from the queue may be added again later, or a node may change priority from within the queue.
  local fringeNodes = dstructs.PriorityQueue:new({startNode}, function(a, b) return nodePriorities[a] > nodePriorities[b] end)

  -- Iterates from stopNode to startNode following the trail left in cameFrom.
  -- The solution (which is the reversed path) is then returned.
  local function reconstructPath()
    local pathReversed, pathReversedSize = {}, 0
    local node = stopNode
    while node do
      pathReversedSize = pathReversedSize + 1
      pathReversed[pathReversedSize] = node
      node = cameFrom[node]
    end
    return pathCosts[stopNode], pathReversed, pathReversedSize
  end

  -- Checks a node adjacent to current to see if it corresponds to an unexplored node, or if it improves the path cost of an existing one.
  local function checkNeighbor(current, x, y, z, baseMovementCost)
    local neighbor = (y * grid.zSize + z) * grid.xSize + x + 1
    -- Return early if a barrier is found, or the neighbor is in the reverse direction.
    if grid[neighbor] == -1 or neighbor == cameFrom[current] then
      return
    end
    checkNeighborCounter = checkNeighborCounter + 1
    local oldPathCost = pathCosts[neighbor]
    local newPathCost = pathCosts[current] + baseMovementCost + grid[neighbor]
    --io.write("  checkNeighbor(" .. current .. ", " .. x .. ", " .. y .. ", " .. z .. "), newPathCost = " .. newPathCost .. ", oldPathCost = " .. pathCosts[neighbor])
    if newPathCost < oldPathCost then
      --io.write(" (greater)\n")
      --if oldPathCost ~= math.huge then
        --io.write("score was increased for existing node!\n")
      --end
      cameFrom[neighbor] = current
      pathCosts[neighbor] = newPathCost
      nodePriorities[neighbor] = pathCostHeuristicSum(newPathCost, x, y, z)

      -- The neighbor is added to fringe if it hasn't been found before, or it has and is not currently in the fringe.
      -- Another way to do this is to always add it (allow duplicates), but this doesn't work so well when we store the priorities outside of the queue (the heap can get invalidated).
      if oldPathCost == math.huge or not fringeNodes:updateKey(neighbor, neighbor, true) then
        fringeNodes:push(neighbor)
      end
    --else
      --io.write(" (less)\n")
    end
  end

  while not fringeNodes:empty() do
    -- Get node at the front of queue (smallest total estimated cost) and explore the neighbor nodes if it's not the goal.
    local current = fringeNodes:pop()
    if current == stopNode then
      return reconstructPath()
    end
    --io.write("current = " .. current .. ", (" .. x .. ", " .. y .. ", " .. z .. "), g = " .. pathCosts[current] .. ", f = " .. nodePriorities[current] .. "\n")

    -- Prefer straight line paths by adding a small path cost when turning.
    local xMovementCost, yMovementCost, zMovementCost = turnCost, 1, turnCost
    if cameFrom[current] then
      if math.abs(current - cameFrom[current]) == 1 then
        xMovementCost = 1
        yMovementCost = turnCost
      elseif math.abs(current - cameFrom[current]) == grid.xSize then
        yMovementCost = turnCost
        zMovementCost = 1
      end
    end

    local divisor = (current - 1) / grid.xSize
    local x = (current - 1) % grid.xSize
    local z = math.floor(divisor) % grid.zSize
    local y = math.floor(divisor / grid.zSize)

    if x < grid.xSize - 1 then
      checkNeighbor(current, x + 1, y, z, xMovementCost)
    end
    if x > 0 then
      checkNeighbor(current, x - 1, y, z, xMovementCost)
    end
    if y < grid.ySize - 1 then
      checkNeighbor(current, x, y + 1, z, yMovementCost)
    end
    if y > 0 then
      checkNeighbor(current, x, y - 1, z, yMovementCost)
    end
    if z < grid.zSize - 1 then
      checkNeighbor(current, x, y, z + 1, zMovementCost)
    end
    if z > 0 then
      checkNeighbor(current, x, y, z - 1, zMovementCost)
    end
  end
  return math.huge, {}, 0
end

-- Draw the layers of the 3D grid for visualization purposes.
local function printGrid(grid, paths)
  io.write("\ngrid.xSize = ", tostring(grid.xSize), ", grid.ySize = ", tostring(grid.ySize), ", grid.zSize = ", tostring(grid.zSize), "\n\n")
  local i = 1
  for y = 0, grid.ySize - 1 do
    print("layer y = " .. y)
    for z = 0, grid.zSize - 1 do
      for x = 0, grid.xSize - 1 do
        if grid[i] == -1 then
          io.write("██")
        else
          io.write(paths[i] or " ")
          io.write(tostring(grid[i]))
        end
        i = i + 1
      end
      io.write("\n")
    end
    io.write("\n")
  end
  io.write("\n")
end




-- Create the 3D grid and fill it with random weights and barriers (the value -1 marks a barrier).
local grid = {}
grid.xSize = 32
grid.ySize = 4
grid.zSize = 16
local paths = {}
math.randomseed(456)
for i = 1, grid.xSize * grid.ySize * grid.zSize do
  local num = math.random()
  grid[i] = num < 0.1 and -1 or math.random(0, 5)
end

-- Define start/stop points and other parameters for the search.
local start = {0, 0, 0}
local stop = {grid.xSize-1, grid.ySize-1, grid.zSize-1}
local w = 1.5
print("finding path from (" .. start[1] .. ", " .. start[2] .. ", " .. start[3] .. ") to (" .. stop[1] .. ", " .. stop[2] .. ", " .. stop[3] .. ")")

-- Find the path.
local pathCost, pathReversed, pathReversedSize = findPath(grid, start, stop)
--local pathCost, pathReversed, pathReversedSize = findPathWeighted(grid, start, stop, 1.0001, w)
--local pathCost, pathReversed, pathReversedSize = findPathPWXD(grid, start, stop, 1.0001, w)

print("pathCost = ", pathCost)
if pathCost ~= math.huge then
  io.write("pathReversed (indices in grid):\n")
  for i, node in ipairs(pathReversed) do
    io.write(node .. "  ")
    assert(grid[node] >= 0 and grid[node] < 10 and not paths[node])
    paths[node] = "#"
  end
  assert(paths[pathReversed[pathReversedSize]] == "#" and paths[pathReversed[1]] == "#")
  paths[pathReversed[pathReversedSize]] = "A"
  paths[pathReversed[1]] = "B"
  io.write("\n")
else
  paths[(start[2] * grid.zSize + start[3]) * grid.xSize + start[1] + 1] = "A"
  paths[(stop[2] * grid.zSize + stop[3]) * grid.xSize + stop[1] + 1] = "B"
  print("no path!\n")
end
-- Debugging info:
--reprioritizeCounterAvgIndex = reprioritizeCounterAvgIndex / reprioritizeCounter
--print("reprioritizeCounter = " .. reprioritizeCounter .. ", reprioritizeCounterFalse = " .. reprioritizeCounterFalse .. ", reprioritizeCounterAvgIndex = " .. reprioritizeCounterAvgIndex)
--print("checkNeighborCounter = " .. checkNeighborCounter .. "\n")

printGrid(grid, paths)