/**
 * Persistent AA Tree tree.
 * 
 * @author Andreas van Cranenburgh
 * 
 * Uses node-copying (as opposed to path-copying, so has O(1) amortized space
 * complexity for each update operation instead of O(log n)).
 * 
 * Each node is augmented with a mutation field, containing either an updated
 * left or right pointer. This way, nodes along the path of a node to be 
 * updated can be re-used many times.
 * 
 * For more details, see: Driscoll et al., 1989, "Making Data Structures
 * Persistent", journal of computer and system sciences, vol. 38, no. 1,
 * february 1989. http://www.cs.cmu.edu/~sleator/papers/Persistence.htm
 * 
 * Reference for AA trees: Andersson, 1993, "Balanced Search Trees Made Simple"
 * In Proc. Workshop on Algorithms and Data Structures, pages 60--71. Springer
 * Verlag. http://user.it.uu.se/~arnea/abs/simp.html
 */

public class NodeCopyingAATree<K extends Comparable<? super K>, V> implements
		BST<K, V> {
	/** Constants to signify whether a mutation is for the left or right
	  * child. */
	final boolean LEFT = java.lang.Boolean.TRUE;

	final boolean RIGHT = java.lang.Boolean.FALSE;

	/** Array to contain root node for each version of the tree */
	private java.util.ArrayList<AANode<K, V>> versions;

	/** Array to contain mutations, in order for rollback() to be able to
	  * erase them. */
	private java.util.ArrayList<AANode<K, V>> mutations;

	/** Reference to the current (i.e., writable) root of the tree. */
	private AANode<K, V> root;

	/** Version number of the tree. */
	protected int livever = 0;

	/** Default constructor */
	public NodeCopyingAATree() {
		root = null;
		versions = new java.util.ArrayList<AANode<K, V>>();
		mutations = new java.util.ArrayList<AANode<K, V>>();
	}

	/** find key and return corresponding value
	 * 
	 * @param key the key to find
	 * @return value associated with key, or null if not found.
	 */
	public V search(K key) {
		if (key == null)
			throw new NullPointerException("tried to search with null key");
		else
			return search(root, key);
	}

	/** find key and return corresponding value from a specific version
	 * 
	 * @param key the key to find
	 * @param version the version number of the tree.
	 * @return value associated with key, or null if not found / non-existing
	 *         version.
	 */
	public V search(K key, int version) {
		if (key == null)
			throw new NullPointerException("tried to search with null key");
		if (version > livever)
			return null;

		/* temporarily set current version to requested version, and restore it
		 * after performing search. */
		int temp = livever;
		livever = version;
		V value = search(root, key);
		livever = temp;
		return value;
	}

	/** find key in a specific subtree, return value pass a root node from an
 	  * earlier insert or delete operation to search in that version of the 
 	  * tree.
	  * 
	  * @param node the (sub)tree where lookup starts
	  * @param key the key to find
	  * @return the value associated with key, in the specified tree
	  */
	public V search(AANode<K, V> n, K key) {
		/* key not found */
		if (n == null)
			return null;

		int comp = key.compareTo(n.key);
		if (comp < 0)
			return search(n.getLeft(livever), key);
		else if (comp > 0)
			return search(n.getRight(livever), key);
		else
			/* key found */
			return n.value;
	}

	/**
	 * associate key with value
	 * 
	 * @param key the key to insert
	 * @param value the value to associate with key
	 * @return the version spawned by this operation.
	 */
	public int insert(K key, V value) {
		versions.add(root);
		livever++;
		if (key == null)
			throw new NullPointerException("tried to insert with null key");

		root = insert(root, key, value);
		return livever;
	}

	/** insert key in a specific subtree */
	private AANode<K, V> insert(AANode<K, V> node, K key, V value) {
		/* place new node here, end recursion */
		if (node == null)
			return new AANode<K, V>(key, value);

		/* recurse until we reach the desired node */
		AANode<K, V> next, result;
		int comp = key.compareTo(node.key);
		if (comp < 0) {
			/* traverse newest version of the tree */
			next = node.getLeft(livever);
			result = insert(next, key, value);

			/* update reference in case a new node was spawned */
			if (result != next)
				node = modLeft(node, result);

		} else if (comp > 0) {
			/* traverse newest version of the tree */
			next = node.getRight(livever);
			result = insert(next, key, value);

			/* update reference in case a new node was spawned */
			if (result != next)
				node = modRight(node, result);

		} else { /* old key found, add new value */
			node = new AANode<K, V>(node);
			node.key = key;
			node.value = value;
		}

		/* before returning, first skew and then split */
		return split(skew(node));
	}

	/**
	 * find key and remove corresponding element
	 * 
	 * @param key the key to delete
	 * @return the version spawned by this operation.
	 */
	public int delete(K key) {
		versions.add(root);
		livever++;
		if (key == null)
			throw new NullPointerException("tried to delete with null key");
		root = delete(root, key);
		return livever;
	}

	/**
	 * find key and remove corresponding element from specified (sub)tree.
	 */
	private AANode<K, V> delete(AANode<K, V> node, K key) {
		/* key not found, end recursion */
		if (node == null)
			return null;

		/* recurse until we reach the desired node */
		AANode<K, V> next, result;
		int comp = key.compareTo(node.key);
		if (comp < 0) {
			/* traverse newest version of the tree */
			next = node.getLeft(livever);

			/* recurse */
			result = delete(next, key);

			/* update reference in case a new node was spawned */
			if (result != next)
				node = modLeft(node, result);
			return node;

		} else if (comp > 0) {
			/* traverse newest version of the tree */
			next = node.getRight(livever);

			/* recurse */
			result = delete(next, key);

			/* update reference in case a new node was spawned */
			if (result != next)
				node = modRight(node, result);
			return node;
		}

		/*
		 * recursion reached the node to be deleted
		 * (since comp is neither < 0, nor > 0)
		 */
		AANode<K, V> left = node.getLeft(livever),
			right = node.getRight(livever);

		/* simple cases, 0 or 1 child: */
		if (left == null && right == null)
			return null;
		else if (left == null)
			result = right;
		else if (right == null)
			result = left;
		/* 2 children, swap node with in-order predecessor and delete it. to
		 * simplify matters the path to the predecessor is walked twice, first
		 * looking it up and then copying its path minus the node to be
		 * deleted. */
		else {
			/* find in-order predecessor */
			AANode<K, V> pred = predNode(left);
			/* delete it from this subtree */
			AANode<K, V> temp = delete(left, pred.key);
			/* return predecessor with children of old node */
			result = new AANode<K, V>(temp, pred.key, pred.value, right);
		}

		/* slightly cumbersome version of AA tree re-balancing, since
		 * we have to avoid destructive assignments and side-effects.
		 * */
		result = decrease_level(result);
		result = skew(result);
		/* skew(node.right) */
		result = modRight(result, skew(result.getRight(livever)));
		/* skew(node.right.right) */
		if (result.getRight(livever) != null)
			result = modRight(result, 
				modRight(result.getRight(livever), 
				skew(result.getRight(livever).
					getRight(livever))));

		result = split(result);

		result = modRight(result, split(result.
			getRight(livever)));
		
		return result;
	}

	/** return in-order predecessor of specified node */
	private AANode<K, V> predNode(AANode<K, V> node) {
		AANode<K, V> next;
		/* traverse newest version of the tree */
		if (node.ver != 0 && node.ver <= livever && node.lr == RIGHT)
			next = node.mutation;
		else
			next = node.right;

		if (next == null)
			return node;
		else

			return predNode(next);
	}
	
	private AANode<K,V> decrease_level(AANode<K,V> node) {
		node = new AANode<K,V>(node);
		int correctlevel = Math.max(node.left == null ? 0 :
			node.left.level, node.right == null ||
			node.right.getRight(livever) == null ? 0 :
			node.right.getRight(livever).level) + 1; 
		if (correctlevel < node.level) {
			node.level = correctlevel;
			if (correctlevel < (node.right == null ? 0 : 
				node.right.level))
				node.right.level = correctlevel;
		}
		return node;
	}

	private AANode<K, V> skew(AANode<K, V> node) {
		if (node == null)
			return null;
		if (node.getLeft(livever) != null &&
			node.getLeft(livever).level == node.level) {
			AANode<K, V> temp;
			/* rotate right */
			if (node.ver == 0) {
				temp = node;
				node = new AANode<K, V>(node.getLeft(livever));
				temp.mutation = node.getRight(livever);
				temp.lr = LEFT;
				temp.ver = livever;
			} else {
				temp = new AANode<K, V>(node);
				if (node.getLeft(livever) == null)
					node = temp.left = null;
				else {
					node = new AANode<K, V>(node.getLeft(livever));
					temp.left = node.getRight(livever);
				}
			}
			node.right = temp;
		}
		return node;
	}

	private AANode<K, V> split(AANode<K, V> node) {
		if (node == null)
			return null;
		/* all bow to NullPointerExceptions! */
		if (node.getRight(livever) != null && 
			node.getRight(livever).getRight(livever) != null && 
			node.getRight(livever).getRight(livever).level == node.level) {
			AANode<K, V> temp;
			/* rotate left */
			if (node.ver == 0) {
				temp = node;
				node = new AANode<K, V>(node.getRight(livever));
				temp.mutation = node.getLeft(livever);
				temp.lr = RIGHT;
				temp.ver = livever;
			} else {
				temp = new AANode<K, V>(node);
				node = new AANode<K, V>(node.getRight(livever));
				temp.right = node.getLeft(livever);
			}
			node.left = temp;
			node.level += 1;
		}
		return node;
	}

	/** undo last mutation. undoing itself cannot be undone, this operation is
	 * destructive, since the versions are stored in a list which can not keep
	 * track of multiple branches. */
	public void rollback() {
		/* get last version from stack and make it current */
		root = versions.remove(versions.size() - 1);

		/* pop mutations of current version and erase them */
		while (mutations.get(mutations.size() - 1).ver == livever)
			mutations.remove(mutations.size() - 1).ver = 0;

		livever--;
	}

	/** return an in-order traversal of the tree */
	public String toString() {
		if (root == null)
			return "empty tree";
		else
			return inOrder(root);
	}

	/** recursive in-order helper function */
	private String inOrder(AANode<K, V> node) {
		if (node == null)
			return "";

		AANode<K, V> left = node.left, right = node.right;

		/* traverse newest version of tree */
		if (node.ver != 0 && node.ver <= livever)
			if (node.lr == LEFT)
				left = node.mutation;
			else
				right = node.mutation;

		return inOrder(left) + " " + node.key + "=" + node.value
				+ inOrder(right);
	}

	/**
	 * Modify left child for a given node. This will either store the new child
	 * in the mutation field, or spawn a new node.
	 */
	private AANode<K,V> modLeft(AANode<K,V> node, AANode<K,V> left) {
		AANode<K,V> result = node;
		if (node.ver != 0) {
			result = new AANode<K,V>(node);
			result.ver = 0;
			result.left = left;
		} else {
			result.mutation = left;
			result.lr = LEFT;
			result.ver = livever;
			mutations.add(result); 
		}

		return result;
	}
	 

	/** Modify right child for a given node. This will either store the new
	  * child in the mutation field, or spawn a new node. */
	private AANode<K, V> modRight(AANode<K, V> node, AANode<K, V> right) {
		AANode<K, V> result = node;
		if (node.ver != 0) {
			result = new AANode<K, V>(node);
			result.ver = 0;
			result.right = right;
		} else {
			result.mutation = right;
			result.lr = RIGHT;
			result.ver = livever;
			mutations.add(result);
		}

		return result;
	}

	public static void main(String[] args) {
		NodeCopyingAATree t = new NodeCopyingAATree<Integer, Integer>();
		t.insert(5, 5);
		System.out.println(t);
		t.insert(4, 4);
		System.out.println(t);
		t.insert(8, 8);
		System.out.println(t);
		t.insert(1, 1);
		System.out.println(t);
		t.delete(5);
		System.out.println(t);
		t.delete(4);
		System.out.println(t);
		t.delete(8);
		System.out.println(t);
		t.delete(1);
		System.out.println(t);
	}
}

/** AA tree node with extra level pointer, as well as extra pointer for
 * persistence. */
class AANode<K extends Comparable<? super K>, V> {
	protected K key;

	protected V value;

	protected AANode<K, V> left, right, mutation;

	/* version stamp of mutation (0 means no mutation) */
	protected int ver = 0;

	/* whether the mutation is for the left or right child */
	protected boolean lr;

	/* level of node */
	protected int level = 1;
	
	/** Constants to signify whether a mutation is for the left or right child.
	 */
	final boolean LEFT = java.lang.Boolean.TRUE;

	final boolean RIGHT = java.lang.Boolean.FALSE;

	public AANode(AANode<K, V> left, K key, V value, AANode<K, V> right) {
		this(key, value);
		this.left = left;
		this.right = right;
	}

	public AANode(K key, V value) {
		this.key = key;
		this.value = value;
	}

	public AANode(AANode<K, V> node) {
		this(node.left, node.key, node.value, node.right);
		if (node.ver != 0) {
			if (node.lr == LEFT)
				this.left = node.mutation;
			else
				this.right = node.mutation;
		}
		this.level = node.level;
	}
	
	/** Get left child, or updated left child if applicable. */
	public AANode<K, V> getLeft(int livever) {
		if (this.ver != 0 && this.ver <= livever && this.lr == LEFT)
			return this.mutation;
		else
			return this.left;
	}

	/** Get right child, or updated left child if applicable. */
	public AANode<K, V> getRight(int livever) {
		if (this.ver != 0 && this.ver <= livever && this.lr == RIGHT)
			return this.mutation;
		else
			return this.right;
	}
}