package pick

import "math/rand"

// Stream retains a random sample of the elements it's been given.
type Stream struct {
	seen  int
	max   int
	set   []interface{}
	randn func(n int) int
}

// New returns a streaming sampler that will store at most max elements.
//
// If no randn function is provided, the sampler will use math/rand.
func New(max int, randn func(n int) int) *Stream {
	if randn == nil {
		randn = rand.Intn
	}
	return &Stream{
		max:   max,
		randn: randn,
	}
}

// Add includes the provided element in the Picker's set, subject to sampling.
func (s *Stream) Add(element interface{}) {
	s.AddFn(func() interface{} { return element })
}

// AddFn determines whether a sample will be needed, calling the provided
// function to get the element when it is required. If the function returns
// nil, its result is ignored.
func (s *Stream) AddFn(fn func() interface{}) {
	s.seen++

	if len(s.set) < s.max {
		// If there's still space available, save every element
		if element := fn(); element != nil {
			s.set = append(s.set, element)
		}
		return
	}
	if n := s.randn(s.seen); n < s.max {
		// When the set is "full", pick an element to replace from the
		// imaginary series of all elements we've seen. If the selected
		// element is past the end of the slice of saved elements, we'll
		// discard the provided element.
		//
		// If the selected element is within the slice, write the new element
		// to that location.
		//
		// This results in the set remaining a uniformly random sample of the
		// inbound element stream (to the extent that the provided RNG is
		// fair).
		if element := fn(); element != nil {
			s.set[n] = element
		}
	}
}

// Picked returns the sampled values.
func (s *Stream) Picked() []interface{} {
	return s.set
}
