// SPDX-License-Identifier: Apache-2.0

package bounded_tracking

import (
	"container/heap"

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

type minHeap[P problem.Problem[State], State comparable] struct {
	problem problem.Problem[State]
	known   map[State]int
	items   []State
}

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

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

func (h minHeap[P, State]) Swap(i, j int) {
	h.items[i], h.items[j] = h.items[j], h.items[i]
	h.known[h.items[i]] = i
	h.known[h.items[j]] = j
}

func (h *minHeap[P, State]) Push(x any) {
	state := x.(State)
	h.items = append(h.items, state)
	h.known[state] = len(h.items)
}

func (h *minHeap[P, State]) Pop() any {
	n := len(h.items)
	state := h.items[n-1]
	h.items = h.items[:n-1]
	delete(h.known, state)
	return state
}

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

func (s *solver[P, State]) Push(state State) {
	if i, ok := s.known[state]; ok {
		// The state is already in the heap, update its position instead.
		heap.Fix(&s.minHeap, i)
		return
	}

	heap.Push(&s.minHeap, state)
}

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

	return heap.Pop(&s.minHeap).(State), true
}

func (s *solver[P, State]) Reset() {
	for _, state := range s.items {
		s.problem.Discard(state)
	}

	s.items = s.items[:0]
	clear(s.known)
}

// Returns a solver that for bounded problems that can update their states.
//
// Submitting a state that is already in the heap will update its position,
// rather than adding it again. problem.Discard will only be invoked once.
func New[P problem.Problem[State], State comparable](problem P) *solver[P, State] {
	return &solver[P, State]{
		minHeap: minHeap[P, State]{
			problem: problem,
			known:   make(map[State]int),
		},
	}
}