#pragma once
#include<iostream>
#include<sstream>

#ifdef GRAPHVIZ
#include<graphviz/cgraph.h>
#include<graphviz/gvc.h>
#include<string>
#endif

// Now X should be a tree_node*
#define IS_RED(X) (X->parent_color & 1)
#define PARENT(X) ((tree_node*)(X->parent_color & ~1))
#define SET_RED(X) (X->parent_color = (X->parent_color | 1))
#define SET_BLACK(X) (X->parent_color = (X->parent_color & ~1))
#define SET_PARENT(X, P) (X->parent_color = ((size_t)P | (X->parent_color & 1)))

template<typename T>
class rbtree_set {
	struct tree_node {
		T data;
		tree_node *left = 0, *right = 0;
		size_t parent_color = 0;
	};
	tree_node *root = 0;
public:
	void insert(const T& new_item){
		tree_node *ntn = new tree_node;
		ntn->data = new_item;
		if(!root){
			root = ntn;
			return;
		}
		tree_node *place = root;
		tree_node *last_place;
		while(place){
			last_place = place;
			if(place->data < new_item) {
				place = place->right;
				
			}
			else if(place->data == new_item){
				delete ntn;
				return;
			}
			else {
				place = place->left;
			}
		}
		// Node goes under last_place
		ntn->parent_color = (size_t)last_place; // Maybe type error
		if(last_place->data < new_item){
			last_place->right = ntn;
		} else {
			last_place->left = ntn;
		}
		SET_RED(ntn);
		insert_rebalance(ntn);
	}

	void insert_rebalance(tree_node *current  ){
		tree_node *parent = PARENT(current);
		// Case I3
		if(!parent)
			return;

		// Case I1
		if(!IS_RED(parent))
			return;

		tree_node *uncle;
		tree_node *grandparent = PARENT(parent);
		// Case I2
		if(grandparent){
			uncle = grandparent->left;
			if(uncle == parent)
				uncle = grandparent->right;
			if(uncle && IS_RED(uncle)) {
				SET_BLACK(uncle);
				SET_BLACK(parent);
				SET_RED(grandparent);
				insert_rebalance(grandparent);
				return;
			}
		} else { // Case I4
			SET_BLACK(parent);
			return;
		}
		// What if we're moving the root of the tree?
		// A:  We have a grandparent that we're not moving
		// Case I5 (preparation for I6)
		if(parent->right == current && grandparent->left == parent){
			// If the above is true, we're an inner child
			grandparent->left = current;
			parent->right = current->left;
			if(parent->right)
				SET_PARENT(parent->right, parent);
			current->left = parent;
			SET_PARENT(current, grandparent);
			SET_PARENT(parent, current);
			current = parent;
		}
		else if(parent->left == current && grandparent->right == parent){
			// If the above is true, we're an inner child
			grandparent->right = current;
			parent->left = current->right;
			if(parent->left)
				SET_PARENT(parent->left, parent);
			current->right = parent;
			SET_PARENT(current, grandparent);
			SET_PARENT(parent, current);
			current = parent;
		}

		// Case I6	
		// We need to restore these if I5 happened
		
		// Great Grandparent needs to point to the parent
		parent = PARENT(current);
		grandparent = PARENT(parent);
		tree_node *great_grandparent = PARENT(grandparent);
		if(!great_grandparent)
			root = parent;
		else 
			if(great_grandparent->left == grandparent)
				great_grandparent->left = parent;
			else
				great_grandparent->right = parent;

		if(parent->right == current){
			grandparent->right = parent->left;
			parent->left = grandparent;
			if(grandparent->right)
				SET_PARENT(grandparent->right, grandparent);
		} else if(parent->left == current){
			grandparent->left = parent->right;
			parent->right = grandparent;
			if(grandparent->left)
				SET_PARENT(grandparent->left, grandparent);
		}
		SET_PARENT(parent, PARENT(grandparent));
		SET_PARENT(grandparent, parent);
		SET_RED(grandparent);
		SET_BLACK(parent);
	}
	
	void preorder_traversal_from(auto F, tree_node *place){
		F(place->data);
		if(place->left)
			preorder_traversal_from(F, place->left);
		if(place->right)
			preorder_traversal_from(F, place->right);
	}
	void preorder_traversal(auto F){
		if(root)
			preorder_traversal_from(F, root);
	}

	void inorder_traversal_from(auto F, tree_node *place){
		if(place->left)
			inorder_traversal_from(F, place->left);
		F(place->data);
		if(place->right)
			inorder_traversal_from(F, place->right);
	}
	void inorder_traversal(auto F){
		if(root)	
			inorder_traversal_from(F, root);
	}

	void postorder_traversal_from(auto F, tree_node *place){
		if(place->left)
			postorder_traversal_from(F, place->left);
		if(place->right)
			postorder_traversal_from(F, place->right);
		F(place->data);
	}
	void postorder_traversal(auto F){
		if(root)
			postorder_traversal_from(F, root);
	}
	
	bool contains(const T& target){
		tree_node *place = root;
		while(place){
			if(place->data == target)
				return true;
			if(place->data < target)
				place = place->right;
			else
				place = place->left;
			// this would work too:
			// place = (place->data < target)? place->right:place->left;	
		}
		// If we're here, place was null, so we didn't find the target
		return false;
	}

	void delete_node(const T& target){
		tree_node *place = root;
		tree_node **parent = &root;
		while(place){
			if(place->data == target) 
				break;
			parent = &place;
			if(place->data < target) {
				parent = &(place->right);
				place = place->right;
			}
			else {
				parent = &(place->left);
				place = place->left;
			}
		}
		if(!place)
			return;
		// We know the node is in the tree and place points to it
		if(place->right && place->left){ // Go find a different node to delete
			tree_node *new_place = place->right;
			parent = &place->right;
			while(new_place->left){
				parent = &(new_place->left);
				new_place = new_place->left;
			}
			// Another =.  Again, should we call memcpy or move pointers around?
			place->data = new_place->data;
			place = new_place;
		}
		if(place->right){
			place->data = place->right->data; // Is it ok to call = on our data?
			delete place->right;              // Maybe call memcpy or change pointers around
			place->right = 0;
			SET_BLACK(place);
			return;
		} else if(place->left){
			place->data = place->left->data;
			delete place->left;
			place->left = 0;
			SET_BLACK(place);
			return;
		}
		if(place == root){
			delete root;
			root = 0;
			return;
		}
		if(IS_RED(place)){
			*parent = 0;
			delete place;
			return;
		}
		// Fix the imbalance that's about to happen
		delete_fix(place);
		
		// Now we have the problem case!
		*parent = 0;
		delete place;
	}

	void delete_fix(tree_node *current){
		// Case D1
		if(current == root)
			return;
		tree_node *parent = PARENT(current);
		tree_node *sibling = parent->left == current? parent->right:parent->left;
		// Case D2
		if(!IS_RED(parent) && sibling && !IS_RED(sibling)){
			if(((sibling->left && !IS_RED(sibling->left)) || !sibling->left) &&
					((sibling->right && !IS_RED(sibling->right)) || !sibling->right)){
				SET_RED(sibling);
				delete_fix(parent);
				return;
			}
		}
		// Case D3
		// What points to the parent?
		tree_node **pgp = &root;
		tree_node *grandparent = PARENT(parent);
		if(grandparent){
			if(grandparent->left == parent)
				pgp = &(grandparent->left);
			else
				pgp = &(grandparent->right);
		}

		if(sibling && IS_RED(sibling)) {
			if(parent->left == current){
				parent->right = sibling->left;
				if(parent->right)
					SET_PARENT(parent->right, parent);
				sibling->left = parent;
				SET_PARENT(sibling->left, sibling);
			} else { // We're a right child
				parent->left = sibling->right;
				if(parent->left)
					SET_PARENT(parent->left, parent);
				sibling->right = parent;
				SET_PARENT(sibling->right, sibling);
			}
			*pgp = sibling;
			if(grandparent)
				SET_PARENT(sibling, grandparent);
			SET_BLACK(sibling);
			SET_RED(parent);
			delete_fix(current);
			return;
		}
		// Case D4
		if(IS_RED(parent) && sibling && !IS_RED(sibling)){
			if( (!sibling->left || !IS_RED(sibling->left)) && 
					(!sibling->right || !IS_RED(sibling->right))){
				SET_RED(sibling);
				SET_BLACK(parent);
				return;
			}
		}
		// Case D5 (preparation for D6)
		if(sibling && !IS_RED(sibling)){
			if(current == parent->left){
				if(sibling->left && IS_RED(sibling->left) && (!sibling->right || !IS_RED(sibling->right))){
					tree_node *C = sibling->left;
					sibling->left = C->right;
					if(sibling->left)
						SET_PARENT(sibling->left, sibling);
					C->right = sibling;
					SET_PARENT(C->right, C);
					parent->right = C;
					SET_PARENT(parent->right, parent);
					SET_RED(sibling);
					SET_BLACK(C);
					sibling = C;
				}
			} else {
				if(sibling->right && IS_RED(sibling->right) && (!sibling->left || !IS_RED(sibling->left))){
					tree_node *C = sibling->right;
					sibling->right = C->left;
					if(sibling->right)
						SET_PARENT(sibling->right, sibling);
					C->left = sibling;
					SET_PARENT(C->left, C);
					parent->left = C;
					SET_PARENT(parent->left, parent);
					SET_RED(sibling);
					SET_BLACK(C);
					sibling = C;
				}
			}
		}
		// Case D6
		if(!sibling)
			return; // No need for case D6, we don't have a sibling
		tree_node *D = parent->left == current? sibling->right : sibling->left;
		// Re-figure these in case we did a D3
		pgp = &root;
		grandparent = PARENT(parent);
		if(grandparent){
			if(grandparent->left == parent)
				pgp = &(grandparent->left);
			else
				pgp = &(grandparent->right);
		}

		if(!IS_RED(sibling) && D && IS_RED(D)){
			if(parent->left == current){
				parent->right = sibling->left;
				if(parent->right)
					SET_PARENT(parent->right, parent);
				sibling->left = parent;
				SET_PARENT(sibling->left, sibling);
			} else {
				parent->left = sibling->right;
				if(parent->left)
					SET_PARENT(parent->left, parent);
				sibling->right = parent;
				SET_PARENT(sibling->right, sibling);
			}
			*pgp = sibling;
			if(grandparent)
				SET_PARENT(sibling, grandparent);
			if(IS_RED(parent))
				SET_RED(sibling);
			else
				SET_BLACK(sibling);
			SET_BLACK(parent);
			SET_BLACK(D);
		}
	}

	void deallocate_tree(tree_node *place){
		if(place->left)
			deallocate_tree(place->left);
		if(place->right)
			deallocate_tree(place->right);
		delete place;
	}
	~rbtree_set(){
		if(root)
			deallocate_tree(root);
	}

	/* Two things to look for:
	 * 	Are there two red nodes in a row?
	 * 	Are all the paths the same length counting just the black nodes?
	 */
	int expected_depth;
	bool is_valid_recur(tree_node *place, int depth){
		if(IS_RED(place)){
			if(place->left && IS_RED(place->left)) {
				std::cout << place->data << " and " << place->left->data << " are both red\n";
				return false;
			}
			if(place->right && IS_RED(place->right)){
				std::cout << place->data << " and " << place->right->data << " are both red\n";
				return false;	
			}
		}
		if(!IS_RED(place))
			depth++;
		if(!place->left || !place->right){ // We're kind of like a leaf
			if(expected_depth == -1)
				expected_depth = depth;
			if(depth != expected_depth){
				std::cout << place->data << " is at unexpected depth " << depth << std::endl;
				return false;
			}
		}
		if(place->left)
			if(!is_valid_recur(place->left, depth))
				return false;
		if(place->right)
			if(!is_valid_recur(place->right, depth))
				return false;
		return true;
	}
	bool is_valid(){
		expected_depth = -1;
		if(root)
			return is_valid_recur(root, 0);
		else
			return true;
	}

#ifdef GRAPHVIZ	
	void make_graph_recur(Agraph_t *g, Agnode_t *parent, tree_node *place){
		std::stringstream converter;
		converter << place->data;
		std::string label;
		converter >> label;

		Agnode_t *graphviz_place = agnode(g, (char*)label.c_str(), 1);
		if(IS_RED(place))
			agset(graphviz_place, (char*)"color", (char*)"red");
		if(parent)
			agedge(g, parent, graphviz_place, (char*)"", 1);
		if(place->left)
			make_graph_recur(g, graphviz_place, place->left);
		if(place->right)
			make_graph_recur(g, graphviz_place, place->right);
	}

	void make_image(const char* name){
		if(!root){
			std::cerr << "No tree exists, refusing to make graph\n";
			return;
		}
		Agraph_t *g;
		g = agopen((char*)"G", Agdirected, 0);
		Agsym_t *color_attr = agattr(g, AGNODE, (char*)"color",  (char*)"black");

		make_graph_recur(g, 0, root); 
		
		GVC_t* gvc = gvContext();
		gvLayout(gvc, g, "dot");
		gvRenderFilename(gvc, g, "png", name);
		gvFreeLayout(gvc, g);
		agclose(g);
	}
#endif
};