package heap

import (
	"cmp"
	"container/heap"
	"slices"
)

type Heap[T comparable] struct {
	Data     []T
	Index    map[T]map[int]bool
	lessFunc func(a, b T) bool
}

func More[T cmp.Ordered](x, y T) bool {
	return !cmp.Less(x, y) && x != y
}

func NewBy[T comparable](lessFunc func(a, b T) bool, items ...T) *Heap[T] {
	h := &Heap[T]{slices.Clone(items), map[T]map[int]bool{}, lessFunc}
	heap.Init(h)

	for i, v := range h.Data {
		if h.Index[v] == nil {
			h.Index[v] = map[int]bool{}
		}

		h.Index[v][i] = true
	}

	return h
}

func New[T cmp.Ordered](items ...T) *Heap[T] {
	return NewBy(func(a, b T) bool { return a < b }, items...)
}

func (h *Heap[T]) putIndex(item T, index int) {
	if h.Index == nil {
		h.Index = map[T]map[int]bool{}
	}

	if h.Index[item] == nil {
		h.Index[item] = map[int]bool{}
	}

	h.Index[item][index] = true
}

func (h *Heap[T]) dropIndex(item T, index int) {
	if h.Index == nil {
		h.Index = map[T]map[int]bool{}
	}

	if h.Index[item] == nil {
		h.Index[item] = map[int]bool{}
	}

	delete(h.Index[item], index)
}

func (h Heap[T]) Len() int           { return len(h.Data) }
func (h Heap[T]) Less(i, j int) bool { return h.lessFunc(h.Data[i], h.Data[j]) }
func (h Heap[T]) Swap(i, j int) {
	h.dropIndex(h.Data[i], i)
	h.dropIndex(h.Data[j], j)

	h.Data[i], h.Data[j] = h.Data[j], h.Data[i]

	h.putIndex(h.Data[i], i)
	h.putIndex(h.Data[j], j)
}

func (h *Heap[T]) Push(x any) {
	datum := x.(T)
	h.Data = append(h.Data, datum)
	h.putIndex(datum, len(h.Data)-1)
}

func (h *Heap[T]) Pop() any {
	index := len(h.Data) - 1
	datum := h.Data[index]
	h.dropIndex(datum, index)

	h.Data = h.Data[:index]
	return datum
}

func (h *Heap[T]) GetIndex(item T) int {
	if h.Index[item] == nil {
		return -1
	}

	var key int
	for k := range h.Index[item] {
		key = k
		break
	}

	return key
}

func (h *Heap[T]) Drop(item T) {
	heap.Remove(h, h.GetIndex(item))
}

func (h *Heap[T]) Put(item T) {
	heap.Push(h, item)
}

func (h *Heap[T]) Top() T {
	return h.Data[0]
}