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.ProbabilityTree;

public class AlarmNetBuilderTree {

	// 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 sprinklerNet
	 */
	private static void burglary(BayesNet sprinklerNet) {
		Probability p_bulglary = new Probability(0.001);

		ProbabilityTree burglary = new ProbabilityTree();

		LinkedList<Assignment> assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), TRUE));
		burglary.setProbabilityForAssignment(assignment, p_bulglary);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), FALSE));
		burglary.setProbabilityForAssignment(assignment, p_bulglary.rest());

		sprinklerNet.setDistribution(new Variable(BURGLARY, DOMAIN), burglary);
	}

	/**
	 * Build earthquake prior
	 * 
	 * @param sprinklerNet
	 */
	private static void earthquake(BayesNet sprinklerNet) {
		Probability p_earthquake = new Probability(0.002);

		ProbabilityTree earthquake = new ProbabilityTree();

		LinkedList<Assignment> assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), TRUE));
		earthquake.setProbabilityForAssignment(assignment, p_earthquake);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), FALSE));
		earthquake.setProbabilityForAssignment(assignment, p_earthquake.rest());

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

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

		// P(ALARM) == false
		Probability f = new Probability(.05);

		ProbabilityTree jhon = new ProbabilityTree();

		LinkedList<Assignment> assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(JOHN, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), TRUE));
		jhon.setProbabilityForAssignment(assignment, t);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(JOHN, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), FALSE));
		jhon.setProbabilityForAssignment(assignment, f);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(JOHN, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), TRUE));
		jhon.setProbabilityForAssignment(assignment, t.rest());

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(JOHN, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), FALSE));
		jhon.setProbabilityForAssignment(assignment, f.rest());

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

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

		// P(ALARM) == false
		Probability f = new Probability(.01);
		ProbabilityTree mary = new ProbabilityTree();

		LinkedList<Assignment> assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(MARRY, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), TRUE));
		mary.setProbabilityForAssignment(assignment, t);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(MARRY, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), FALSE));
		mary.setProbabilityForAssignment(assignment, f);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(MARRY, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), TRUE));
		mary.setProbabilityForAssignment(assignment, t.rest());

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(MARRY, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), FALSE));
		mary.setProbabilityForAssignment(assignment, f.rest());
		sprinklerNet.setDistribution(new Variable(MARRY, DOMAIN), mary);
	}

	/**
	 * The alarm goes off!
	 * 
	 * @param sprinklerNet
	 */
	private static void alarm(BayesNet sprinklerNet) {
		// 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);
		ProbabilityTree alarm = new ProbabilityTree();

		LinkedList<Assignment> assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), TRUE));
		alarm.setProbabilityForAssignment(assignment, tt);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), TRUE));
		alarm.setProbabilityForAssignment(assignment, tf);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), TRUE));
		alarm.setProbabilityForAssignment(assignment, ft);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), TRUE));
		alarm.setProbabilityForAssignment(assignment, ff);

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), FALSE));
		alarm.setProbabilityForAssignment(assignment, tt.rest());

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), FALSE));
		alarm.setProbabilityForAssignment(assignment, tf.rest());

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), TRUE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), FALSE));
		alarm.setProbabilityForAssignment(assignment, ft.rest());

		assignment = new LinkedList<Assignment>();
		assignment.add(new Assignment(new Variable(BURGLARY, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(EARTHQUAKE, DOMAIN), FALSE));
		assignment.add(new Assignment(new Variable(ALARM, DOMAIN), FALSE));
		alarm.setProbabilityForAssignment(assignment, ff.rest());

		sprinklerNet.setDistribution(new Variable(ALARM, DOMAIN), alarm);
	}

	public static BayesNet sprinkler() {
		BayesNet sprinklerNet = new BayesNet(VARIABLES);

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

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

		return sprinklerNet;
	}
}