#ifndef CLASS_TREE_H
#define CLASS_TREE_H

// Remove this if you want to use the tree without graphviz support
#define GRAPHVIZ

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

#include<iostream>
using namespace std;

template<class T>
struct node {
	bool isred;
	T data;
	node *left, *right, *parent;
};

template<class T>
class Tree {
	node<T> *root = 0;
#ifdef GRAPHVIZ
	void make_image_recur(Agraph_t*, Agnode_t*, node<T>*);
#endif
	node<T>* find_uncle(node<T>*);
	void rebalance(node<T>*);
public:
	bool find_reference(T data, T **found_place);
	void insert(T data);
	T& search(T data);

#ifdef GRAPHVIZ
	void make_image(string);
#endif
	void deallocate_tree(node<T>* place);
	~Tree(){	deallocate_tree(root);	}
	void print_inorder(node<T>* place = 0);
	void remove(T data);
	void rotate_left(node<T> **place = 0);
	void rotate_right(node<T> **place = 0);
	void improve_balance(node<T> **place = 0);
	int depth_from_place(node<T> *place);
};

template<class T>
void Tree<T>::deallocate_tree(node<T>* place){
	// This would be where code for preorder traversal goes
	if(place->left)
		deallocate_tree(place->left);
	// This place for inorder
	if(place->right)
		deallocate_tree(place->right);
	// This place down here for postorder traversal
	delete place;  // So we do deletion postorder
}

template<class T>
bool Tree<T>::find_reference(T data, T **found_place){
	node<T> *place = root;
	while(place){
		if(place->data == data){
			*found_place = &(place->data);
			return true;
		}
		if(place->data > data)
			place = place->left;
		if(place && place->data < data)
			place = place->right;
	}
	return false;
}

template<class T>
T& Tree<T>::search(T data){
	node<T> *place = root;
	while(place){
		if(place->data == data)
			return place->data;
		if(place->data > data)
			place = place->left;
		if(place->data < data)
			place = place->right;
	}
	cout << "Oh no!  We didn't find it!\n"; // TODO:  Handle this somehow
	return data;
}

template<class T>
void Tree<T>::insert(T data){
	node<T> *new_node = new node<T>;
	new_node->data = data;
	new_node->left = new_node->right = new_node->parent = 0;
	new_node->isred = 1;
	if(!root){
		root = new_node;
		new_node->isred = 0;
		return;
	}
	node<T> *place = root;
	while(place){
		new_node->parent = place;
		if(data < place->data)
			place = place->left;
		else if(data > place->data)
			place = place->right;
		else 
			cout << "Help!  A duplicate!\n";
	}
	if(new_node->parent->data < data)
		new_node->parent->right = new_node;
	else
		new_node->parent->left = new_node;
	rebalance(new_node);
}

template<class T>
node<T>* Tree<T>::find_uncle(node<T> *startpoint){
	if(!startpoint->parent)
		return 0;
	node<T> *grandparent = startpoint->parent->parent;
	if(!grandparent)
		return 0;
	if(grandparent->left == startpoint->parent)
		return grandparent->right;
	else
		return grandparent->left;
}

template<class T>
void Tree<T>::rebalance(node<T> *new_node){
	// Balancing section
	// Case 3:  We're the root (Yes, I realize this is out of order)
	if(!new_node->parent)
		return;
	
	// Case 1:  The parent was not painted red
	if(new_node->parent && !new_node->parent->isred) 
		return;

	// Case 2:  Parent and uncle are red
	node<T> *uncle = find_uncle(new_node);
	node<T> *grandparent = new_node->parent->parent;
	if(grandparent && uncle && uncle->isred){
		new_node->parent->isred = false;
		uncle->isred = false;
		new_node->parent->parent->isred = true;
		rebalance(grandparent);
		return;
	}

	// Case 3:  We're the root of the tree
	
	// Case 4:  Parent is red, and is also the root of the tree
	if(!new_node->parent->parent){
		new_node->parent->isred = false;
		return;
	}

	// Case 5 and 6:  Parent is red, uncle is black, new_node is an inner grandchild
	if(!uncle || !uncle->isred){
		if	((new_node->parent == grandparent->left && new_node == new_node->parent->right) || 
			 (new_node->parent == grandparent->right && new_node == new_node->parent->left)){
			cout << "Now we do case 5!\n";
			make_image("pre_6");
			bool go_left = new_node == new_node->parent->right;
			node<T> **down_ptr = &grandparent->left;
			if(*down_ptr != new_node->parent)
				down_ptr = &grandparent->right;
			new_node = new_node->parent; // Moved this two lines later
			if(go_left)
				rotate_left(down_ptr);
			else
				rotate_right(down_ptr);
//			new_node->parent = go_left? grandparent->left : grandparent->right;
			grandparent = new_node->parent->parent;
			make_image("post_5");
		}
		
		node<T> **down_ptr;
		if(grandparent->parent)
			down_ptr = (grandparent->parent->left == grandparent)? &(grandparent->parent->left) : &(grandparent->parent->right);
		else
			down_ptr = &root;
		if(new_node->parent->left == new_node)
			rotate_right(down_ptr); // Use the pointer to the grandparent, what a pain!
		else
			rotate_left(down_ptr);
		
		if(new_node->parent)	
			new_node->parent->isred = false;
		grandparent->isred = true;
		return;
	}
}

#ifdef GRAPHVIZ
/* To build this program remember to add -lcgraph and -lgvc to your compile line
	If there's an error about missing header files, maybe libgraphviz-dev isn't installed:
	sudo apt install libgraphviz-dev
	*/
template<class T>
void Tree<T>::make_image_recur(Agraph_t *g, Agnode_t *parent, node<T> *place){
	Agnode_t *us = agnode(g, (char*)to_string(place->data).c_str(), 1); // problem later because a T isn't a char*
	if(place->isred)
		agset(us, (char*)"color", (char*)"red");
	if(parent)
		Agedge_t *e = agedge(g, parent, us, (char*)"", 1);
	if(place->left)
		make_image_recur(g, us, place->left);
	if(place->right)
		make_image_recur(g, us, place->right);
}

template<class T>
void Tree<T>::make_image(string name){
	Agraph_t *g;
	g = agopen((char*)"G", Agdirected, 0);
	Agsym_t *color_attr = agattr(g, AGNODE, (char*)"color", (char*)"black");

	// Go through the tree, and add all our nodes!
	// Add our own node
	// If there is a left subtree, add the left subtree
	// If there is a right subtree, add the right subtree
	make_image_recur(g, 0, root);

	GVC_t* gvc = gvContext();
	gvLayout(gvc, g, "dot");
	agwrite(g, stdout);	;
	gvRenderFilename(gvc, g, "png", (name + ".png").c_str());
	gvFreeLayout(gvc, g);
	agclose(g);
}

#endif
template<class T>
void Tree<T>::remove(T data){
	node<T> *place = root;
	node<T> *parent = 0;
	while(place && place->data != data){
		parent = place;
		if(place->data > data)
			place = place->left;
		else
			place = place->right;
	}
	if(!place)
		return;
	// Find the predecessor
	if(place->left){
		node<T> *predecessor = place->left;
		while(predecessor->right){
			predecessor = predecessor->right;
		}
		T hold_data = predecessor->data;
		remove(predecessor->data);
		place->data = hold_data; // Will this cause an extra destructor call?
		return;
	}
	if(place->right){
		node<T> *predecessor = place->right;
		while(predecessor->left){
			predecessor = predecessor->left;
		}
		T hold_data = predecessor->data;
		remove(predecessor->data);
		place->data = hold_data; // Will this cause an extra destructor call?
		return;
	}

	// This only works if place is a leaf
	// That's the only case left!  place->left and place->right are both null
	if(parent->right == place)
		parent->right = 0;
	else
		parent->left = 0;
	delete place;
}


// A workaround for the parameter
// Caution:  Don't use this if you expect the recursion to stop when place == null
template<class T>
void Tree<T>::print_inorder(node<T>* place){
	if(!place){
		place = root;
		if(place->parent)
			cout << "Error:  Root has parent!\n";
	}
	// This would be where code for preorder traversal goes
	if(place->left)
		print_inorder(place->left);
	// This place for inorder

	if(place->left && place->left->parent != place)
		cout << "  LM";
	cout << place->data << " ";
	if(place->right && place->right->parent != place)
		cout << "RM  ";

	if(place->right)
		print_inorder(place->right);
	// This place down here for postorder traversal
	if(place == root)
		cout << endl;
}

template<class T>
void Tree<T>::rotate_left(node<T> **place){
	cout << "Rotate Left\n";
	if(place == 0)
		place = &root;
	node<T> *old_root = *place;
	node<T> *new_root = (*place)->right;
	if(!new_root)
		return;
	new_root->parent = old_root->parent;
	old_root->parent = new_root;
	old_root->right = new_root->left;
	if(old_root->right)
		old_root->right->parent = old_root;
	new_root->left = old_root;
	*place = new_root;
	if(*place == root)
		root = new_root;
	if(root->parent)
		cout << "Oh no!  The root has a parent!\n";
	
}

template<class T>
void Tree<T>::rotate_right(node<T> **place){
	cout << "Rotate Right\n";
	if(place == 0)
		place = &root;
	node<T> *old_root = *place;
	node<T> *new_root = (*place)->left;
	if(!new_root)
		return;
	new_root->parent = old_root->parent;
	old_root->parent = new_root;
	old_root->left = new_root->right;
	if(old_root->left)
		old_root->left->parent = old_root;
	new_root->right = old_root;
	*place = new_root;
	if(*place == root)
		root = new_root;
	if(root->parent)
		cout << "Oh no!  The root has a parent!\n";
}

/* Sorry, it's not this simple! */
template<class T>
void Tree<T>::improve_balance(node<T> **place){
	/* if the balance is really bad:
	 * 	Do a rotation the correct way to make it better!
	 */
	if(place == 0)
		place = &root;
	if(! *place)
		return;
	improve_balance(&(*place)->right);
	if(! *place)
		return;
	improve_balance(&(*place)->left);
	if(! *place)
		return;
	int left_depth = depth_from_place((*place)->left);
	int right_depth = depth_from_place((*place)->right);
	if(left_depth > right_depth + 1)
		rotate_right(&root);
	else if(right_depth > left_depth + 1)
		rotate_left(&root);
	else
		cout << "No rotation:  " << left_depth << " " << right_depth << "\n";
}

int maximum(int a, int b){
	return (a>b)? a:b;
}
template<class T>
int Tree<T>::depth_from_place(node<T> *place){
	if(!place)
		return 0;
	//	return maximum(depth_from_place(place->left), depth_from_place(place->right)) + 1;
	return depth_from_place(place->left) + depth_from_place(place->right) + 1;
}
#endif