package dkohl.bayes.probability.distribution;

import java.util.LinkedList;

import dkohl.bayes.probability.Assignment;
import dkohl.bayes.probability.Probability;
import dkohl.bayes.probability.distribution.tree.DecisionNode;
import dkohl.bayes.probability.distribution.tree.ProbabilityLeaf;

/**
 * Probability distribution in descision tree representation
 * 
 * DANGER: No node ordering during creation yet.
 * 
 * @author Daniel Kohlsdorf
 */
public class ProbabilityTree implements ProbabilityDistribution {

    /**
     * The trees root node
     */
    private DecisionNode root = null;
    
    /**
     * Initialize the tree 
     * by creating one path, along an assignment of
     * variables and add a probability node at the end.
     * 
     * @param assignment
     * @param probability
     */
    private void initTree(LinkedList<Assignment> assignment, Probability probability) {
	    root = new DecisionNode(assignment.get(0).getVariable().getName());
	    
	    DecisionNode parent = root;
	    for(int i = 1; i < assignment.size(); i++) {
		// add the child to the parent with the key of the parents assignment
		DecisionNode child = new DecisionNode(assignment.get(i).getVariable().getName());
		parent.put(assignment.get(i - 1).getValue(), child);
		parent = child;
	    }
	    parent.put(assignment.getLast().getValue(), new ProbabilityLeaf(probability));
    }
    
    /**
     * Follow existing paths until
     * successors do not contain the current assignment,
     * then start inserting.
     * 
     * @param assignment
     * @param probability
     */
    public void setProbabilityForAssignment(LinkedList<Assignment> assignment, Probability probability) {
	if(root == null) {
	    initTree(assignment, probability);
	} else {
	    DecisionNode parent = root;
	    for(int i = 1; i < assignment.size(); i++) {
		if(parent.getSuccessors().containsKey(assignment.get(i - 1).getValue())) {
		    parent = (DecisionNode) parent.getSuccessors().get(assignment.get(i - 1).getValue());
		} else {
		    DecisionNode child = new DecisionNode(assignment.get(i).getVariable().getName());
		    parent.put(assignment.get(i - 1).getValue(), child);
		    parent = child;
		}
	    }
	    parent.put(assignment.getLast().getValue(), new ProbabilityLeaf(probability));
	}
    }
    
    @Override
    public Probability eval(LinkedList<Assignment> assignment) {
	return root.eval(assignment);
    }

}
