package dijkstra

import (
	"container/heap"
	"reflect"
	"testing"
)

// Edge is a directed, weighted edge to another vertex.
type Edge struct {
	To     int
	Weight int
}

// Dijkstra returns shortest distances and predecessors from source.
// Unreachable vertices have distance MaxInt and predecessor -1.
func Dijkstra(graph [][]Edge, source int) (dist, prev []int) {
	const infinity = int(^uint(0) >> 1)

	dist = make([]int, len(graph))
	prev = make([]int, len(graph))
	for i := range graph {
		dist[i] = infinity
		prev[i] = -1
	}

	if source < 0 || source >= len(graph) {
		return dist, prev
	}

	dist[source] = 0
	queue := priorityQueue{{vertex: source, distance: 0}}
	heap.Init(&queue)

	for queue.Len() > 0 {
		current := heap.Pop(&queue).(item)
		if current.distance != dist[current.vertex] {
			continue // Ignore stale queue entries.
		}

		for _, edge := range graph[current.vertex] {
			if edge.To < 0 || edge.To >= len(graph) || edge.Weight < 0 {
				continue
			}
			if dist[current.vertex] > infinity-edge.Weight {
				continue
			}

			candidate := dist[current.vertex] + edge.Weight
			if candidate < dist[edge.To] {
				dist[edge.To] = candidate
				prev[edge.To] = current.vertex
				heap.Push(&queue, item{vertex: edge.To, distance: candidate})
			}
		}
	}

	return dist, prev
}

// Path reconstructs a source-to-target path from Dijkstra's predecessors.
func Path(prev []int, source, target int) []int {
	if source < 0 || source >= len(prev) || target < 0 || target >= len(prev) {
		return nil
	}

	path := make([]int, 0)
	for vertex, steps := target, 0; vertex != -1 && steps <= len(prev); steps++ {
		path = append(path, vertex)
		if vertex == source {
			for left, right := 0, len(path)-1; left < right; left, right = left+1, right-1 {
				path[left], path[right] = path[right], path[left]
			}
			return path
		}
		vertex = prev[vertex]
	}
	return nil
}

type item struct {
	vertex   int
	distance int
}

type priorityQueue []item

func (p priorityQueue) Len() int           { return len(p) }
func (p priorityQueue) Less(i, j int) bool { return p[i].distance < p[j].distance }
func (p priorityQueue) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

func (p *priorityQueue) Push(value any) {
	*p = append(*p, value.(item))
}

func (p *priorityQueue) Pop() any {
	old := *p
	last := old[len(old)-1]
	*p = old[:len(old)-1]
	return last
}

func TestDijkstra(t *testing.T) {
	graph := [][]Edge{
		{{To: 1, Weight: 4}, {To: 2, Weight: 1}},
		{{To: 3, Weight: 1}},
		{{To: 1, Weight: 2}, {To: 3, Weight: 5}},
		{},
	}

	dist, prev := Dijkstra(graph, 0)

	if want := []int{0, 3, 1, 4}; !reflect.DeepEqual(dist, want) {
		t.Fatalf("distances = %v, want %v", dist, want)
	}
	if want := []int{0, 2, 1, 3}; !reflect.DeepEqual(Path(prev, 0, 3), want) {
		t.Fatalf("path = %v, want %v", Path(prev, 0, 3), want)
	}
}

func TestDijkstraUnreachable(t *testing.T) {
	graph := [][]Edge{
		{{To: 1, Weight: 2}},
		{},
		{},
	}

	_, prev := Dijkstra(graph, 0)
	if path := Path(prev, 0, 2); path != nil {
		t.Fatalf("path to unreachable vertex = %v, want nil", path)
	}
}