package heap

import "golang.org/x/exp/constraints"

type HeapType bool

const (
	MINHEAP HeapType = false
	MAXHEAP HeapType = true
)

// Heap object for generic orderable (i.e. can be compared with < or >) data types 
type Heap[Orderable constraints.Ordered] struct {
	data     []Orderable
	heapSize int
}

// Parent index is returned for an input index
func Parent(index int) (parent int) {
	return (index - 1) / 2
}

// Left index is returned for an input index
func Left(index int) (left int) {
	return 2*index + 1
}

// Right index is returned for an input index
func Right(index int) (right int) {
	return 2*index + 2
}

// exchange two items within the context of the heap. requires max/min heapify afterwards
func exchange[Orderable constraints.Ordered](h *Heap[Orderable], x, y int) {
	h.data[x], h.data[y] = h.data[y], h.data[x]
}

// IsMaxHeap returns true when the heap satisfies the max heap property i.e. holds a valid structure for a max heap
func IsMaxHeap[Orderable constraints.Ordered](h *Heap[Orderable]) bool {
	for index, val := range h.data {
		if h.data[Parent(index)] < val {
			return false
		}
	}
	return true
}

// IsMinHeap returns true when the heap satisfies the min heap property i.e. holds a valid structure for a min heap
func IsMinHeap[Orderable constraints.Ordered](h *Heap[Orderable]) bool {
	for index, val := range h.data {
		if h.data[Parent(index)] > val {
			return false
		}
	}
	return true
}


// maxHeapify reorders into a min heap starting at the input index i
func maxHeapify[Orderable constraints.Ordered](h *Heap[Orderable], i int) {
	l, r := Left(i), Right(i)

	var largest int
	if l < h.heapSize && h.data[l] > h.data[i] {
		largest = l
	} else {
		largest = i
	}
	if r < h.heapSize && h.data[r] > h.data[largest] {
		largest = r
	}

	if largest != i {
		exchange(h, i, largest)
		maxHeapify(h, largest)
	}
}

// minHeapify reorders into a min heap starting at the input index i
func minHeapify[Orderable constraints.Ordered](h *Heap[Orderable], i int) {
	l, r := Left(i), Right(i)

	var smallest int
	if l < h.heapSize && h.data[l] < h.data[i] {
		smallest = l
	} else {
		smallest = i
	}
	if r < h.heapSize && h.data[r] < h.data[smallest] {
		smallest = r
	}

	if smallest != i {
		exchange(h, i, smallest)
		minHeapify(h, smallest)
	}
}

// BuildHeap takes a generic slice (any type that can be compared with < and >) and the type of heap and returns a pointer to a heap structure
func BuildHeap[Orderable constraints.Ordered](d []Orderable, maxHeap HeapType) *Heap[Orderable] {
	heap := &Heap[Orderable]{d, len(d)}
	for i := len(d) / 2; i >= 0; i-- {
		if maxHeap {
			maxHeapify(heap, i)
		} else {
			minHeapify(heap, i)
		}
	}
	return heap

}

// HeapSort takes a generic slice (any type that can be compared with < and >) and returns the sorted slice
func HeapSort[Orderable constraints.Ordered](d []Orderable, maxHeap HeapType) []Orderable {
    heap := BuildHeap(d, maxHeap)

	for i := len(d) - 1; i > 0; i-- {
		exchange(heap, 0, i)
		heap.heapSize = heap.heapSize - 1
		if maxHeap {
			maxHeapify(heap, 0)
		} else {
			minHeapify(heap, 0)
		}
	}
	return heap.data
}