sequence-js/src/random.ts

import { BitArray } from "./bitarray.js";
import { asArray } from "./utils.js";

export type ElementPredicate<T = any> = (index: number, obj: T) => boolean;
export type ElementWeight<T = any> = (index: number, obj: T) => number;
export type RandomGenerator = () => number;

export interface RandomOptions<T = any> {
  predicate?: ElementPredicate<T>;
  weight?: ElementWeight<T>;
  random?: RandomGenerator;
};

export const alwaysTrue: ElementPredicate = () => true;
export const weightOfOne: ElementWeight = () => 1.0;
export const mathRandom: RandomGenerator = () => Math.random();

const defaultOptions = Object.freeze<Required<RandomOptions>>({
  predicate: alwaysTrue,
  weight: weightOfOne,
  random: mathRandom
});

function mergeOptions<T>(first: RandomOptions<T> | undefined, second: RandomOptions<T> | undefined): RandomOptions<T> | undefined {
  if (!first) {
    return second;
  }

  if (!second) {
    return first;
  }

  const firstPredicate = first.predicate;
  const secondPredicate = second.predicate;

  return {
    predicate: firstPredicate ? secondPredicate ? (i, o) => firstPredicate(i, o) && secondPredicate(i, o) : firstPredicate : secondPredicate,
    weight: first.weight ?? second.weight,
    random: first.random ?? second.random
  };
}

function withDefaultOptions<T>(options: RandomOptions<T> | undefined): Required<RandomOptions<T>> {
  if (!options) {
    return defaultOptions;
  }

  return {
    predicate: options.predicate ?? defaultOptions.predicate,
    weight: options.weight ?? defaultOptions.weight,
    random: options.random ?? defaultOptions.random
  };
}

export function getRandomElement<T>(sequence: Iterable<T>, options?: RandomOptions<T>) {
  const { predicate, weight, random } = withDefaultOptions(options);

  let result: T | undefined = undefined;
  let resultIndex = -1;
  let index = 0;
  let weightAcc = 0.0;

  for (const element of sequence) {
    const currentIndex = index++;

    if (predicate(currentIndex, element)) {
      const w = weight(currentIndex, element);

      if (w <= 0.0) {
        continue;
      }

      weightAcc += w;

      if (random() * weightAcc < w) {
        result = element;
        resultIndex = currentIndex;
      }
    }
  }

  return { element: result, index: resultIndex };
}

export class RandomPicker<T> {
  readonly #elements: Iterable<T>;
  readonly #flags: BitArray;
  readonly #options: RandomOptions<T>;

  public constructor(elements: T[], options?: RandomOptions<T>) {
    this.#elements = elements;
    this.#flags = new BitArray(elements.length);
    this.#options = withDefaultOptions(mergeOptions({ predicate: i => this.#flags.get(i) }, options));

    this.reset();
  }

  public static from<T>(iterable: Iterable<T>, options?: RandomOptions<T>) {
    return new this<T>(asArray(iterable), options);
  }

  public static of<T>(options?: RandomOptions<T>, ...values: T[]) {
    return new this<T>(values, options);
  }

  public get state() {
    return this.#flags;
  }

  public reset() {
    this.#flags.fill(true);
  }

  public next() {
    const result = getRandomElement(this.#elements, this.#options);

    if (result.index < 0) {
      this.reset();
    } else {
      this.#flags.set(result.index, false);

      if (this.#flags.isEmpty()) {
        this.reset();
      }
    }

    return result.element;
  }
}