/*
 * Copyright (c) 2009 University of Tartu
 */
package org.jpmml.manager;

import java.util.*;

import org.dmg.pmml.*;

public class TreeModelManager extends ModelManager<TreeModel> {

	private TreeModel treeModel = null;

	private Node root = null;


	public TreeModelManager(){
	}

	public TreeModelManager(PMML pmml){
		this(pmml, find(pmml.getContent(), TreeModel.class));
	}

	public TreeModelManager(PMML pmml, TreeModel treeModel){
		super(pmml);

		this.treeModel = treeModel;
	}

	public String getSummary(){
		return "Tree";
	}

	@Override
	public TreeModel getModel(){
		ensureNotNull(this.treeModel);

		return this.treeModel;
	}

	public TreeModel createClassificationModel(){
		return createModel(MiningFunctionType.CLASSIFICATION);
	}

	/**
	 * @throws ModelManagerException If the Model already exists
	 *
	 * @see #getModel()
	 */
	public TreeModel createModel(MiningFunctionType miningFunction){
		ensureNull(this.treeModel);

		this.treeModel = new TreeModel(new MiningSchema(), new Node(), miningFunction);

		getModels().add(this.treeModel);

		return this.treeModel;
	}

	/**
	 * @return The root Node
	 */
	public Node getOrCreateRoot(){

		if(this.root == null){
			TreeModel treeModel = getModel();

			this.root = treeModel.getNode();
			if(this.root == null){
				this.root = new Node();

				treeModel.setNode(this.root);
			}

			Predicate predicate = this.root.getPredicate();
			if(predicate == null){
				this.root.setPredicate(new True());
			}
		}

		return this.root;
	}

	/**
	 * Adds a new Node to the root Node.
	 *
	 * @return The newly added Node
	 *
	 * @see #getOrCreateRoot()
	 */
	public Node addNode(Predicate predicate){
		return addNode(getOrCreateRoot(), predicate);
	}

	/**
	 * Adds a new Node to the specified Node.
	 *
	 * @return The newly added Node
	 */
	public Node addNode(Node parentNode, Predicate predicate){
		Node node = new Node();
		node.setPredicate(predicate);

		parentNode.getNodes().add(node);

		return node;
	}

	public ScoreDistribution getOrAddScoreDistribution(Node node, String value){
		List<ScoreDistribution> scoreDistributions = node.getScoreDistributions();

		for(ScoreDistribution scoreDistribution : scoreDistributions){

			if((scoreDistribution.getValue()).equals(value)){
				return scoreDistribution;
			}
		}

		ScoreDistribution scoreDistribution = new ScoreDistribution(value, 0);
		scoreDistributions.add(scoreDistribution);

		return scoreDistribution;
	}
}
@@ -246,16 +246,50 @@
  null;%0A%0A
+%09private Set%3CString%3E ids = null;%0A%0A
 %0A%09public
@@ -2477,9 +2477,461 @@
 ion;%0A%09%7D%0A
+%0A%09public Set%3CString%3E getAllNodesId() %7B%0A%09%09if (ids == null) %7B%0A%09%09%09// Fill ids%0A%09%09%09Node root = getOrCreateRoot();%0A%09%09%09ids = new HashSet%3CString%3E();%0A%09%09%09// This function fills ids.%0A%09%09%09getAllNodesIdFrom(root);%0A%09%09%7D%0A%0A%0A%09%09return ids;%0A%09%7D%0A%0A%09private void getAllNodesIdFrom(Node current) %7B%0A%09%09if (current != null) %7B%0A%09%09%09if (current.getId() != null) %7B%0A%09%09%09%09ids.add(current.getId());%0A%09%09%09%7D%0A%09%09%09for (Node child : current.getNodes()) %7B%0A%09%09%09%09getAllNodesIdFrom(child);%0A%09%09%09%7D%0A%09%09%7D%0A%09%7D%0A%0A
 %7D

Add a method to get all the node id.