package types

// Heap is a max heap
type Heap struct {
	buf []int
}

func (h *Heap) Peak() int {
	return h.buf[0]
}

func (h *Heap) up(idx int) {
	if idx == 0 {
		return
	}
	parent := (idx - 1) / 2
	if h.buf[parent] < h.buf[idx] {
		h.buf[parent], h.buf[idx] = h.buf[idx], h.buf[parent]
		h.up(parent)
		return
	}
}

func (h *Heap) Insert(key int) {
	h.buf = append(h.buf, key)
	h.up(len(h.buf) - 1)
}

func (h *Heap) down(idx int) {
	leftChild := idx*2 + 1
	rightChild := idx*2 + 2
	if leftChild >= len(h.buf) {
		return
	}

	swap := -1
	if h.buf[idx] < h.buf[leftChild] {
		if rightChild >= len(h.buf) || h.buf[leftChild] >= h.buf[rightChild] {
			swap = leftChild
		} else {
			swap = rightChild
		}
	} else if rightChild < len(h.buf) && h.buf[idx] < h.buf[rightChild] {
		swap = rightChild
	}

	if swap < 0 {
		return
	}

	h.buf[swap], h.buf[idx] = h.buf[idx], h.buf[swap]
	h.down(swap)
}

func (h *Heap) Delete() int {
	max := h.buf[0]
	h.buf[0], h.buf[len(h.buf)-1] = h.buf[len(h.buf)-1], h.buf[0]
	h.buf = h.buf[:len(h.buf)-1]
	h.down(0)
	return max
}

func (h *Heap) Build() {
	for i := len(h.buf) - 1; i >= 0; i -= 2 {
		h.down((i - 1) / 2)
	}
}

// put the following to heap_test.go
package types

import (
	"math/rand"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestHeap(t *testing.T) {
	const total = 101
	h := &Heap{}
	perm := rand.Perm(total)
	for _, v := range perm {
		h.Insert(v)
	}
	assert.Equal(t, total-1, h.Peak())
	for i := total - 1; i >= 0; i-- {
		assert.Equal(t, i, h.Delete())
	}
}

func TestBuildHeap(t *testing.T) {
	const total = 101
	perm := rand.Perm(total)
	h := &Heap{buf: perm}
	h.Build()
	assert.Equal(t, total-1, h.Peak())
	for i := total - 1; i >= 0; i-- {
		assert.Equal(t, i, h.Delete())
	}
}

func TestHeapSort(t *testing.T) {
	const total = 101
	perm := rand.Perm(total)
	h := &Heap{buf: perm}
	h.Build()
	assert.Equal(t, total-1, h.Peak())
	for i := total - 1; i >= 0; i-- {
		assert.Equal(t, i, h.Delete())
	}
	for i := 0; i < total; i++ {
		assert.Equal(t, i, perm[i])
	}
}