package unbounded

import (
	"container/heap"

	"smariot.com/tsp/internal/solver/problem"
)

type heapEntry[State comparable] struct {
	state    State
	minIndex int
	maxIndex int
}

type minHeap[P problem.Problem[State], State comparable] struct {
	problem problem.Problem[State]
	entries []heapEntry[State]
	indexes []int
}

func (h minHeap[P, State]) Len() int {
	return len(h.indexes)
}

func (h minHeap[P, State]) Less(i, j int) bool {
	return h.problem.OptimisticLess(h.entries[h.indexes[i]].state, h.entries[h.indexes[j]].state)
}

func (h minHeap[P, State]) Swap(i, j int) {
	h.entries[h.indexes[i]].minIndex = j
	h.entries[h.indexes[j]].minIndex = i
	h.indexes[i], h.indexes[j] = h.indexes[j], h.indexes[i]
}

func (h *minHeap[P, State]) Push(x any) {
	index := x.(int)
	h.entries[index].minIndex = len(h.indexes)
	h.indexes = append(h.indexes, index)
}

func (h *minHeap[P, State]) Pop() any {
	n := len(h.indexes)
	index := h.indexes[n-1]
	h.entries[index].minIndex = -1
	h.indexes = h.indexes[:n-1]
	return index
}

type maxHeap[P problem.Problem[State], State comparable] struct {
	problem problem.Problem[State]
	entries []heapEntry[State]
	indexes []int
}

func (h maxHeap[P, State]) Len() int {
	return len(h.indexes)
}

func (h maxHeap[P, State]) Less(i, j int) bool {
	return h.problem.PessimisticLess(h.entries[h.indexes[j]].state, h.entries[h.indexes[i]].state)
}

func (h maxHeap[P, State]) Swap(i, j int) {
	h.entries[h.indexes[i]].maxIndex = j
	h.entries[h.indexes[j]].maxIndex = i
	h.indexes[i], h.indexes[j] = h.indexes[j], h.indexes[i]
}

func (h *maxHeap[P, State]) Push(x any) {
	index := x.(int)
	h.entries[index].maxIndex = len(h.indexes)
	h.indexes = append(h.indexes, index)
}

func (h *maxHeap[P, State]) Pop() any {
	n := len(h.indexes)
	index := h.indexes[n-1]
	h.entries[index].maxIndex = -1
	h.indexes = h.indexes[:n-1]
	return index
}

type solver[P problem.Problem[State], State comparable] struct {
	minHeap[P, State]
	maxHeap maxHeap[P, State]
	free    []int
}

func (s *solver[P, State]) Push(state State) {
	if len(s.free) == 0 {
		// if this is worse than the worst state, discard it.
		if !s.problem.PessimisticLess(state, s.entries[s.maxHeap.indexes[0]].state) {
			s.problem.Discard(state)
			return
		}

		// otherwise, discard and replace the worst state.
		index := s.maxHeap.indexes[0]
		s.problem.Discard(s.entries[index].state)
		s.entries[index].state = state
		heap.Fix(&s.minHeap, s.entries[index].minIndex)
		heap.Fix(&s.maxHeap, 0)
		return
	}

	index := s.free[len(s.free)-1]
	s.free = s.free[:len(s.free)-1]

	s.entries[index].state = state

	heap.Push(&s.minHeap, index)
	heap.Push(&s.maxHeap, index)
}

func (s *solver[P, State]) Pop() (State, bool) {
	if s.Len() == 0 {
		var zero State
		return zero, false
	}

	index := heap.Pop(&s.minHeap).(int)
	s.free = append(s.free, index)

	heap.Remove(&s.maxHeap, s.entries[index].maxIndex)

	return s.entries[index].state, true
}

func (s *solver[P, State]) Reset() {
	for _, index := range s.minHeap.indexes {
		s.problem.Discard(s.entries[index].state)
		s.free = append(s.free, index)
	}

	s.minHeap.indexes = s.minHeap.indexes[:0]
	s.maxHeap.indexes = s.maxHeap.indexes[:0]
}

// Returns a solver for unbounded problems.
//
// It maintains both a min and a max heap, and will automatically discard states once it reaches a maximum capacity.
//
// It doesn't keep track of known states. Submitting a state multiple times will result in multiple copies being stored,
// and problem.Discard being called multiple times.
func New[P problem.Problem[State], State comparable](problem P, capacity int) *solver[P, State] {
	if capacity <= 0 {
		panic("unbounded.New: capacity must be greater than 0")
	}

	free := make([]int, capacity)
	entries := make([]heapEntry[State], capacity)

	for i := 0; i < capacity; i++ {
		free[i] = capacity - i - 1
		entries[i].minIndex = -1
		entries[i].maxIndex = -1
	}

	indexes := make([]int, capacity*2)

	return &solver[P, State]{
		free: free,
		minHeap: minHeap[P, State]{
			problem: problem,
			entries: entries,
			indexes: indexes[0:0:capacity],
		},
		maxHeap: maxHeap[P, State]{
			problem: problem,
			entries: entries,
			indexes: indexes[capacity : capacity : capacity*2],
		},
	}
}