// SPDX-License-Identifier: Apache-2.0

package maze

import (
	"cmp"
	"fmt"
	"image"
)

type problemCellState uint8

const (
	// zero value, for unvisited cells.
	stateUnvisited problemCellState = iota

	// we considered the cell and set its heuristic, but then the movement
	// test failed, we don't know how to reach it still. distance is still uninitialized.
	stateUnexplored

	// this is the starting cell.
	// distance is 0.
	stateStart

	// the cell was reached via the respective direction.
	// distance is the number of cells traversed to reach this cell.
	stateRight
	stateDown
	stateLeft
	stateUp

	// the number of directions, for array bounds.
	// not used as a state itself.
	stateLen
)

var stateStrings = [stateLen]string{
	"unvisited",
	"unexplored",
	"start",
	"right",
	"down",
	"left",
	"up",
}

// lookup table to reverse the direction.
//
// used to backtrack from the end to the start once the maze is solved.
var dirReverse = [stateLen]problemCellState{
	stateUnvisited,  // stateUnvisited
	stateUnexplored, // stateUnexplored
	stateStart,      // stateStart
	stateLeft,       // stateRight
	stateUp,         // stateDown
	stateRight,      // stateLeft
	stateDown,       // stateUp
}

// lookup table to advance in the given direction.
var dirAdvance = [stateLen]image.Point{
	{0, 0},  // stateUnvisited
	{0, 0},  // stateUnexplored
	{0, 0},  // stateStart
	{1, 0},  // stateRight
	{0, 1},  // stateDown
	{-1, 0}, // stateLeft
	{0, -1}, // stateUp
}

// lookup table to offset a point during movement tests
// for left and up, we need to test the cell to the left or above.
var dirTestOffset = [stateLen]image.Point{
	{0, 0},  // stateUnvisited
	{0, 0},  // stateUnexplored
	{0, 0},  // stateStart
	{0, 0},  // stateRight
	{0, 0},  // stateDown
	{-1, 0}, // stateLeft
	{0, -1}, // stateUp
}

func (d problemCellState) String() string {
	if d < stateLen {
		return stateStrings[d]
	}

	// The Stringer is the only case where we do bounds checking, the other functions can just panic.
	return fmt.Sprintf("unknown(%d)", d)
}

type problemCell struct {
	state problemCellState

	// the number of cells traversed to reach this cell.
	// only valid if state for stateStart, stateRight, stateDown, stateLeft, or stateUp.
	distance uint

	// the optimistic distance estimate to the end point.
	// valid for any state except stateUnvisited.
	heuristic uint
}

// A Problem is intended to be used by the Solve function to find a path through a maze.
//
// This is more of a demonstration of how to use the Solve function, a dedicated maze solving
// algorithm would more efficient. In particular, the solver manages a max heap for discarding
// states when at capacity, which a maze solver would not need - a maze has a finite number of cells,
// which for our purposes represent a state, and we'd never exceed (and thus need to keep track
// of and discard) that many states.
type Problem struct {
	// the bounds of the maze.
	size image.Rectangle

	// a lookup table for the bounds that a cell can be transverse in.
	//
	// for the movement directions, this will be identical to size, but contracted by one cell.
	// for example, legalPositions[stateRight] will have the rightmost column removed, as you can't
	// move further right from the rightmost column.
	//
	// for other possible states, the rectangle will be empty (and realistically, never used).
	legalPositions [stateLen]image.Rectangle

	// the functions that test if movement is allowed in the given direction.
	// this needs to be combined with dirTestOffset to get the correct cell to test.
	testFuncs [stateLen]func(image.Point) bool

	// the start and end points of the maze.
	start, end image.Point

	// the cells in the maze, starting at size.Min, going left to right, top to bottom.
	cells []problemCell
}

// convert a coordinate to an index in the cells slice.
//
// it is assumed that c.In(p.size) is true, otherwise this
// will yield a wrong, and possibly out of bounds index.
func (p Problem) coordToIndex(c image.Point) int {
	return (c.X - p.size.Min.X) + (c.Y-p.size.Min.Y)*p.size.Dx()
}

func absDelta(a, b int) uint {
	if a > b {
		return uint(a - b)
	}
	return uint(b - a)
}

// an optimistic distance estimate based on the Manhattan distance to the end point.
func (p Problem) heuristic(c image.Point) uint {
	return absDelta(c.X, p.end.X) + absDelta(c.Y, p.end.Y)
}

// Initializes the maze search problem, and appends the start point to the given slice.
func (p *Problem) Initialize(out []image.Point) ([]image.Point, error) {
	p.cells = make([]problemCell, p.size.Dx()*p.size.Dy())

	p.cells[p.coordToIndex(p.start)] = problemCell{
		state:     stateStart,
		heuristic: p.heuristic(p.end),
	}

	return append(out, p.start), nil
}

// Appends the next positions to consider from the given position.
func (p Problem) Next(c image.Point, out []image.Point) ([]image.Point, error) {
	distance := p.cells[p.coordToIndex(c)].distance + 1

	for _, d := range [...]problemCellState{stateRight, stateDown, stateLeft, stateUp} {
		if !c.In(p.legalPositions[d]) {
			// can't move in this direction, stop immediately.
			continue
		}

		next := c.Add(dirAdvance[d])
		nextIndex := p.coordToIndex(next)

		switch {
		case p.cells[nextIndex].state == stateUnvisited:
			// calculate the heuristic for the first time and change its state to unexplored.
			p.cells[nextIndex].heuristic = p.heuristic(next)
			p.cells[nextIndex].state = stateUnexplored
			fallthrough

		case p.cells[nextIndex].state == stateUnexplored:
			// similar to the unvisited case, but don't calculate the heuristic again.
			fallthrough

		// if the cell has been visited, but we found a shorter path to it, then
		// update the direction and distance.
		case distance < p.cells[nextIndex].distance:
			// if we reach this point, then one of the above cases has occurred, and
			// this is a cell we're interested in exploring further.
			//
			// test to make sure we can actually move there.
			if p.testFuncs[d](c.Add(dirTestOffset[d])) {
				p.cells[nextIndex].state = d
				p.cells[nextIndex].distance = distance
				out = append(out, next)
			}
		}
	}

	return out, nil
}

// Returns true if the given position is the end of the maze.
//
// We're also assuming that the point was returned by Initialize or Next,
// otherwise we wouldn't know how to reach the start position, and this
// wouldn't actually be solved.
func (p Problem) Solved(c image.Point) bool {
	return c == p.end
}

// Converts the given position, the one that passed the Solved test above,
// into a list of points that form a path from the start to the end.
func (p Problem) Finish(c image.Point) ([]image.Point, error) {
	if c != p.end {
		return nil, fmt.Errorf("not at the end: %v", c)
	}

	switch p.cells[p.coordToIndex(c)].state {
	case stateUnvisited, stateUnexplored:
		return nil, fmt.Errorf("no path to the end")
	}

	length := p.cells[p.coordToIndex(c)].distance
	path := make([]image.Point, length+1)
	path[length] = c

	for ; length > 0; length-- {
		c = c.Add(dirAdvance[dirReverse[p.cells[p.coordToIndex(c)].state]])
		path[length-1] = c
	}

	return path, nil
}

// This function would be used to release any resources associated with a given search state,
// but since a state for this problem is just a point, and all the actual state information
// is stored in the cells slice, there's nothing to do here.
func (p Problem) Discard(image.Point) {
	// nop
}

// Returns true if state a should be explored before b.
func (p Problem) OptimisticLess(a, b image.Point) bool {
	ca := p.cells[p.coordToIndex(a)]
	cb := p.cells[p.coordToIndex(b)]

	if r := cmp.Compare(ca.distance+ca.heuristic, cb.distance+cb.heuristic); r != 0 {
		return r < 0
	}

	// If the estimated total path length for both states seems to be equal,
	// prefer the one with the lower heuristic (closest to the end).
	//
	// since distance+heuristic for both states is equal, this conversely means that
	// we prefer the one with the higher distance traveled so far.
	return ca.heuristic < cb.heuristic
}

// Similar to OptimisticLess, but we scale the heuristic to penalize uncertainty
// about the actual remaining distance.
//
// Note that the solver uses this with a max heap to determine which states to prune when at capacity.
//
// This will typically be similar to OptimisticLess, but with a penalty for uncertainty.
//
// For our purposes, we're just going to use the same heuristic as OptimisticLess. Assuming that you give
// the solver a capacity equal to the number of cells in the maze, it won't need to prune any states.
func (p Problem) PessimisticLess(a, b image.Point) bool {
	return p.OptimisticLess(a, b)
}

// NewMazeProblem creates a new maze problem with the given size, start and end points,
// and well as two functions that test if movement to the right and down is allowed from a given point.
func NewProblem(size image.Rectangle, start, end image.Point, rightTest func(image.Point) bool, downTest func(image.Point) bool) *Problem {
	var legalPositions [stateLen]image.Rectangle

	legalPositions[stateRight] = image.Rect(size.Min.X, size.Min.Y, size.Max.X-1, size.Max.Y)
	legalPositions[stateDown] = image.Rect(size.Min.X, size.Min.Y, size.Max.X, size.Max.Y-1)
	legalPositions[stateLeft] = image.Rect(size.Min.X+1, size.Min.Y, size.Max.X, size.Max.Y)
	legalPositions[stateUp] = image.Rect(size.Min.X, size.Min.Y+1, size.Max.X, size.Max.Y)

	var testFuncs [stateLen]func(image.Point) bool

	testFuncs[stateRight] = rightTest
	testFuncs[stateDown] = downTest

	// note that the dirTestOffset lookup table has non-zero values for left and up,
	// to compensate for the fact that we need to test the cell to the left or above.
	testFuncs[stateLeft] = rightTest
	testFuncs[stateUp] = downTest

	return &Problem{
		size:           size,
		legalPositions: legalPositions,
		testFuncs:      testFuncs,
		start:          start,
		end:            end,
	}
}