package ec.gp.syntactic.simplification;

import java.util.Collections;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.Set;

import ec.Individual;
import ec.gp.GPIndividual;
import ec.gp.GPNode;
import ec.gp.GPTree;

/**
 * Rewrites trees in population, such that a duplicated (equivalent) subtrees are represented by only one instance of
 * it. Note that the resulting population can be incompatible with some pipelines/operators/etc. Don't use it unless you
 * are completely sure that it won't damage your experiment!
 * 
 * 
 * @author Tomasz Pawlak
 */
public class SubtreeRewriter {

	private static final int MAX_DEPTH = 18;
	private static final int MAX_HEIGHT = 9;

	private Set<GPNode> visited;
	private HashMap<Wrapper, Wrapper> treeCache;

	//private long replacements;
	//private long reducedNodeCount;

	public void rewrite(final Individual[] individuals) {
		//this.replacements = 0;
		//this.reducedNodeCount = 0;

		System.gc();
		System.out.println(String.format("Before memory usage\t%dMB", (Runtime.getRuntime().totalMemory() - Runtime
				.getRuntime().freeMemory()) / 1024 / 1024));

		this.visited = Collections.newSetFromMap(new IdentityHashMap<GPNode, Boolean>(0x00040000));
		this.treeCache = new HashMap<Wrapper, Wrapper>(0x00040000);

		for (final Individual ind : individuals) {
			final GPIndividual gpInd = (GPIndividual) ind;
			for (final GPTree tree : gpInd.trees) {
				this.rewriteSubtree(tree.child, 1);
			}
		}

		this.visited = null;
		this.treeCache = null;

		System.gc();
		System.out.println(String.format("After memory usage\t%dMB", (Runtime.getRuntime().totalMemory() - Runtime
				.getRuntime().freeMemory()) / 1024 / 1024));

		//System.out.println(String.format("Number of replacements: %d, reducedNodeCount: %d", replacements,
		//		reducedNodeCount));
	}

	private int rewriteSubtree(GPNode node, int myDepth) {
		int myHeight = 1;
		int childHeight;

		if (myDepth > MAX_DEPTH)
			return 0x10000000; // just a big number

		for (int ch = 0; ch < node.children.length; ++ch) {
			childHeight = this.rewriteSubtree(node.children[ch], myDepth + 1);

			if (childHeight >= myHeight)
				myHeight = childHeight + 1;

			if (childHeight > MAX_HEIGHT)
				continue; // do not try to rewrite such a big tree

			if (this.visited.contains(node.children[ch]))
				continue; //already visited this node

			Wrapper wrapped = new Wrapper(node.children[ch]);
			Wrapper inCache = this.treeCache.get(wrapped);
			if (inCache == null) {
				// insert into cache
				this.treeCache.put(wrapped, wrapped);
				this.visited.add(node.children[ch]);
			} else if (wrapped.getNode() != inCache.getNode()) {
				assert this.visited.contains(inCache.getNode());
				assert !this.visited.contains(node.children[ch]);

				// replace subtree with a cached subtree
				node.children[ch] = inCache.getNode();

				// statistics
				//++this.replacements;
				//this.reducedNodeCount += inCache.getNode().numNodes(GPNode.NODESEARCH_ALL);
			}
		}

		return myHeight;
	}
}

final class Wrapper {
	private final GPNode node;

	public Wrapper(GPNode node) {
		this.node = node;
	}

	public final GPNode getNode() {
		return this.node;
	}

	@Override
	public final boolean equals(Object obj) {
		assert obj instanceof Wrapper;

		return this.node.rootedTreeEquals(((Wrapper) obj).getNode());
	}

	public final int hashCode() {
		return this.node.rootedTreeHashCode();
	}
}
