package cs211.hw8;

import java.util.Set;
import java.util.TreeSet;

/**
 * This class provides a binary search tree implementation of the Map interface.
 * 
 * @author C. Andrews
 *
 * @param <K> the key type of the Map
 * @param <V> the value type associated with the key
 */
public class BSTMap<K extends Comparable<K>, V> implements Map<K, V> {
	BTNode _root =null;
	Set<K> _keys;
	
	/**
	 * Create a new empty BSTMap.
	 */
	public BSTMap(){
		_keys = new TreeSet<K>();
	}
	

	
	
	/**
	 * This function returns true if the Map contains the provided {@code key}.
	 * 
	 * @param key the key value to look for in the map
	 * @return a boolean value indicating if the key is present
	 */
	@Override
	public boolean containsKey(K key) {
		return get(key) != null;
	}

	/**
	 * Perform a lookup in the map based on the input key.
	 * 
	 * @param key the key associated with the value to be returned
	 */
	@Override
	public V get(K key) {
		BTNode tmp = find(key);
		
		if(tmp == null)
			return null;
		else
			return tmp.value;
	}

	
	/**
	 * Return the set of all keys present in the map.
	 * 
	 * @return the {@code Set} of all keys in the map
	 */
	@Override
	public Set<K> keySet() {
		return _keys;
	}

	/**
	 * Add an element to the map.
	 * 
	 * If the key is already present in the map, the value is replaced with {@code value}.
	 * 
	 * @param key the key to associate the value with
	 * @param value the value being added to the map
	 */
	@Override
	public void put(K key, V value) {
		BTNode tmp = new BTNode(key, value);
		BTNode current = _root;
		BTNode parent = null;
		_keys.add(key);
		
		while (current != null){
			parent = current;
			if (key.compareTo(current.key) == 0){ // they are the same, just update the value
				current.value = value;
				return;
			}else if (key.compareTo(current.key) < 0){ // the key is smaller, go left
				current = current.left;
			}else{ // the key is larger, go right
				current = current.right;
			}
		}
		
		tmp.parent = parent;
		
		if (parent == null){
			_root = tmp;
		}else{
			if (key.compareTo(parent.key) < 0){
				parent.left = tmp;
			}else{
				parent.right = tmp;
			}
		}

	}
	
	/**
	 * Remove the value associated with {@code key} from the map.
	 * 
	 * Removes the key and value pair and returns the value.
	 * 
	 * @param key the key for the value we are trying to remove
	 * @return the value associated with the key or null if the key is not present in the map
	 */
	@Override
	public V remove(K key) {
		_keys.remove(key);
		BTNode toRemove = find(key);
		V value;
		if (toRemove == null){
			return null;
		}
		value = toRemove.value;
		
		if (toRemove.left == null){
			transplant(toRemove, toRemove.right);
		} else if (toRemove.right == null){
			transplant(toRemove, toRemove.left);
		}else{
			BTNode s = successor(toRemove);
			if (s != toRemove.right){
				transplant(s,s.right);
				s.right = toRemove.right;
				s.right.parent = s;
			}
			transplant(toRemove, s);
			s.left =toRemove.left;
			s.left.parent = s;
			
		}
		
		
		return value;
	}

	
	/**
	 * This private method finds nodes in the tree based on the key value.
	 * 
	 * If the key exists in the map, this returns the node associated with it (not just the value). If the 
	 * key is not present, this returns null. This is used to implement get and remove.
	 * 
	 * @param key the key that we are looking for
	 * @return the node containing the key or null if the key is not present
	 */
	private BTNode find(K key){
		BTNode current = _root;
		
		while (current != null && key.compareTo(current.key) != 0){
			if (key.compareTo(current.key) < 0){
				current = current.left;
			}else{
				current = current.right;
			}
		}
		
		if (current != null)
			return current;
		else
			return null;
	}
	
	
	/**
	 * A private helper that finds the minimum value in the tree rooted at t.
	 * 
	 * This will fail if t is null since that would be an invalid operation.
	 * 
	 * @param t the subtree we are exploring
	 * @return the minimum value in t
	 */
	private BTNode minimum(BTNode t){
		while (t.left != null){
			t = t.left;
		}
		return t;
	}
	
	/**
	 * A private helper that finds the immediate successor of a node in the tree.
	 * 
	 * @param t the node we are trying to find the successor of
	 * @return the immediate successor of the node or null if there isn't one
	 */
	private BTNode successor(BTNode t){
		if (t.right != null){
			return minimum(t.right);
		}
		
		BTNode parent = t.parent;
		
		while (parent != null && parent.right == t){
			t = parent;
			parent = parent.parent;
		}
		return parent;
	}


	/**
	 * This function replaces a node in the tree with a different node.
	 * 
	 * This function is used primarily for removals. Note that this only handles the connection to the parent. 
	 * Handling the children should be handled elsewhere.
	 * 
	 * @param oldNode the node being replaced
	 * @param newNode the node replacing the {@code oldNode}
	 */
	private void transplant(BTNode oldNode, BTNode newNode){
		if (oldNode.parent == null){
			_root = newNode;
		}else if (oldNode == oldNode.parent.left){
			oldNode.parent.left = newNode;
		}else{
			oldNode.parent.right = newNode;
		}
		
		if (newNode != null){
			newNode.parent = oldNode.parent;
		}
	}
	


	
	/**
	 * This is a simple node class that provides fields for a key and value, as well as references to 
	 * the parent and the two children.
	 */
	private class BTNode{
		private BTNode left, right, parent;
		
		private K key;
		private V value;
		
		private BTNode(K key, V value){
			this.key = key;
			this.value =value;
		}
		
	}
}