package dkohl.bayes.example.builders;

import java.util.LinkedList;

import dkohl.bayes.bayesnet.BayesNet;
import dkohl.bayes.probability.Assignment;
import dkohl.bayes.probability.Probability;
import dkohl.bayes.probability.Variable;
import dkohl.bayes.probability.distribution.ProbabilityTable;

/**
 * The alarm example for baysian nets.
 * 
 * Stuart Russel, Peter Norvig: Artificial Intelligence: A Modern Approach, 3ed
 * Edition, Prentice Hall, 2010
 * 
 * @author Daniel Kohlsdorf
 */
public class AlarmNetBuilderTable {

	// Variable names
	public static final String BURGLARY = "Burglary";
	public static final String EARTHQUAKE = "Earthquake";
	public static final String ALARM = "Alarm";
	public static final String JOHN = "John";
	public static final String MARRY = "Marry";

	// Possible outcomes
	public static final String TRUE = "true";
	public static final String FALSE = "false";

	// A variables domain
	public static final String DOMAIN[] = { TRUE, FALSE };

	// A set of variables
	public static final String VARIABLES[] = { BURGLARY, EARTHQUAKE, ALARM,
			JOHN, MARRY };

	/**
	 * Builds the query from the book: P(B| j, m)
	 */
	public static LinkedList<Assignment> completeQueryBulgary() {
		LinkedList<Assignment> assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(JOHN, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(MARRY, DOMAIN), TRUE));
		return new LinkedList<Assignment>(assignment);
	}

	/**
	 * Build burglary prior
	 * 
	 * @param alarmNet
	 */
	private static void burglary(BayesNet alarmNet) {
		Probability p_bulglary = new Probability(0.001);
		String names[] = { BURGLARY };
		ProbabilityTable burglary = new ProbabilityTable(names);
		burglary.setProbabilityForAssignment("true;", p_bulglary);
		burglary.setProbabilityForAssignment("false;", p_bulglary.rest());
		alarmNet.setDistribution(new Variable(BURGLARY, DOMAIN), burglary);
	}

	/**
	 * Build earthquake prior
	 * 
	 * @param alarmNet
	 */
	private static void earthquake(BayesNet alarmNet) {
		Probability p_earthquake = new Probability(0.002);
		String names[] = { EARTHQUAKE };
		ProbabilityTable earthquake = new ProbabilityTable(names);
		earthquake.setProbabilityForAssignment("true;", p_earthquake);
		earthquake.setProbabilityForAssignment("false;", p_earthquake.rest());

		alarmNet.setDistribution(new Variable(EARTHQUAKE, DOMAIN), earthquake);
	}

	/**
	 * Jhon calls!
	 * 
	 * @param alarmNet
	 */
	private static void jhon(BayesNet alarmNet) {
		// P(ALARM) == true
		Probability t = new Probability(.90);

		// P(ALARM) == false
		Probability f = new Probability(.05);
		String names[] = { ALARM, JOHN };
		ProbabilityTable jhon = new ProbabilityTable(names);
		jhon.setProbabilityForAssignment("true;true;", t);
		jhon.setProbabilityForAssignment("false;true;", f);
		jhon.setProbabilityForAssignment("true;false", t.rest());
		jhon.setProbabilityForAssignment("false;false;", f.rest());

		alarmNet.setDistribution(new Variable(JOHN, DOMAIN), jhon);
	}

	/**
	 * Marry calls!
	 * 
	 * @param alarmNet
	 */
	private static void mary(BayesNet alarmNet) {
		// P(ALARM) == true
		Probability t = new Probability(.70);

		// P(ALARM) == false
		Probability f = new Probability(.01);
		String names[] = { ALARM, MARRY };
		ProbabilityTable marry = new ProbabilityTable(names);
		marry.setProbabilityForAssignment("true;true;", t);
		marry.setProbabilityForAssignment("false;true;", f);
		marry.setProbabilityForAssignment("true;false", t.rest());
		marry.setProbabilityForAssignment("false;false;", f.rest());
		alarmNet.setDistribution(new Variable(MARRY, DOMAIN), marry);
	}

	/**
	 * The alarm goes off!
	 * 
	 * @param alarmNet
	 */
	private static void alarm(BayesNet alarmNet) {
		// P(ALARM | BURGLARY = true, EARTHQUAKE = true)
		Probability tt = new Probability(.95);

		// P(ALARM | BURGLARY = true, EARTHQUAKE = false)
		Probability tf = new Probability(.94);

		// P(ALARM | BURGLARY = false, EARTHQUAKE = true)
		Probability ft = new Probability(.29);

		// P(ALARM | BURGLARY = false, EARTHQUAKE = false)
		Probability ff = new Probability(.001);
		String names[] = { BURGLARY, EARTHQUAKE, ALARM };
		ProbabilityTable alarm = new ProbabilityTable(names);
		alarm.setProbabilityForAssignment(TRUE + ";" + TRUE + ";" + TRUE + ";",
				tt);
		alarm.setProbabilityForAssignment(
				TRUE + ";" + FALSE + ";" + TRUE + ";", tf);
		alarm.setProbabilityForAssignment(
				FALSE + ";" + TRUE + ";" + TRUE + ";", ft);
		alarm.setProbabilityForAssignment(FALSE + ";" + FALSE + ";" + TRUE
				+ ";", ff);

		alarm.setProbabilityForAssignment(
				TRUE + ";" + TRUE + ";" + FALSE + ";", tt.rest());
		alarm.setProbabilityForAssignment(TRUE + ";" + FALSE + ";" + FALSE
				+ ";", tf.rest());
		alarm.setProbabilityForAssignment(FALSE + ";" + TRUE + ";" + FALSE
				+ ";", ft.rest());
		alarm.setProbabilityForAssignment(FALSE + ";" + FALSE + ";" + FALSE
				+ ";", ff.rest());
		alarmNet.setDistribution(new Variable(ALARM, DOMAIN), alarm);
	}

	public static BayesNet alarm() {
		BayesNet alarmNet = new BayesNet(VARIABLES);

		// set probability tables and priors
		burglary(alarmNet);
		earthquake(alarmNet);
		alarm(alarmNet);
		jhon(alarmNet);
		mary(alarmNet);

		// construct the graph
		alarmNet.connect(ALARM, BURGLARY);
		alarmNet.connect(ALARM, EARTHQUAKE);
		alarmNet.connect(JOHN, ALARM);
		alarmNet.connect(MARRY, ALARM);

		return alarmNet;
	}
}