package queue

import (
	"codeberg.org/snonux/algorithms/ds"
)

type HeapPriority[T ds.Number] struct {
	ElementaryPriority[T]
}

func NewHeapPriority[T ds.Number](capacity int) *HeapPriority[T] {
	q := HeapPriority[T]{
		ElementaryPriority: ElementaryPriority[T]{make(ds.ArrayList[T], 0, capacity), capacity},
	}

	// Index 0 unused
	q.a = append(q.a, 0)
	return &q
}

func (q *HeapPriority[T]) Empty() bool {
	return q.Size() == 0
}

func (q *HeapPriority[T]) Size() int {
	return len(q.a) - 1
}

func (q *HeapPriority[T]) Insert(a T) {
	q.a = append(q.a, a)
	q.swim(len(q.a) - 1)
}

func (q *HeapPriority[T]) Max() (max T) {
	if q.Empty() {
		return 0
	}

	return q.a[1]
}

func (q *HeapPriority[T]) DeleteMax() T {
	switch q.Size() {
	case 0:
		return 0
	case 1:
		max := q.a[1]
		q.a = q.a[0:1]
		return max
	default:
		a := q.a[1]
		// Last index
		max := len(q.a) - 1

		q.a.Swap(1, max)
		q.a = q.a[0:max]
		q.sink(1)

		return a
	}
}

func (q *HeapPriority[T]) swim(k int) {
	for k > 1 && q.a[k/2] < q.a[k] {
		q.a.Swap(k/2, k)
		k = k / 2
	}
}

func (q *HeapPriority[T]) sink(k int) {
	// Last index
	max := len(q.a) - 1

	for 2*k <= max {
		// Left child
		j := 2 * k

		// Go to right child?
		if j < max && q.a[j] < q.a[j+1] {
			j++
		}

		// Found my spot in the heap
		if q.a[k] >= q.a[j] {
			break
		}

		q.a.Swap(k, j)
		k = j
	}
}