package sampling

import (
	"math"
	"math/rand"
	"sort"
)

// Candidate is a question ID paired with its sampling weight.
type Candidate struct {
	ID     string
	Weight float64
}

// SelectWeighted picks n distinct candidates using the A-Res weighted
// reservoir algorithm (Efraimidis–Spirakis). Each item's selection
// probability is proportional to its weight. O(m log m) time.
func SelectWeighted(candidates []Candidate, n int, rng *rand.Rand) []Candidate {
	if n >= len(candidates) {
		out := make([]Candidate, len(candidates))
		copy(out, candidates)
		return out
	}

	type keyed struct {
		c   Candidate
		key float64
	}

	keys := make([]keyed, len(candidates))
	for i, c := range candidates {
		u := rng.Float64()
		if u == 0 {
			u = 1e-12 // avoid log(0) / pow weirdness
		}
		keys[i] = keyed{c, math.Pow(u, 1.0/c.Weight)}
	}

	sort.Slice(keys, func(i, j int) bool { return keys[i].key > keys[j].key })

	out := make([]Candidate, n)
	for i := range out {
		out[i] = keys[i].c
	}
	return out
}