package net.maizegenetics.pal.tree;

import java.util.Arrays;

import net.maizegenetics.pal.ids.Identifier;

/**
 * @author Peter Bradbury 2/23/2010
 * For a Tree and a given number of clusters, this class finds the members of those clusters.
 */
public class TreeClusters {
	Tree theTree;
	int nTaxa;
	double[] sortedHeight;
	int nNodes;
	
	public TreeClusters(Tree theTree) {
		this.theTree = theTree;
		nTaxa = theTree.getIdCount();
		nNodes = theTree.getInternalNodeCount();
		sortedHeight = new double[nNodes];
		for (int n = 0; n < nNodes; n++) {
			sortedHeight[n] = theTree.getInternalNode(n).getNodeHeight();
		}
		Arrays.sort(sortedHeight);
	}
	
	public int[] getGroups (int numberOfGroups) {
		int[] groups = new int[nTaxa];
		if (numberOfGroups == 1) {
			for (int t = 0; t < nTaxa; t++) groups[t] = 0;
		} else if (numberOfGroups == nTaxa) {
			for (int t = 0; t <nTaxa; t++) groups[t] = t;
		} else {
			double maxHeight = sortedHeight[nNodes - numberOfGroups];
			for (int t = 0; t <nTaxa; t++) groups[t] = -1;
			int group = 0;
			for (int t = 0; t < nTaxa; t++) {
				if (groups[t] == -1) {
					Node aNode = theTree.getExternalNode(t);
					while (aNode.getParent().getNodeHeight() <= maxHeight) {
						aNode = aNode.getParent();
					}
					setNodeToGroup(aNode, group, groups);
					group++;
				}
			}
		}
		return groups;
	}
	
	public void setNodeToGroup(Node aNode, int group, int[] groups) {
		if (aNode.isLeaf()) {
			Identifier nodeId = aNode.getIdentifier();
			int index = theTree.whichIdNumber(nodeId);
			groups[index] = group;
		}
		int nChildren = aNode.getChildCount();
		for (int c = 0; c < nChildren; c++) {
			setNodeToGroup(aNode.getChild(c), group, groups);
		}
	}
	
	public Identifier getTaxon(int whichTaxon) { return theTree.getIdentifier(whichTaxon);}
	public Tree getTree() {return theTree;}
}

@@ -668,32 +668,262 @@
 eight);%0D%0A%09%7D%0D%0A%09%0D%0A
+%09/**%0D%0A%09 * For a given number of groups this function returns the group membership for each taxon.%0D%0A%09 * @param numberOfGroups the number of groups desired%0D%0A%09 * @return an array of integers indicating the group of each taxon%0D%0A%09 */%0D%0A
 %09public int%5B%5D ge
@@ -1616,32 +1616,888 @@
 groups;%0D%0A%09%7D%0D%0A%09%0D%0A
+%09/**%0D%0A%09 * This function returns taxa group membership when the tree is cut at the designated height.%0D%0A%09 * @param height the height at which the tree is cut%0D%0A%09 * @return the group membership of each taxon%0D%0A%09 */%0D%0A%09public int%5B%5D getGroups(double height) %7B%0D%0A%09%09int%5B%5D groups = new int%5BnTaxa%5D;%0D%0A%09%09double maxHeight = sortedHeight%5BnNodes -1%5D;%0D%0A%09%09if (maxHeight %3C height) %7B //all individuals are in a single group%0D%0A%09%09%09for (int t = 0; t %3CnTaxa; t++) groups%5Bt%5D = 0;%0D%0A%09%09%7D else %7B%0D%0A%09%09%09for (int t = 0; t %3CnTaxa; t++) groups%5Bt%5D = -1;%0D%0A%09%09%09int group = 0;%0D%0A%09%09%09for (int t = 0; t %3C nTaxa; t++) %7B%0D%0A%09%09%09%09if (groups%5Bt%5D == -1) %7B%0D%0A%09%09%09%09%09Node aNode = theTree.getExternalNode(t);%0D%0A%09%09%09%09%09while (aNode.getParent().getNodeHeight() %3C= height) %7B%0D%0A%09%09%09%09%09%09aNode = aNode.getParent();%0D%0A%09%09%09%09%09%7D%0D%0A%09%09%09%09%09setNodeToGroup(aNode, group, groups);%0D%0A%09%09%09%09%09group++;%0D%0A%09%09%09%09%7D%0D%0A%09%09%09%7D%0D%0A%09%09%7D%0D%0A%09%09%0D%0A%09%09return groups;%0D%0A%09%7D%0D%0A%09%0D%0A
 %09public void set

add function to cut the tree at a specific height