package dijkstra

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

type Edge struct {
	To     int
	Weight int
}

// Graph maps each node to its outgoing weighted edges.
type Graph map[int][]Edge

// Dijkstra returns shortest distances from start to every reachable node.
// Edge weights must be non-negative.
func Dijkstra(g Graph, start int) map[int]int {
	dist := map[int]int{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 _, e := range g[cur.node] {
			next := cur.dist + e.Weight
			if old, ok := dist[e.To]; !ok || next < old {
				dist[e.To] = next
				heap.Push(pq, item{node: e.To, dist: next})
			}
		}
	}

	return dist
}

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
}

func Distance(dist map[int]int, node int) int {
	if d, ok := dist[node]; ok {
		return d
	}
	return math.MaxInt
}

func TestDijkstraFindsShortestPaths(t *testing.T) {
	g := Graph{
		0: {{To: 1, Weight: 4}, {To: 2, Weight: 1}},
		1: {{To: 3, Weight: 1}},
		2: {{To: 1, Weight: 2}, {To: 3, Weight: 5}},
		3: nil,
	}

	dist := Dijkstra(g, 0)

	want := map[int]int{0: 0, 1: 3, 2: 1, 3: 4}
	for node, expected := range want {
		if got := dist[node]; got != expected {
			t.Fatalf("dist[%d] = %d, want %d", node, got, expected)
		}
	}
}

func TestDijkstraLeavesUnreachableNodesAbsent(t *testing.T) {
	g := Graph{
		0: {{To: 1, Weight: 7}},
		1: nil,
		2: {{To: 3, Weight: 1}},
	}

	dist := Dijkstra(g, 0)

	if _, ok := dist[2]; ok {
		t.Fatal("unreachable node 2 should not be present")
	}
	if got := Distance(dist, 2); got != math.MaxInt {
		t.Fatalf("Distance for unreachable node = %d, want %d", got, math.MaxInt)
	}
}