package dijkstra

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

type Edge struct {
	To     int
	Weight int
}

type item struct {
	node int
	dist int
}

type priorityQueue []item

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.(item)) }
func (pq *priorityQueue) Pop() any {
	old := *pq
	n := len(old)
	x := old[n-1]
	*pq = old[:n-1]
	return x
}

// ShortestPaths returns the minimum distance from start to every node.
// Unreachable nodes keep distance math.MaxInt.
func ShortestPaths(graph [][]Edge, start int) []int {
	dist := make([]int, len(graph))
	for i := range dist {
		dist[i] = math.MaxInt
	}
	if start < 0 || start >= len(graph) {
		return dist
	}

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

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

		for _, edge := range graph[cur.node] {
			if edge.Weight < 0 || edge.To < 0 || edge.To >= len(graph) {
				continue
			}

			nextDist := cur.dist + edge.Weight
			if nextDist < dist[edge.To] {
				dist[edge.To] = nextDist
				heap.Push(pq, item{node: edge.To, dist: nextDist})
			}
		}
	}

	return dist
}

// ShortestPath returns the minimum distance from start to target.
func ShortestPath(graph [][]Edge, start, target int) (int, bool) {
	dist := ShortestPaths(graph, start)
	if target < 0 || target >= len(dist) || dist[target] == math.MaxInt {
		return 0, false
	}
	return dist[target], true
}

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}},
		{},
	}

	got := ShortestPaths(graph, 0)
	want := []int{0, 3, 1, 4}

	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("dist[%d] = %d, want %d", i, got[i], want[i])
		}
	}
}

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

	if _, ok := ShortestPath(graph, 0, 2); ok {
		t.Fatal("expected node 2 to be unreachable")
	}
}