package dkohl.bayes.builders;

import java.util.HashSet;

import net.woodyfolsom.cs6601.p2.Diner;
import net.woodyfolsom.cs6601.p2.Ingredient.TYPE;
import net.woodyfolsom.cs6601.p2.Ingredients;
import net.woodyfolsom.cs6601.p2.Recipe;
import net.woodyfolsom.cs6601.p2.RecipeBook;
import net.woodyfolsom.cs6601.p2.Survey;
import dkohl.bayes.bayesnet.BayesNet;
import dkohl.bayes.estimation.MaximumLikelihoodEstimation;
import dkohl.bayes.example.builders.FoodExampleBuilder;
import dkohl.bayes.probability.Assignment;
import dkohl.bayes.probability.Probability;
import dkohl.bayes.probability.Variable;
import dkohl.bayes.probability.distribution.ContinousDistribution;
import dkohl.bayes.probability.distribution.ProbabilityDistribution;
import dkohl.bayes.probability.distribution.ProbabilityTable;
import dkohl.bayes.statistic.DataPoint;
import dkohl.bayes.statistic.DataSet;
import dkohl.onthology.Ontology;

public class FoodNetBuilder {
    public static final String TASTE = "Taste";

    public static final String SOMEONE_VEGETARIAN = "vegetarian";

    public static final String CONTAINS_MEAT = "Meat";
    public static final String CONTAINS_VEGETABLE = "Vegetable";
    public static final String CONTAINS_BEEF = TYPE.BEEF.toString();
    public static final String CONTAINS_PORK = TYPE.PORK.toString();
    public static final String CONTAINS_TOMATOS = TYPE.TOMATO.toString();
    public static final String CONTAINS_POTATOS = TYPE.POTATO.toString();

    public static final String TRUE_VALUE = "true";
    public static final String FALSE_VALUE = "false";

    public static final String DOMAIN[] = { TRUE_VALUE, FALSE_VALUE };

    public static final String RATING_DOMAIN[] = { "1", "2", "3", "4", "5",
	    "6", "7", "8", "9", "10" };

    private static final String[] VARIABLES = { SOMEONE_VEGETARIAN, 
    CONTAINS_BEEF, CONTAINS_MEAT, CONTAINS_PORK, CONTAINS_POTATOS,
	    CONTAINS_TOMATOS, CONTAINS_VEGETABLE, TASTE };

    private static final String[] OBSERVED = { CONTAINS_BEEF, CONTAINS_PORK,
	    CONTAINS_POTATOS, CONTAINS_TOMATOS };

    /**
     * Build the onthology:
     * 
     * 	Ingredient -> category
     * 
     * @return
     */
    public static Ontology createOntology() {
	HashSet<String> classes = new HashSet<String>();

	classes.add(CONTAINS_MEAT);
	classes.add(CONTAINS_VEGETABLE);

	Ontology ontology = new Ontology(classes);

	ontology.define(CONTAINS_PORK, CONTAINS_MEAT);
	ontology.define(CONTAINS_BEEF, CONTAINS_MEAT);

	ontology.define(CONTAINS_TOMATOS, CONTAINS_VEGETABLE);
	ontology.define(CONTAINS_POTATOS, CONTAINS_VEGETABLE);
	return ontology;
    }

    /**
     * Build an assignment given a variable and a value
     * @param varible 
     * @param value
     * @return
     */
    private static Assignment build(String varible, String value) {
	return new Assignment(new Variable(varible, DOMAIN), value);
    }
    
    /**
     * Resolve onthology by adding category for
     * each ingredient and add all categories 
     * that are not in a dish as false.
     * 
     * @param point
     * @param onto
     * @return
     */
    public static DataPoint normalize(DataPoint point, Ontology onto) {
	// resolve onthology
	DataPoint normPoint = new DataPoint(point);
	for (String key : point.keySet()) {
	    if (onto.getInheritance().containsKey(key)) {
		if(point.get(key).getValue().equals(TRUE_VALUE)){
		    normPoint.add(build(onto.getInheritance().get(key), TRUE_VALUE));
		} else {
		    normPoint.add(build(onto.getInheritance().get(key), FALSE_VALUE));
		}
	    }
	}

	// implement closed world assumption
	// everything unknown is false
	for (String variable : VARIABLES) {
	    if (!normPoint.containsKey(variable)) {
		normPoint.add(build(variable, FALSE_VALUE));
	    }
	}
	return normPoint;
    }
    
    /**
     * Build training data from survey.
     * 
     * @param survey
     * @param recipeBook
     * @param startIndex
     * @param endIndex
     * @return
     */
    public static DataSet getSurveyDataSet(Survey survey,
	    RecipeBook recipeBook, int startIndex, int endIndex) {
	DataSet data = FoodExampleBuilder.examples();
	Ontology onto = createOntology();

	for (int dinerIndex = 0; dinerIndex < survey.getDinerCount(); dinerIndex++) {
	    Diner diner = survey.getDiner(dinerIndex);
	    for (int dishIndex = startIndex; dishIndex < endIndex; dishIndex++) {
		data.add(normalize(
			createDataPoint(recipeBook, survey.getDish(dishIndex),
				diner.getRating(dishIndex)), onto));
	    }
	}
	return data;
    }

    /**
     * Build a data point with as rating.
     * If something rated above a threshold
     * it is true, else it is false
     * 
     * @param recipeBook
     * @param recipeName
     * @param weight
     * @return
     */
    public static DataPoint createDataPoint(RecipeBook recipeBook,
	    String recipeName, int weight) {
	DataPoint point = new DataPoint();

	Recipe recipe = recipeBook.getRecipe(recipeName);
	Ingredients ingredients = recipe.getIngredients();

	if (ingredients.contains(TYPE.BEEF)) {
	    if(weight > 5) {
		point.add(build(CONTAINS_BEEF, TRUE_VALUE));
	    } else {
		point.add(build(CONTAINS_BEEF, FALSE_VALUE));
	    }
	}
	if (ingredients.contains(TYPE.PORK)) {
	    if(weight > 5) {
		point.add(build(CONTAINS_PORK, TRUE_VALUE));
	    } else {
		point.add(build(CONTAINS_PORK, FALSE_VALUE));
	    }
	}
	if (ingredients.contains(TYPE.POTATO)) {
	    if(weight > 5) {
		point.add(build(CONTAINS_POTATOS, TRUE_VALUE));
	    } else {
		point.add(build(CONTAINS_POTATOS, FALSE_VALUE));
	    }
	}
	if (ingredients.contains(TYPE.TOMATO)) {
	    if(weight > 5) {
		point.add(build(CONTAINS_TOMATOS, TRUE_VALUE));
	    } else {
		point.add(build(CONTAINS_TOMATOS, FALSE_VALUE));
	    }	
	}
	point.add(build(TASTE, "" + weight));

	return point;
    }

    /**
     * Vegetarian control variable
     * 
     * @param survey
     * @return
     */
    public static ProbabilityDistribution vegi(Survey survey) {
	String names[] = {CONTAINS_MEAT, SOMEONE_VEGETARIAN};
	ProbabilityTable table = new ProbabilityTable(names);
	table.setProbabilityForAssignment("true;true;", new Probability(0));
	table.setProbabilityForAssignment("false;true;", new Probability(1));
	table.setProbabilityForAssignment("true;false;", new Probability(1));
	table.setProbabilityForAssignment("false;false;", new Probability(1));
	return table;
    }
    
    /**
     * Initialize probability distributions
     */
    
    public static ProbabilityDistribution beef() {
	String names[] = { CONTAINS_MEAT, CONTAINS_BEEF };
	ProbabilityTable table = new ProbabilityTable(names);
	return table;
    }

    public static ProbabilityDistribution pork() {
	String names[] = { CONTAINS_MEAT, CONTAINS_PORK };
	ProbabilityTable table = new ProbabilityTable(names);
	return table;
    }

    public static ProbabilityDistribution meat() {
	String names[] = { CONTAINS_MEAT };
	ProbabilityTable table = new ProbabilityTable(names);
	return table;
    }

    public static ProbabilityDistribution tomatos() {
	String names[] = { CONTAINS_VEGETABLE, CONTAINS_TOMATOS };
	ProbabilityTable table = new ProbabilityTable(names);
	return table;
    }

    public static ProbabilityDistribution potatos() {
	String names[] = { CONTAINS_VEGETABLE, CONTAINS_POTATOS };
	ProbabilityTable table = new ProbabilityTable(names);
	return table;
    }

    public static ProbabilityDistribution vegetables() {
	String names[] = { CONTAINS_VEGETABLE };
	ProbabilityTable table = new ProbabilityTable(names);
	return table;
    }

    public static ProbabilityDistribution taste() {
	String names[] = { TASTE, CONTAINS_BEEF, CONTAINS_PORK,
		CONTAINS_POTATOS, CONTAINS_TOMATOS /* , CONTAINS_GENERIC_NUTS */};
	ContinousDistribution distribution = new ContinousDistribution(names, 0);
	return distribution;
    }

    public static BayesNet createDishNet(Survey survey, RecipeBook recipeBook,
	    int startIndex, int endIndex) {
	BayesNet net = new BayesNet(VARIABLES);
	
	// initialize distributions
	net.setDistribution(new Variable(CONTAINS_MEAT, DOMAIN), meat());
	net.setDistribution(new Variable(SOMEONE_VEGETARIAN, DOMAIN),vegi(survey));
	net.setDistribution(new Variable(CONTAINS_VEGETABLE, DOMAIN),vegetables());
	net.setDistribution(new Variable(CONTAINS_BEEF, DOMAIN), beef());
	net.setDistribution(new Variable(CONTAINS_PORK, DOMAIN), pork());
	net.setDistribution(new Variable(CONTAINS_POTATOS, DOMAIN), potatos());
	net.setDistribution(new Variable(CONTAINS_TOMATOS, DOMAIN), tomatos());
	net.setDistribution(new Variable(TASTE, RATING_DOMAIN), taste());

	Ontology ontology = createOntology();

	// Build net
	net.connect(SOMEONE_VEGETARIAN, CONTAINS_MEAT);

	for (String thing : OBSERVED) {
	    net.connect(thing, ontology.getInheritance().get(thing));
	    net.connect(TASTE, thing);
	}

	DataSet dataSet = getSurveyDataSet(survey, recipeBook, startIndex,
		endIndex);

	for (String category : ontology.getClasses()) {
	    MaximumLikelihoodEstimation.estimate(dataSet, net, category);
	    for (String thing : ontology.getClasses2thing().get(category)) {
		MaximumLikelihoodEstimation.estimate(dataSet, net, thing);
	    }
	}

	MaximumLikelihoodEstimation.estimate(dataSet, net, TASTE);
	ContinousDistribution distribturion = (ContinousDistribution) net
		.getNodes().get(TASTE);

	distribturion.estimate();

	return net;
    }
}