package aoc

type PriorityQueue[T comparable] struct {
	items      []T
	comparator func(a, b T) int
}

func NewPriorityQueue[T comparable](comparator func(a, b T) int) PriorityQueue[T] {
	return PriorityQueue[T]{comparator: comparator}
}
func (pq *PriorityQueue[T]) Push(item T) {
	pq.items = append(pq.items, item)
	pq.heapifyUp(len(pq.items) - 1)
}
func (pq *PriorityQueue[T]) Pop() T {
	top := pq.items[0]
	lastIndex := len(pq.items) - 1
	pq.items[0], pq.items[lastIndex] = pq.items[lastIndex], pq.items[0]
	pq.items = pq.items[:lastIndex]
	pq.heapifyDown(0)
	return top
}
func (pq *PriorityQueue[T]) Peek() T {
	return pq.items[0]
}
func (pq *PriorityQueue[T]) Len() int {
	return len(pq.items)
}
func (pq *PriorityQueue[T]) IsEmpty() bool {
	return len(pq.items) == 0
}
func (pq *PriorityQueue[T]) IsNotEmpty() bool {
	return !pq.IsEmpty()
}
func (pq *PriorityQueue[T]) heapifyUp(index int) {
	for index > 0 {
		parentIndex := (index - 1) / 2
		if pq.comparator(pq.items[index], pq.items[parentIndex]) < 0 {
			pq.items[index], pq.items[parentIndex] = pq.items[parentIndex], pq.items[index]
			index = parentIndex
		} else {
			break
		}
	}
}
func (pq *PriorityQueue[T]) heapifyDown(index int) {
	for {
		leftChild, rightChild, smallest := 2*index+1, 2*index+2, index
		if leftChild < len(pq.items) && pq.comparator(pq.items[leftChild], pq.items[smallest]) < 0 {
			smallest = leftChild
		}
		if rightChild < len(pq.items) && pq.comparator(pq.items[rightChild], pq.items[smallest]) < 0 {
			smallest = rightChild
		}
		if smallest != index {
			pq.items[index], pq.items[smallest] = pq.items[smallest], pq.items[index]
			index = smallest
		} else {
			break
		}
	}
}