path: root/base/src/main/java/bjc/utils/gen/WeightedRandom.java

package bjc.utils.gen;

import java.util.Random;

import bjc.data.IPair;
import bjc.data.Pair;
import bjc.funcdata.FunctionalList;
import bjc.funcdata.IList;

/**
 * Represents a random number generator where certain results are weighted more
 * heavily than others.
 *
 * @author ben
 *
 * @param <E>
 *                The type of values that are randomly selected.
 */
public class WeightedRandom<E> {
	private final IList<IPair<Integer, E>> values;

	/* The source for any needed random numbers */
	private Random source;
	/* The total chance for all values. */
	private int totalChance;

	private final static Random BASE = new Random();

	private boolean exhaust;

	/**
	 * Create a new weighted random generator with the specified source of
	 * randomness.
	 *
	 * @param src
	 *                The source of randomness to use.
	 */
	public WeightedRandom(Random src) {
		values = new FunctionalList<>();

		if (src == null) throw new NullPointerException("Source of randomness must not be null");

		source = src;
	}

	/**
	 * Create a new weighted random generator.
	 */
	public WeightedRandom() {
		this(BASE);
	}

	private WeightedRandom(Random src, IList<IPair<Integer, E>> vals, int chance) {
		source = src;

		values = vals;

		totalChance = chance;
	}

	/**
	 * Add a probability for a specific result to be given.
	 *
	 * @param chance
	 *                The chance to get this result.
	 *
	 * @param result
	 *                The result to get when the chance comes up.
	 */
	public void addProbability(final int chance, final E result) {
		values.add(new Pair<>(chance, result));

		totalChance += chance;
	}

	/**
	 * Generate a weighted random value.
	 *
	 * @return A random value selected in a weighted fashion.
	 */
	public E generateValue() {
		return generateValue(source);
	}

	/**
	 * Generate a random value, using the specified Random.
	 * 
	 * @param rn
	 *                The Random instance to use.
	 * @return A random value.
	 */
	public E generateValue(Random rn) {
		int target = rn.nextInt(totalChance);
		for (IPair<Integer, E> val : values) {
			int prob = val.getLeft();

			if (target < prob) {
				if (exhaust) {
					totalChance -= val.getLeft();
					values.removeMatching(val);
				}

				return val.getRight();
			}

			target -= prob;
		}

		return null;
	}

	/**
	 * Return a list of values that can be generated by this generator
	 *
	 * @return A list of all the values that can be generated
	 */
	public IList<E> getResults() {
		return values.map(IPair::getRight);
	}

	/**
	 * Return a list containing values that can be generated paired with the
	 * probability of those values being generated
	 *
	 * @return A list of pairs of values and value probabilities
	 */
	public IList<IPair<Integer, E>> getValues() {
		return values;
	}

	/**
	 * Get a descending value.
	 * 
	 * Descending values are quite simple. You have a 1 in factor chance to
	 * advance to the next value, otherwise, the current value is the one
	 * returned.
	 * 
	 * @param factor
	 *                The descent factor to use.
	 * @return A random value.
	 */
	public E getDescent(int factor) {
		return getDescent(factor, source);
	}

	/**
	 * Get a descending value.
	 * 
	 * Descending values are quite simple. You have a 1 in factor chance to
	 * advance to the next value, otherwise, the current value is the one
	 * returned.
	 * 
	 * @param factor
	 *                The descent factor to use.
	 * @param rn
	 *                The Random instance to use.
	 * @return A random value.
	 */
	public E getDescent(int factor, Random rn) {
		if (values.getSize() == 0) return null;

		for (IPair<Integer, E> val : values) {
			if (rn.nextInt(factor) == 0) continue;

			if (exhaust) {
				totalChance -= val.getLeft();

				values.removeMatching(val);
			}

			return val.getRight();
		}

		IPair<Integer, E> val = values.getByIndex(values.getSize() - 1);
		if (exhaust) values.removeMatching(val);

		return val.getRight();
	}

	/**
	 * Get a value, treating this as a set of binomial trials.
	 * 
	 * Essentially, the system rolls a bound-sided dice trials times, and
	 * marks a success for every roll less than or equal to target.
	 * 
	 * @param target
	 *                The number to roll at or under.
	 * @param bound
	 *                The maximum roll value.
	 * @param trials
	 *                The number of times to roll.
	 * @return The value at the index corresponding to the number of
	 *         successes.
	 */
	public E getBinomial(int target, int bound, int trials) {
		return getBinomial(target, bound, trials, source);
	}

	/**
	 * Get a value, treating this as a set of binomial trials.
	 * 
	 * Essentially, the system rolls a bound-sided dice trials times, and
	 * marks a success for every roll less than or equal to target.
	 * 
	 * @param target
	 *                The number to roll at or under.
	 * @param bound
	 *                The maximum roll value.
	 * @param trials
	 *                The number of times to roll.
	 * @param rn
	 *                The Random instance to use.
	 * @return The value at the index corresponding to the number of
	 *         successes.
	 */
	public E getBinomial(int target, int bound, int trials, Random rn) {
		if (values.getSize() == 0) return null;

		int numSuc = 0;

		for (int i = 0; i < trials; i++) {
			/*
			 * Adjust for zero, because it's easy to think of this
			 * as rolling a bound-sided dice and marking a success
			 * for every roll less than or equal to target.
			 */
			int num = rn.nextInt(bound) + 1;

			if (num <= target) {
				//System.err.printf("\t\tTRACE: mark binomial success (%d <= 1d%d, %d)\n", target, bound, num);
				numSuc += 1;
			}
		}

		//System.err.printf("\tTRACE: got %d success for binomial trials (%d <= 1d%d, %d times)\n", numSuc, target, bound, trials);
		IPair<Integer, E> val = values.getByIndex(Math.min(numSuc, values.getSize() - 1));
		if (exhaust) {
			totalChance -= val.getLeft();

			values.removeMatching(val);
		}

		return val.getRight();
	}

	/**
	 * Return a new WeightedRandom that is exhaustible (only returns one value once).
	 * 
	 * @return A new WeightedRandom that is exhaustible.
	 */
	public WeightedRandom<E> exhaustible() {
		IList<IPair<Integer, E>> lst = new FunctionalList<>();
		for (IPair<Integer, E> val : values) {
			lst.add(val);
		}

		WeightedRandom<E> res = new WeightedRandom<>(source, lst, totalChance);

		res.exhaust = true;

		return res;
	}
}