package dkohl.bayes.example.builders;

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;
import dkohl.bayes.statistic.DataPoint;
import dkohl.bayes.statistic.DataSet;

/**
 * The Sprinkler net example
 * 
 * http://www.cs.ubc.ca/~murphyk/Bayes/bnintro.html
 * 
 * @author Daniel Kohlsdorf
 */
public class EstimateSprinklerNetBuilderTable {

	// Variable names
	public static final String CLOUDY = "Cloudy";
	public static final String SPRINKLER = "Sprinkler";
	public static final String GRASS_WET = "GrassWet";
	public static final String RAIN = "Rain";

	// 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[] = { CLOUDY, SPRINKLER, GRASS_WET,
			RAIN };

	private static void cloudy(BayesNet sprinklerNet) {
		Probability p_cloudy = new Probability(0.5);
		String names[] = { CLOUDY };
		ProbabilityTable cloudy = new ProbabilityTable(names);
		cloudy.setProbabilityForAssignment("true;", p_cloudy);
		cloudy.setProbabilityForAssignment("false;", p_cloudy.rest());
		sprinklerNet.setDistribution(new Variable(CLOUDY, DOMAIN), cloudy);
	}

	private static void rain(BayesNet sprinklerNet) {
		Probability p_cloudy = new Probability(0.8);
		Probability p_notcloudy = new Probability(0.2);

		String names[] = { CLOUDY, RAIN };
		ProbabilityTable rain = new ProbabilityTable(names);
		rain.setProbabilityForAssignment("true;true;", p_cloudy);
		rain.setProbabilityForAssignment("false;false;", p_notcloudy);

		rain.setProbabilityForAssignment("false;true;", p_notcloudy.rest());
		rain.setProbabilityForAssignment("true;false;", p_cloudy.rest());
		sprinklerNet.setDistribution(new Variable(RAIN, DOMAIN), rain);
	}

	private static void sprinkler(BayesNet sprinklerNet) {
		Probability p_cloudy = new Probability(0.1);
		Probability p_notcloudy = new Probability(0.5);

		String names[] = { CLOUDY, SPRINKLER };
		ProbabilityTable sprinkler = new ProbabilityTable(names);
		sprinkler.setProbabilityForAssignment("true;true;", p_cloudy);
		sprinkler.setProbabilityForAssignment("false;false;", p_notcloudy);

		sprinkler
				.setProbabilityForAssignment("false;true;", p_notcloudy.rest());
		sprinkler.setProbabilityForAssignment("true;false;", p_cloudy.rest());
		sprinklerNet
				.setDistribution(new Variable(SPRINKLER, DOMAIN), sprinkler);
	}

	private static void grass(BayesNet sprinklerNet) {
		String names[] = { RAIN, SPRINKLER, GRASS_WET };
		ProbabilityTable sprinkler = new ProbabilityTable(names);
		sprinklerNet
				.setDistribution(new Variable(GRASS_WET, DOMAIN), sprinkler);
	}

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

		cloudy(sprinkler);
		rain(sprinkler);
		sprinkler(sprinkler);
		grass(sprinkler);

		sprinkler.connect(RAIN, CLOUDY);
		sprinkler.connect(SPRINKLER, CLOUDY);
		sprinkler.connect(GRASS_WET, RAIN);
		sprinkler.connect(GRASS_WET, SPRINKLER);

		return sprinkler;
	}

	private static Assignment build(String varible, String value) {
		return new Assignment(new Variable(varible, DOMAIN), value);
	}

	public static DataSet dataSet() {
		DataSet dataSet = new DataSet();

		/**
		 * If rain is false and sprinkler is false, grass is never wet.
		 */
		DataPoint one = new DataPoint();
		one.add(build(RAIN, FALSE));
		one.add(build(SPRINKLER, FALSE));
		one.add(build(GRASS_WET, FALSE));
		dataSet.add(one);

		/**
		 * 1 / 10 times the grass is not wet when the sprinkler is on
		 */
		DataPoint two = new DataPoint();
		two.add(build(RAIN, FALSE));
		two.add(build(SPRINKLER, TRUE));
		two.add(build(GRASS_WET, FALSE));
		dataSet.add(two);

		/**
		 * 9 / 10 times the sprinkler is on and the grass is wet
		 */
		for (int i = 0; i < 9; i++) {
			DataPoint point = new DataPoint();
			point.add(build(RAIN, FALSE));
			point.add(build(SPRINKLER, TRUE));
			point.add(build(GRASS_WET, TRUE));
			dataSet.add(point);
		}

		/**
		 * 1 / 10 times the grass is not wet when it rains
		 */
		DataPoint three = new DataPoint();
		three.add(build(RAIN, TRUE));
		three.add(build(SPRINKLER, FALSE));
		three.add(build(GRASS_WET, FALSE));
		dataSet.add(three);

		/**
		 * 9 / 10 times it rains and the grass is wet
		 */
		for (int i = 0; i < 9; i++) {
			DataPoint point = new DataPoint();
			point.add(build(RAIN, TRUE));
			point.add(build(SPRINKLER, FALSE));
			point.add(build(GRASS_WET, TRUE));
			dataSet.add(point);
		}

		/**
		 * 1 / 100 times the grass is not wet when it rains and the sprinkler is
		 * on
		 */
		DataPoint four = new DataPoint();
		four.add(build(RAIN, TRUE));
		four.add(build(SPRINKLER, TRUE));
		four.add(build(GRASS_WET, FALSE));
		dataSet.add(four);

		/**
		 * 99 / 100 times it rains and the grass is wet
		 */
		for (int i = 0; i < 99; i++) {
			DataPoint point = new DataPoint();
			point.add(build(RAIN, TRUE));
			point.add(build(SPRINKLER, TRUE));
			point.add(build(GRASS_WET, TRUE));
			dataSet.add(point);
		}

		return dataSet;
	}

}