#include<graphviz/cgraph.h>
#include<graphviz/gvc.h>
#include<string>
#include<ostream>
#include "better_linkedlist.h"
#include<iostream>

template<typename keytype, typename valuetype>
class kv_rbtree {
	struct node {
		keytype key;
		valuetype value;
		bool isred = true;
		node *left = 0, *right = 0, *parent = 0; // Will we need a parent?
	};
	class kv_iterator {
		// We need to keep track of the state of our traversal
		list<node*> stack;
		node *prev = 0;

		public:
		kv_iterator(node* root, bool end) {
			if(!end){
				stack.push(root);
				++(*this);
			} // We'll leave the stack empty if end is true
		}
		kv_iterator& operator++(){
			// Advance our traversal by one
			if(stack.empty()) // We can't traverse anymore
				return *this;
			while(true){
				node *current = stack[0];
				// No previous node, we're starting here
				// Or we're on our way down
				if(!prev || prev->left == current || prev->right == current){
					if(current->left){
						stack.push(current->left);
						prev = current;
					} else if (current->right) {
						stack.push(current->right);
						prev = current;
					} else {
						stack.pop();
						prev = current;
						return *this;
					}
				} else if (current->left == prev) {
					if(current->right) {
						stack.push(current->right);
						prev = current;
					}
					else {
						stack.pop();
						prev = current;
						return *this;
					}
				} else if (current->right == prev) {
					stack.pop();
					prev = current;
					return *this;
				}
			}
		}
		bool operator!=(const kv_iterator &other) const {
			// Return true if other is a traversal in a different spot
			if(stack.empty() && other.stack.empty())
				return false;
			else if(stack.empty() || other.stack.empty())
				return true;
			else if(stack[0] != other.stack[0])
				return true;
			// Swept under rug:  If the stacks are different heights, maybe 
			// one is on the way down, and the other on the way up
			// To fix:  Just check if the stacks are the same size
			// (better_linkedlist.h doesn't actually have a method for this)
			return false;
		}
		keytype& operator*(){
			return prev->key;
		}
	};

	node* root = nullptr; // nullptr means 0 in C++
	public:

	valuetype& operator[](const keytype& key){
		if(!root){
			root = new node;
			root->key = key;
			root->isred = false;
			return root->value;
		}
		node *current = root;
		while(true){
			if(current->key == key) // This was in the wrong spot last time
				return current->value;
			if(key < current->key){
				if(current->left) {
					current = current->left;
				} else {
					node *nn = new node;
					current->left = nn;
					current->left->key = key;
					current->left->parent = current;
					insert_rebalance(current->left);
					return nn->value;
				}	
			} else { // Right branch (key indicates we go right)
				if(current->right) {
					current = current->right;
				} else {
					node *nn = new node;
					current->right = nn;
					current->right->key = key;
					current->right->parent = current;
					insert_rebalance(current->right);
					return nn->value; // Problem here:  current might not have right anymore
				}	
			}
		}
	}
	void insert_rebalance(node* new_node){
		// Case I3
		if(!new_node->parent)
			return;
		// Case I1
		if(new_node->parent->isred == false)
			return;
		// Case I2
		node *uncle = 0;
		if(new_node->parent->parent){ // There's a grandparent
			uncle = new_node->parent->parent->left;
			if(new_node->parent->parent->left == new_node->parent)
				uncle = new_node->parent->parent->right;
			if(uncle && uncle->isred) { // We should do case I2
				uncle->isred = false;
				new_node->parent->isred = false;
				new_node->parent->parent->isred = true;
				insert_rebalance(new_node->parent->parent);
				return;
			}
		} else { // Case I4
			new_node->parent->isred = false;
			return;
		}
		// Case I5
		// Because of other cases, we know either there's no uncle or the uncle is black
		if(new_node->parent->right == new_node && new_node->parent->parent->left == new_node->parent){
			// Rotate us to the left 
			node *G = new_node->parent->parent;
			node *P = new_node->parent;
			G->left = new_node;
			P->right = new_node->left;  // Problem 1:  These were out of order
			if(P->right)
				P->right->parent = P;   // Problem 2:  This parent isn't set
			new_node->left = P;
			new_node->parent = G;
			P->parent = new_node;
			new_node = P;
		}
		else if(new_node->parent->left == new_node && new_node->parent->parent->right == new_node->parent){
			// Rotate us to the right
			node *G = new_node->parent->parent;
			node *P = new_node->parent;
			G->right = new_node;
			P->left = new_node->right; // Problem 1:  These were are out of order
			if(P->left)
				P->left->parent = P; // Problem 2:  This parent isn't set
			new_node->right = P;
			new_node->parent = G;
			P->parent = new_node;
			new_node = P;
		}
		// Case I6

		node **GG = &root;
		node *G = new_node->parent->parent;
		if(G->parent){
			if(G->parent->left == G)
				GG = &(G->parent->left);
			else
				GG = &(G->parent->right);
		}
		node *P = new_node->parent;
		*GG = P;
		if(new_node->parent->left == new_node){	
			G->left = P->right;
			P->right = G;
			if(G->left)
				G->left->parent = G;
		} else if(new_node->parent->right == new_node){	
			G->right = P->left;
			P->left = G;
			if(G->right)
				G->right->parent = G;
		}
		P->parent = G->parent;
		G->parent = P;
		P->isred = false;
		G->isred = true;
	}

	void make_image_recur(Agraph_t *g, Agnode_t *parent, node *place, int depth = -1){
		Agnode_t *us = agnode(g, (char*)to_string(place->key).c_str(), 1);
		if(place->isred)
			agset(us, (char*)"color", (char*)"red");
		if(parent)
			Agedge_t *e = agedge(g, parent, us, (char*)"", 1);
		if(depth == 0)
			return;
		if(depth != -1)
			depth--;
		if(place->left)
			make_image_recur(g, us, place->left, depth);
		if(place->right)
			make_image_recur(g, us, place->right, depth);
	}

	void make_image(std::string name, int depth = -1){
		Agraph_t *g;
		g = agopen((char*)"G", Agdirected, 0);
		Agsym_t *color_attr = agattr(g, AGNODE, (char*)"color", (char*)"black");

		make_image_recur(g, 0, root, depth);

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

	void print_inorder(std::ostream& out) const { // Recursive, so O(log(n) time and O(log(n)) memory
		print_inorder_r(root, out);
	}
	void print_inorder_r(node *place, std::ostream& out) const{
		if(!place)
			return;
		// Preorder
		print_inorder_r(place->left, out);
		// Inorder
		out << " (" << place->key << ", " << place->value << ") "; // Avoid cout in header files
		print_inorder_r(place->right, out);
		// Postorder 
	}

	kv_iterator begin(){
		return kv_iterator(root, false);
	}

	kv_iterator end() {
		return kv_iterator(root, true);
	}

	// With stack, not recursive.  Recursive is way shorter.  Look at the Fall 21 example for that.
	~kv_rbtree() {
		list<node*> stack;
		stack.push(root);
		node *prev = 0;
		while(!stack.empty()){
			// print current state
			for(node* c : stack)
				std::cout << c->key << " ";
			std::cout << std::endl;

			node *current = stack[0];

			// No previous node, we're starting here
			// Or we're on our way down
			if(!prev || prev->left == current || prev->right == current){
				if(current->left)
					stack.push(current->left);
				else if (current->right)
					stack.push(current->right);
				else {
					delete current;
					stack.pop();
				}
			} else if (current->left == prev) {
				if(current->right)
					stack.push(current->right);
				else {
					delete current;
					stack.pop();
				}
			} else if (current->right == prev) {
				delete current;
				stack.pop();
			}
			prev = current;
		}	
	}
};

template<typename keytype, typename valuetype>
std::ostream& operator<<(std::ostream& out, const kv_rbtree<keytype, valuetype>& T){
	return out << T;
}