path: root/BJC-Utils2/src/main/java/bjc/utils/gen/WeightedGrammar.java

package bjc.utils.gen;

import java.util.Hashtable;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.function.Function;

import bjc.utils.data.Pair;
import bjc.utils.funcdata.FunctionalList;

/**
 * A random grammar, where certain rules will come up more often than
 * others.
 * 
 * @author ben
 *
 * @param <E>
 *            The values that make up sentances of this grammar.
 */
public class WeightedGrammar<E> {
	/**
	 * The initial rule of the grammar
	 */
	protected String									initialRule;

	/**
	 * The rules currently in this grammar
	 */
	protected Map<E, WeightedRandom<FunctionalList<E>>>	rules;

	/**
	 * The random number generator used for random numbers
	 */
	private Random										rng;

	/**
	 * All of the subgrammars of this grammar
	 */
	protected Map<E, WeightedGrammar<E>>				subgrammars;

	/**
	 * Create a new weighted grammar.
	 */
	public WeightedGrammar() {
		rules = new Hashtable<>();
		subgrammars = new Hashtable<>();
	}

	/**
	 * Create a new weighted grammar that uses the specified source of
	 * randomness.
	 * 
	 * @param source
	 *            The source of randomness to use
	 */
	public WeightedGrammar(Random source) {
		this();

		rng = source;
	}

	/**
	 * Add a case to an already existing rule.
	 * 
	 * @param rule
	 *            The rule to add a case to.
	 * @param probability
	 *            The probability for this rule to be chosen.
	 * @param cse
	 *            The case being added.
	 */
	public void addCase(E rule, int probability, FunctionalList<E> cse) {
		rules.get(rule).addProbability(probability, cse);
	}

	/**
	 * Add a alias for an existing subgrammar
	 * 
	 * @param name
	 *            The name of the subgrammar to alias
	 * @param alias
	 *            The alias of the subgrammar
	 * @return Whether the alias was succesfully created
	 */
	public boolean addGrammarAlias(E name, E alias) {
		if (subgrammars.containsKey(alias)) {
			return false;
		} else {
			if (subgrammars.containsKey(name)) {
				subgrammars.put(alias, subgrammars.get(name));
				return true;
			} else {
				return false;
			}
		}
	}

	/**
	 * Add a new rule with no cases.
	 * 
	 * @param name
	 *            The name of the rule to add.
	 * @return Whether or not the rule was succesfully added.
	 */
	public boolean addRule(E name) {
		if (rng == null) {
			rng = new Random();
		}
		return addRule(name, new WeightedRandom<>(rng));
	}

	/**
	 * Add a new rule with a set of cases.
	 * 
	 * @param name
	 *            The name of the rule to add.
	 * @param cases
	 *            The set of cases for the rule.
	 * @return Whether or not the rule was succesfully added.
	 */
	public boolean addRule(E name,
			WeightedRandom<FunctionalList<E>> cases) {
		if (rules.containsKey(name)) {
			return false;
		} else {
			rules.put(name, cases);
			return true;
		}
	}

	/**
	 * Add a subgrammar.
	 * 
	 * @param name
	 *            The name of the subgrammar.
	 * @param subgrammar
	 *            The subgrammar to add.
	 * @return Whether or not the subgrammar was succesfully added.
	 */
	public boolean addSubgrammar(E name, WeightedGrammar<E> subgrammar) {
		if (subgrammars.containsKey(name)) {
			return false;
		} else {
			subgrammars.put(name, subgrammar);
			return true;
		}
	}

	/**
	 * Generate a set of debug sentences for the specified rule. Only
	 * generates sentances one layer deep.
	 * 
	 * @param ruleName
	 *            The rule to test.
	 * @return A set of sentances generated by the specified rule.
	 */
	public FunctionalList<FunctionalList<E>>
			generateDebugValues(E ruleName) {
		FunctionalList<FunctionalList<E>> returnedList =
				new FunctionalList<>();

		WeightedRandom<FunctionalList<E>> ruleGenerator =
				rules.get(ruleName);

		for (int i = 0; i < 10; i++) {
			returnedList.add(ruleGenerator.generateValue());
		}

		return returnedList;
	}

	/**
	 * Generate a generic sentance from a initial rule.
	 * 
	 * @param <T>
	 *            The type of the transformed output
	 * 
	 * @param initRule
	 *            The initial rule to start with.
	 * @param tokenTransformer
	 *            The function to transform grammar output into something.
	 * @param spacer
	 *            The spacer element to add in between output tokens.
	 * @return A randomly generated sentance from the specified initial
	 *         rule.
	 */
	public <T> FunctionalList<T> generateGenericValues(E initRule,
			Function<E, T> tokenTransformer, T spacer) {
		FunctionalList<T> returnedList = new FunctionalList<>();

		if (subgrammars.containsKey(initRule)) {
			subgrammars.get(initRule).generateGenericValues(initRule,
					tokenTransformer, spacer).forEach(rulePart -> {
						returnedList.add(rulePart);
						returnedList.add(spacer);
					});
		} else if (rules.containsKey(initRule)) {
			rules.get(initRule).generateValue()
					.forEach(rulePart -> generateGenericValues(rulePart,
							tokenTransformer, spacer)
									.forEach(generatedRulePart -> {
										returnedList
												.add(generatedRulePart);
										returnedList.add(spacer);
									}));
		} else {
			returnedList.add(tokenTransformer.apply(initRule));
			returnedList.add(spacer);
		}

		return returnedList;
	}

	/**
	 * Generate a random list of grammar elements from a given initial
	 * rule.
	 * 
	 * @param initRule
	 *            The initial rule to start with.
	 * @param spacer
	 *            The item to use to space the list.
	 * @return A list of random grammar elements generated by the specified
	 *         rule.
	 */
	public FunctionalList<E> generateListValues(E initRule, E spacer) {
		return generateGenericValues(initRule, strang -> strang, spacer);
	}

	/**
	 * Get the initial rule of this grammar
	 * 
	 * @return The initial rule of this grammar
	 */
	public String getInitialRule() {
		return initialRule;
	}

	/**
	 * Get the subgrammar with the specified name.
	 * 
	 * @param name
	 *            The name of the subgrammar to get.
	 * @return The subgrammar with the specified name.
	 */
	public WeightedGrammar<E> getSubgrammar(E name) {
		return subgrammars.get(name);
	}

	/**
	 * Check if this grammar has an initial rule
	 * 
	 * @return Whether or not this grammar has an initial rule
	 */
	public boolean hasInitialRule() {
		return initialRule != null && !initialRule.equalsIgnoreCase("");
	}

	/**
	 * Prefix a given rule with a token multiple times
	 * 
	 * @param ruleName
	 *            The name of the rule to prefix
	 * @param prefixToken
	 *            The token to prefix to the rules
	 * @param additionalProbability
	 *            The additional probability of the tokens
	 * @param numberOfTimes
	 *            The number of times to prefix the token
	 */
	public void multiPrefixRule(E ruleName, E prefixToken,
			int additionalProbability, int numberOfTimes) {
		WeightedRandom<FunctionalList<E>> rule = rules.get(ruleName);

		FunctionalList<Pair<Integer, FunctionalList<E>>> newResults =
				new FunctionalList<>();

		rule.getValues().forEach((pair) -> {
			FunctionalList<FunctionalList<E>> newRule =
					new FunctionalList<>();

			for (int i = 1; i <= numberOfTimes; i++) {
				FunctionalList<E> newCase =
						pair.merge((left, right) -> right.clone());

				for (int j = 1; j <= i; j++) {
					newCase.prepend(prefixToken);
				}

				newRule.add(newCase);
			}

			newRule.forEach(
					(list) -> newResults
							.add(new Pair<>(
									pair.merge((left, right) -> left)
											+ additionalProbability,
									list)));
		});

		newResults.forEach((pair) -> pair
				.doWith((left, right) -> addCase(ruleName, left, right)));
	}

	/**
	 * Create a series of alternatives for a rule by prefixing them with a
	 * given token
	 * 
	 * @param additionalProbability
	 *            The amount to adjust the probability by
	 * @param ruleName
	 *            The name of the rule to prefix
	 * @param prefixToken
	 *            The token to prefix to the rule
	 */
	public void prefixRule(E ruleName, E prefixToken,
			int additionalProbability) {
		WeightedRandom<FunctionalList<E>> rule = rules.get(ruleName);

		FunctionalList<Pair<Integer, FunctionalList<E>>> newResults =
				new FunctionalList<>();

		rule.getValues().forEach((pair) -> {
			FunctionalList<E> newCase =
					pair.merge((left, right) -> right.clone());
			newCase.prepend(prefixToken);

			newResults.add(new Pair<>(pair.merge((left, right) -> left)
					+ additionalProbability, newCase));
		});

		newResults.forEach((pair) -> pair
				.doWith((left, right) -> addCase(ruleName, left, right)));
	}

	/**
	 * Remove a rule with the specified name.
	 * 
	 * @param name
	 *            The name of the rule to remove.
	 */
	public void deleteRule(E name) {
		rules.remove(name);
	}

	/**
	 * Remove a subgrammar with the specified name.
	 * 
	 * @param name
	 *            The name of the subgrammar to remove.
	 */
	public void deleteSubgrammar(E name) {
		subgrammars.remove(name);
	}

	/**
	 * Returns the number of rules in this grammar
	 * 
	 * @return The number of rules in this grammar
	 */
	public int getRuleCount() {
		return rules.size();
	}

	/**
	 * Returns a set containing all of the rules in this grammar
	 * 
	 * @return The set of all rule names in this grammar
	 */
	public Set<E> getRuleNames() {
		return rules.keySet();
	}

	/**
	 * Set the initial rule of the graphic
	 * 
	 * @param initRule
	 *            The initial rule of this grammar
	 */
	public void setInitialRule(String initRule) {
		this.initialRule = initRule;
	}

	/**
	 * Suffix a token to a rule
	 * 
	 * @param ruleName
	 *            The rule to suffix
	 * @param prefixToken
	 *            The token to prefix to the rule
	 * @param additionalProbability
	 *            Additional probability of the prefixed rule
	 */
	public void suffixRule(E ruleName, E prefixToken,
			int additionalProbability) {
		WeightedRandom<FunctionalList<E>> rule = rules.get(ruleName);

		FunctionalList<Pair<Integer, FunctionalList<E>>> newResults =
				new FunctionalList<>();

		rule.getValues().forEach((par) -> {
			FunctionalList<E> newCase =
					par.merge((left, right) -> right.clone());
			newCase.add(prefixToken);

			newResults.add(new Pair<>(par.merge((left, right) -> left)
					+ additionalProbability, newCase));
		});

		newResults.forEach((pair) -> pair
				.doWith((left, right) -> addCase(ruleName, left, right)));
	}
}