package dkohl.bayes.probability.distribution;

import java.util.HashMap;
import java.util.LinkedList;

import dkohl.bayes.probability.Assignment;
import dkohl.bayes.probability.Probability;

/**
 * A probability table defined as a table of all combinations of all possible
 * values of the varibales involved.
 * 
 * @author Daniel Kohlsdorf
 */
public class ProbabilityTable implements ProbabilityDistribution {

	/**
	 * The involved variables mapping to row
	 */
	private HashMap<String, Integer> variable2row;

	/**
	 * An P(assignment key) assignment_key var1_value....varN_value : String
	 */
	private HashMap<String, Probability> assignments;

	public ProbabilityTable(String names[]) {
		assignments = new HashMap<String, Probability>();
		this.variable2row = new HashMap<String, Integer>();
		for (int i = 0; i < names.length; i++) {
			variable2row.put(names[i], i);
		}
	}

	public ProbabilityTable(LinkedList<String> names) {
		assignments = new HashMap<String, Probability>();
		this.variable2row = new HashMap<String, Integer>();
		for (int i = 0; i < names.size(); i++) {
			variable2row.put(names.get(i), i);
			System.out.println(names.get(i));
		}
	}

	/**
	 * Sets the probability for a table entry key. You can generate one from
	 * your assignments using generateKey, or just use this same method with
	 * your assignments.
	 * 
	 * @param key
	 *            the entry key
	 * @param probability
	 *            the associated probability
	 */
	public void setProbabilityForAssignment(String key, Probability probability) {
		assignments.put(key, probability);
	}

	/**
	 * Sets the probability for the given assignment
	 * 
	 * @param assignment
	 *            the assignment
	 * @param probability
	 *            the associated probability
	 */
	public void setProbabilityForAssignment(LinkedList<Assignment> assignment,
			Probability probability) {
		String key = generateKey(assignment);
		assignments.put(key, probability);
	}

	@Override
	public Probability eval(LinkedList<Assignment> assignment) {
		String key = generateKey(assignment);
		return assignments.get(key);
	}

	/**
	 * Generates a table entry key for an assignment
	 * 
	 * @param assignment
	 *            the assignment
	 * 
	 * @return the key
	 */
	public String generateKey(LinkedList<Assignment> assignment) {
		String[] values = new String[variable2row.size()];
		for (Assignment col : assignment) {
			if (variable2row.containsKey(col.getVariable().getName())) {
				int row = variable2row.get(col.getVariable().getName());
				values[row] = col.getValue();
			}
		}
		String key = "";
		for (String entry : values) {
			key += entry + ";";
		}
		return key;
	}

	public HashMap<String, Probability> getAssignments() {
		return assignments;
	}

	public String[] getNames() {
		String[] names = new String[variable2row.size()];
		for (String name : variable2row.keySet()) {
			names[variable2row.get(name)] = name;
		}
		return names;
	}
}
