persona-community-2/pkg/persona/pools.go

package persona

import (
	"fmt"
	"hash/fnv"
	"math/rand"
)

// Pool represents a collection of items that can be selected from.
type Pool[T any] struct {
	Items []T
}

// NewPool creates a new pool with the given items.
func NewPool[T any](items []T) *Pool[T] {
	return &Pool[T]{Items: items}
}

// Random selects a random item from the pool.
// Returns the zero value if the pool is empty.
func (p *Pool[T]) Random() T {
	var zero T
	if len(p.Items) == 0 {
		return zero
	}
	return p.Items[rand.Intn(len(p.Items))]
}

// RandomWithSeed selects an item deterministically based on a seed.
// The same seed will always return the same item.
func (p *Pool[T]) RandomWithSeed(seed int64) T {
	var zero T
	if len(p.Items) == 0 {
		return zero
	}
	r := rand.New(rand.NewSource(seed))
	return p.Items[r.Intn(len(p.Items))]
}

// RandomN selects n random items from the pool.
// May return duplicates. Returns up to n items or all items if less available.
func (p *Pool[T]) RandomN(n int) []T {
	if len(p.Items) == 0 || n <= 0 {
		return nil
	}
	if n > len(p.Items) {
		n = len(p.Items)
	}
	result := make([]T, n)
	for i := 0; i < n; i++ {
		result[i] = p.Items[rand.Intn(len(p.Items))]
	}
	return result
}

// SelectDescription selects a description from a pool using a deterministic seed.
// This ensures the same character always gets the same descriptions.
func SelectDescription(pool []string, characterID string, fieldName string) string {
	if len(pool) == 0 {
		return ""
	}

	// Create a deterministic seed from character ID and field name
	h := fnv.New64a()
	h.Write([]byte(characterID))
	h.Write([]byte(fieldName))
	seed := int64(h.Sum64())

	r := rand.New(rand.NewSource(seed))
	return pool[r.Intn(len(pool))]
}

// SelectDescriptionN selects n unique descriptions from a pool.
func SelectDescriptionN(pool []string, characterID string, fieldName string, n int) []string {
	if len(pool) == 0 || n <= 0 {
		return nil
	}

	// Create a deterministic seed
	h := fnv.New64a()
	h.Write([]byte(characterID))
	h.Write([]byte(fieldName))
	seed := int64(h.Sum64())

	r := rand.New(rand.NewSource(seed))

	// If n >= pool size, shuffle and return all
	if n >= len(pool) {
		result := make([]string, len(pool))
		copy(result, pool)
		r.Shuffle(len(result), func(i, j int) {
			result[i], result[j] = result[j], result[i]
		})
		return result
	}

	// Select n unique items
	indices := r.Perm(len(pool))[:n]
	result := make([]string, n)
	for i, idx := range indices {
		result[i] = pool[idx]
	}
	return result
}

// WeightedPool allows weighted random selection.
type WeightedPool[T any] struct {
	Items   []T
	Weights []float64
	total   float64
}

// NewWeightedPool creates a weighted pool.
// Weights should be positive numbers; they don't need to sum to 1.
// Returns an error if items and weights have different lengths or if any weight is negative.
func NewWeightedPool[T any](items []T, weights []float64) (*WeightedPool[T], error) {
	if len(items) != len(weights) {
		return nil, fmt.Errorf("%w: items length (%d) != weights length (%d)",
			ErrInvalidPoolConfig, len(items), len(weights))
	}

	var total float64
	for i, w := range weights {
		if w < 0 {
			return nil, fmt.Errorf("%w: negative weight at index %d: %f",
				ErrInvalidPoolConfig, i, w)
		}
		total += w
	}

	return &WeightedPool[T]{
		Items:   items,
		Weights: weights,
		total:   total,
	}, nil
}

// MustNewWeightedPool creates a weighted pool or panics on error.
// Use this only at initialization time. For runtime use, prefer NewWeightedPool.
func MustNewWeightedPool[T any](items []T, weights []float64) *WeightedPool[T] {
	pool, err := NewWeightedPool(items, weights)
	if err != nil {
		panic(err)
	}
	return pool
}

// Random selects a weighted random item.
func (p *WeightedPool[T]) Random() T {
	var zero T
	if len(p.Items) == 0 {
		return zero
	}

	r := rand.Float64() * p.total
	var cumulative float64

	for i, w := range p.Weights {
		cumulative += w
		if r <= cumulative {
			return p.Items[i]
		}
	}

	return p.Items[len(p.Items)-1]
}

// RandomWithSeed selects a weighted random item deterministically.
func (p *WeightedPool[T]) RandomWithSeed(seed int64) T {
	var zero T
	if len(p.Items) == 0 {
		return zero
	}

	rng := rand.New(rand.NewSource(seed))
	r := rng.Float64() * p.total
	var cumulative float64

	for i, w := range p.Weights {
		cumulative += w
		if r <= cumulative {
			return p.Items[i]
		}
	}

	return p.Items[len(p.Items)-1]
}