/*******************************************************************************
 * Copyright (c) 2007, 2010 Association for Decentralized Information Management
 * in Industry THTH ry.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     VTT Technical Research Centre of Finland - initial API and implementation
 *******************************************************************************/
package org.simantics.utils.datastructures.persistent;

    import gnu.trove.procedure.TObjectObjectProcedure;

public class ImmutableMap<T extends Comparable<T>, V> {
	
	@SuppressWarnings({ "rawtypes" })
	static final ImmutableMap EMPTY = new EmptyImmutableSet();
	
	private static class EmptyImmutableSet<T extends Comparable<T>, V> extends ImmutableMap<T, V> {
		public EmptyImmutableSet() {
			isBlack = true;
		}
		
		protected ImmutableMap<T, V> putRec(T key, V value) {
			return new ImmutableMap<T, V>(key, value);
		}
		
		protected ImmutableMap<T, V> removeRec(T obj) {
			return null;
		}
		
		public boolean contains(T obj) {
			return false;
		}
	
		public V get(T key) {
			return null;
		}
		
		@Override
		public boolean forEach(TObjectObjectProcedure<T, V> arg0) {
			return true;
		}
	}
	
	ImmutableMap<T, V> left;
	T key;
	V value;
	ImmutableMap<T, V> right;
	boolean isBlack;
	
	protected ImmutableMap(
			ImmutableMap<T, V> left,
			T key,
			V value,
			ImmutableMap<T, V> right,			
			boolean isBlack) {
		this.left = left;
		this.right = right;
		this.key = key;
		this.value = value;
		this.isBlack = isBlack;
	}

	@SuppressWarnings("unchecked")
	public ImmutableMap(T key, V value) {
		this(EMPTY, key, value, EMPTY, false);
	}		

	private ImmutableMap() {	
	}
	
	public boolean contains(T obj) {
		int cmp = obj.compareTo(key);
		if(cmp < 0)
			return left.contains(obj);
		else if(cmp > 0)
			return right.contains(obj);
		else
			return true;
	}	
	
	public V get(T key) {
		int cmp = key.compareTo(key);
		if(cmp < 0)
			return left.get(key);
		else if(cmp > 0)
			return right.get(key);
		else
			return value;
	}	
	
	protected ImmutableMap<T, V> putRec(T obj, V value) {
		int cmp = obj.compareTo(key);
		if(cmp < 0) {
			ImmutableMap<T, V> newLeft = left.putRec(obj, value);
			if(newLeft == left)
				return this;			
			if(isBlack)
				return balance(newLeft, key, value, right);
			else
				return red(newLeft, key, value, right);
		}
		else if(cmp > 0) {
			ImmutableMap<T, V> newRight = right.putRec(obj, value);
			if(newRight == right)
				return this;
			if(isBlack)
				return balance(left, key, value, newRight);
			else
				return red(left, key, value, newRight);
		}
		else
			return this;
	}
	
	/**
	 * Assumes this is a black nonempty node. 
	 * 
	 * Removes obj from the tree. The returned tree has always
	 * one black node less in every branch (even if nothing is
	 * removed).
	 *   
	 * @param obj
	 * @return
	 */
	protected ImmutableMap<T, V> removeRec(T obj) {		
		int cmp = obj.compareTo(key);
		if(cmp < 0) {
			ImmutableMap<T, V> newLeft = left.removeRec(obj);
			if(newLeft == null)
				return null;
			else if(left.isBlack)
				return balleft(newLeft, key, value, right);
			else
				return red(newLeft, key, value, right);
		}
		else if(cmp > 0) {
			ImmutableMap<T, V> newRight = right.removeRec(obj);
			if(newRight == null)
				return null;
			else if(right.isBlack)
				return balright(left, key, value, newRight);
			else
				return red(left, key, value, newRight);			
		}
		else
			return append(left, right);
	}
	
	/**
	 * Assumes a and b have the same black depth and keys in a are strictly less
	 * than keys in b. Creates a new tree with the same black depth as a and b. 
	 * 
	 * @param <T>
	 * @param a
	 * @param b
	 * @return
	 */
	protected static <T extends Comparable<T>, V> ImmutableMap<T, V> append(ImmutableMap<T, V> a, ImmutableMap<T, V> b) {
		if(a==EMPTY)
			return b;
		if(b==EMPTY)
			return a;
		if(a.isBlack) {
			if(b.isBlack) {
				ImmutableMap<T, V> middle = append(a.right, b.left);
				if(middle.isBlack)
					return balleft(a.left, a.key, a.value, black(middle, b.key, b.value, b.right));
				else
					return red(black(a.left, a.key, a.value, middle.left), middle.key, middle.value, black(middle.right, b.key, b.value, b.right));
			}
			else
				return red(append(a, b.left), b.key, b.value, b.right);
		}
		else {
			if(b.isBlack)
				return red(a.left, a.key, a.value, append(a.right, b));
			else {
				ImmutableMap<T, V> middle = append(a.right, b.left);
				if(middle.isBlack)
					return red(a.left, a.key, a.value, red(middle, b.key, b.value, b.right));
				else
					return red(red(a.left, a.key, a.value, middle.left), middle.key, middle.value, red(middle.right, b.key, b.value, b.right));
			}
		}
	}
	
	public T getFirst() {
		if(left == EMPTY)
			return key;
		else
			return left.getFirst();
	}	
	
	static private <T extends Comparable<T>, V> ImmutableMap<T, V> balance(
			ImmutableMap<T, V> left,
			T key,
			V value,
			ImmutableMap<T, V> right) {
		if(!left.isBlack) {
			if(!left.left.isBlack) 
				return red(
					toBlack(left.left),
					left.key,
					left.value,
					black(left.right, key, value, right)
				);
			else if(!left.right.isBlack)
				return red(
					black(left.left, left.key, left.value, left.right.left),
					left.right.key, left.right.value,
					black(left.right.right, key, value, right)					
				);				
		}
		if(!right.isBlack) {
			if(!right.left.isBlack)
				return red(
					black(left, key, value, right.left.left),
					right.left.key, right.left.value,
					black(right.left.right, right.key, right.value, right.right)
				);
			else if(!right.right.isBlack)
				return red(
					black(left, key, value, right.left),
					right.key,
					right.value,
					toBlack(right.right)
				);		
		}
		return black(left, key, value, right);
	}
	
	static private <T extends Comparable<T>, V> ImmutableMap<T, V> black(
			ImmutableMap<T, V> left,
			T key,
			V value,
			ImmutableMap<T, V> right) {
		return new ImmutableMap<T, V>(left, key, value, right, true);
	}
	
	static private <T extends Comparable<T>, V> ImmutableMap<T, V> red(
			ImmutableMap<T, V> left,
			T key,
			V value,
			ImmutableMap<T, V> right) {
		return new ImmutableMap<T, V>(left, key, value, right, false);
	}
	
	static private <T extends Comparable<T>, V> ImmutableMap<T, V> toBlack(ImmutableMap<T, V> tree) {
		if(tree.isBlack)
			return tree;
		else
			return black(tree.left, tree.key, tree.value, tree.right);
	}
	
	static private <T extends Comparable<T>, V> ImmutableMap<T, V> toRed(ImmutableMap<T, V> tree) {
		if(tree.isBlack)
			return red(tree.left, tree.key, tree.value, tree.right);
		else
			return tree;
	}
			
	
	static private <T extends Comparable<T>, V> ImmutableMap<T, V> balleft(
			ImmutableMap<T, V> left,
			T key,
			V value,
			ImmutableMap<T, V> right) {
		if(left.isBlack) {
			if(right.isBlack)
				return balance(left, key, value, toRed(right));
			else
				return red(black(left, key, value, right.left.left), right.left.key, right.left.value,
					balance(right.left.right, right.key, right.value, toRed(right.right)));
		}
		else
			return red(toBlack(left), key, value, right);
	}
	
	static private <T extends Comparable<T>, V> ImmutableMap<T, V> balright(
			ImmutableMap<T, V> left,
			T key,
			V value,
			ImmutableMap<T, V> right) {
		if(right.isBlack) {
			if(left.isBlack)
				return balance(toRed(left), key, value, right);
			else
				return red(balance(toRed(left.left), left.key, left.value, left.right.left), left.right.key, left.right.value,
					black(left.right.right, key, value, right));
		}
		else
			return red(left, key, value, toBlack(right));
	}

	public ImmutableMap<T, V> put(T key, V value) {
		ImmutableMap<T, V> tree = putRec(key, value);
		tree.isBlack = true;
		return tree;
	}

	public ImmutableMap<T, V> remove(T obj) {
		ImmutableMap<T, V> tree = removeRec(obj);
		if(tree == null)
			return this;
		if(tree.isBlack)
			return tree;
		else
			return black(tree.left, tree.key, tree.value, tree.right);
	}
	
	public boolean forEach(TObjectObjectProcedure<T, V> proc) {
		if(left.forEach(proc))
			if(proc.execute(key, value))
				return right.forEach(proc);
		return false;
	}	
			
}