package dijkstra

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

type Edge struct {
	To     int
	Weight int
}

// ShortestPaths returns distances and predecessors from start.
// Distances for unreachable vertices are math.MaxInt.
func ShortestPaths(graph [][]Edge, start int) ([]int, []int) {
	dist := make([]int, len(graph))
	prev := make([]int, len(graph))
	for i := range graph {
		dist[i] = math.MaxInt
		prev[i] = -1
	}
	if start < 0 || start >= len(graph) {
		return dist, prev
	}

	dist[start] = 0
	pq := priorityQueue{{node: start, dist: 0}}
	heap.Init(&pq)

	for pq.Len() > 0 {
		item := heap.Pop(&pq).(queueItem)
		if item.dist != dist[item.node] {
			continue
		}

		for _, edge := range graph[item.node] {
			if edge.Weight < 0 {
				panic("dijkstra: negative edge weight")
			}
			next := item.dist + edge.Weight
			if next < dist[edge.To] {
				dist[edge.To] = next
				prev[edge.To] = item.node
				heap.Push(&pq, queueItem{node: edge.To, dist: next})
			}
		}
	}

	return dist, prev
}

// Path reconstructs a shortest path to target from a predecessor list.
func Path(prev []int, target int) []int {
	if target < 0 || target >= len(prev) {
		return nil
	}

	path := []int{}
	for at := target; at != -1; at = prev[at] {
		path = append(path, at)
	}

	for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
		path[i], path[j] = path[j], path[i]
	}
	return path
}

type queueItem struct {
	node int
	dist int
}

type priorityQueue []queueItem

func (pq priorityQueue) Len() int { return len(pq) }

func (pq priorityQueue) Less(i, j int) bool {
	return pq[i].dist < pq[j].dist
}

func (pq priorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
}

func (pq *priorityQueue) Push(x any) {
	*pq = append(*pq, x.(queueItem))
}

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

func TestShortestPaths(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 := ShortestPaths(graph, 0)

	if dist[3] != 4 {
		t.Fatalf("distance to 3 = %d, want 4", dist[3])
	}

	got := Path(prev, 3)
	want := []int{0, 2, 1, 3}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("path = %v, want %v", got, want)
		}
	}
}

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

	dist, prev := ShortestPaths(graph, 0)

	if dist[2] != math.MaxInt {
		t.Fatalf("distance to unreachable vertex = %d, want math.MaxInt", dist[2])
	}
	if prev[2] != -1 {
		t.Fatalf("predecessor for unreachable vertex = %d, want -1", prev[2])
	}
}