#include <stdio.h>
#include <stdlib.h>

	/* Handy definitions to make working with the AVL Node structure
	   less insane... */
	#define _path_left	0x00
	#define _path_right	0x01

	/* Return definitions */
	#define _def_failure	0x01
	#define _def_success	0x00

	/* Balance Definitions */
	#define _bal_balanced		0
	#define _bal_left_heavy		-1
	#define _bal_right_heavy	1

	/* AVL Node definition */
	typedef struct avl_node_def avl_node_t;

	struct avl_node_def {
		avl_node_t *path[0x02];
		unsigned int key;
		char balance;
	};

	/* This array tracks the Nodes that we pass through during
	   Insertion, Traversal or Removal. */
	avl_node_t *rh_array[0x100];
	unsigned int rh_index = 0;

	unsigned int n_count = 0;

	/* This function will allocate and initialize an AVL Node */
	avl_node_t *create_node_with_key(unsigned int key) {
		avl_node_t *new_node = (avl_node_t*)malloc(sizeof(avl_node_t));
		if(new_node == 0) return (avl_node_t*)_def_failure;

		new_node->balance = _bal_balanced;
		new_node->path[_path_left] = 0;
		new_node->path[_path_right] = 0;
		new_node->key = key;
	
		return new_node;
	}

	avl_node_t *rotate_right(avl_node_t *node) {
		avl_node_t *link_node = node->path[_path_left];
		node->path[_path_left] = link_node->path[_path_right];
		link_node->path[_path_right] = node;
		return link_node;
	}

	avl_node_t *rotate_left(avl_node_t *node) {
		avl_node_t *link_node = node->path[_path_right];
		node->path[_path_right] = link_node->path[_path_left];
		link_node->path[_path_left] = node;
		return link_node;
	}

	avl_node_t *rebalance_left(avl_node_t *node) {
		if(node->path[_path_left]->balance > _bal_balanced) {
			if(node->path[_path_left]->path[_path_right]->balance < _bal_balanced) {
				node->path[_path_left]->balance = _bal_balanced;
				node->path[_path_left]->path[_path_right]->balance = _bal_balanced;
				node->balance = _bal_right_heavy;
			} else if(node->path[_path_left]->path[_path_right]->balance > _bal_balanced) {
				node->path[_path_left]->balance = _bal_left_heavy;
				node->path[_path_left]->path[_path_right]->balance = _bal_balanced;
				node->balance = _bal_balanced;
			}
			node->path[_path_left] = rotate_left(node->path[_path_left]);
		} else {
			node->balance = _bal_balanced;
			node->path[_path_left]->balance = _bal_balanced;
		}
		return rotate_right(node);
	}

	avl_node_t *rebalance_right(avl_node_t *node) {
		if(node->path[_path_right]->balance < _bal_balanced) {
			if(node->path[_path_right]->path[_path_left]->balance < _bal_balanced) {
				node->path[_path_right]->balance = _bal_balanced;
				node->path[_path_right]->path[_path_left]->balance = _bal_balanced;
				node->balance = _bal_left_heavy;
			}  else if(node->path[_path_right]->path[_path_left]->balance > _bal_balanced) {
				node->path[_path_right]->balance = _bal_right_heavy;
				node->path[_path_right]->path[_path_left]->balance = _bal_balanced;
				node->balance = _bal_balanced;
			}
			node->path[_path_right] = rotate_right(node->path[_path_right]);
		} else {
			node->balance = _bal_balanced;
			node->path[_path_right]->balance = _bal_balanced;
		}
		return rotate_left(node);
	}

	avl_node_t *insert_node(avl_node_t **tree, avl_node_t *node) {
		if((tree == 0) || (node == 0)) return (avl_node_t*)_def_failure;

		if(*tree == 0) {
			*tree = node;
			return node;
		} else {
			avl_node_t *insert_ptr = (avl_node_t*)*tree;
			rh_index = 0;

			while(rh_index < 0x100) {
				rh_array[rh_index] = insert_ptr;
				rh_index++;

				if(node->key == insert_ptr->key) {
					free(node);
					return 0;
				} else if(node->key < insert_ptr->key) {
					if(insert_ptr->path[_path_left] != 0) {
						insert_ptr = insert_ptr->path[_path_left];
					} else {
						insert_ptr->path[_path_left] = node;
						break;
					}
				} else if(node->key > insert_ptr->key) {
					if(insert_ptr->path[_path_right] != 0) {
						insert_ptr = insert_ptr->path[_path_right];
					} else {
						insert_ptr->path[_path_right] = node;
						break;
					}
				}
			}
			if(rh_index == 0x100) {
				free(node);
				return (avl_node_t*)0;
			} else {
				while(rh_index != 0) {	
					rh_index--;
					insert_ptr = rh_array[rh_index];
					if(node->key < insert_ptr->key) insert_ptr->balance -= 1;
					else insert_ptr->balance += 1;

					if(insert_ptr->balance < _bal_left_heavy) {
						if((avl_node_t*)*tree == insert_ptr) {
							*tree = rebalance_left(insert_ptr);
							break;
						} else {
							if(insert_ptr->key < rh_array[rh_index-1]->key)
								rh_array[rh_index-1]->path[_path_left] = rebalance_left(insert_ptr);
							else
								rh_array[rh_index-1]->path[_path_right] = rebalance_left(insert_ptr);
							break;
						}
					} else if(insert_ptr->balance > _bal_right_heavy) {
						if((avl_node_t*)*tree == insert_ptr) {
							*tree = rebalance_right(insert_ptr);
							break;
						} else {
							if(insert_ptr->key < rh_array[rh_index-1]->key)
								rh_array[rh_index-1]->path[_path_left] = rebalance_right(insert_ptr);
							else
								rh_array[rh_index-1]->path[_path_right] = rebalance_right(insert_ptr);
							break;
						}
					}	
				}
			}

			return node;
		}
	}
	
	avl_node_t *find_node(avl_node_t **root, unsigned int key) {
		if((root == 0) || (*root == 0)) return (avl_node_t*)_def_failure;
		else {
			avl_node_t *search_ptr = (avl_node_t*)*root;
			rh_index = 0;
		
			while((rh_index < 0x100) || (search_ptr == 0)) {
				rh_array[rh_index] = search_ptr;

				if(key == search_ptr->key) return search_ptr;
				else if(key < search_ptr->key) search_ptr = search_ptr->path[_path_left];
				else if(key > search_ptr->key) search_ptr = search_ptr->path[_path_right];	
				rh_index++;
			}
			return (avl_node_t*)_def_failure;
		}
	}

	inline char isLeaf(avl_node_t *node) {
		if(!node->path[_path_left] && !node->path[_path_right]) return 1;
		else return 0;
	}

	void delete_node(avl_node_t **root, avl_node_t *node) {
		if((root == 0) || (node == 0) || (*root == 0)) return;
		else {
			avl_node_t *node_ptr = 0;
			unsigned int key_storage = 0;

			if(node->path[_path_left] && node->path[_path_right]) {
				printf(" : Subtree\r\n");

				// Locate Replacement Node
				if(isLeaf(node->path[_path_left])==0) {
					node_ptr = node->path[_path_left];
					while(node_ptr->path[_path_right]) node_ptr = node_ptr->path[_path_right];
				} else {
					node_ptr = node->path[_path_right];
					while(node_ptr->path[_path_left]) node_ptr = node_ptr->path[_path_left];
				}

				// Store its Key.
				key_storage = node_ptr->key;

				// Remove the Node.
				delete_node(root, find_node(root, key_storage));

				// Patch the Subtree Node
				node->key = key_storage;

				// Locate the Subtree Node (to restore the rh_array)
				find_node(root, key_storage);
			} else if(node->path[_path_left] && !node->path[_path_right]) {
				printf(" : Left branch\r\n");
				if(node == *root) {
					*root = node->path[_path_left];
					free(node);
					return;
				}

				if(node->key < rh_array[rh_index-1]->key) {
					rh_array[rh_index-1]->path[_path_left] = node->path[_path_left];
					rh_array[rh_index-1]->balance += 1;
				} else {
					rh_array[rh_index-1]->path[_path_right] = node->path[_path_left];
					rh_array[rh_index-1]->balance -= 1;
				}
				free(node);
			} else if(!node->path[_path_left] && node->path[_path_right]) {
				printf(" : Right branch\r\n");
				if(node == *root) {
					*root = node->path[_path_right];
					free(node);
					return;
				}

				if(node->key < rh_array[rh_index-1]->key) {
					rh_array[rh_index-1]->path[_path_left] = node->path[_path_right];
					rh_array[rh_index-1]->balance += 1;
				} else {
					rh_array[rh_index-1]->path[_path_right] = node->path[_path_right];
					rh_array[rh_index-1]->balance -= 1;
				}
				free(node);
			} else {
				printf(" : Leaf\r\n");
				if(node == *root) {
					*root = 0;
					free(node);
					return;
				}

				if(node->key < rh_array[rh_index-1]->key) {
					rh_array[rh_index-1]->balance += 1;
					rh_array[rh_index-1]->path[_path_left] = 0;
				} else {
					rh_array[rh_index-1]->balance -= 1;
					rh_array[rh_index-1]->path[_path_right] = 0;
				}
				free(node);
			}

			while(rh_index > 0) {
				node_ptr = rh_array[--rh_index];
				printf(" ? Node(%d), balance: %d\r\n", node_ptr->key, node_ptr->balance);

				if(node_ptr->balance == 0) {
					if(node_ptr != *root) {
						if(node_ptr->key < rh_array[rh_index-1]->key) 
							rh_array[rh_index-1]->balance += 1;
						else
							rh_array[rh_index-1]->balance -= 1;
					}
				} else if(node_ptr->balance < _bal_left_heavy) {
					if(node_ptr == *root) {
						*root = rebalance_left(node_ptr);
						break;
					} else {
						if(node_ptr->key < rh_array[rh_index-1]->key)
							rh_array[rh_index-1]->path[_path_left] = rebalance_left(node_ptr);
						else
							rh_array[rh_index-1]->path[_path_right] = rebalance_left(node_ptr);
						break;
					}
				} else if(node_ptr->balance > _bal_right_heavy) {
					if(node_ptr == *root) {
						*root = rebalance_right(node_ptr);
						break;
					} else {
						if(node_ptr->key < rh_array[rh_index-1]->key)
							rh_array[rh_index-1]->path[_path_left] = rebalance_right(node_ptr);
						else
							rh_array[rh_index-1]->path[_path_right] = rebalance_right(node_ptr);
						break;
					}
				} else if((node_ptr->balance == _bal_left_heavy) ||\
					  (node_ptr->balance == _bal_right_heavy)) {
					break;
				} 
			}
		}
	}

	void show_tree(avl_node_t **root, avl_node_t *tree) {
		if(tree != 0) {
			show_tree(root, tree->path[_path_left]);

			printf(" - Node 0x%x K %d, L: 0x%x, R: 0x%x, B: %d ", (unsigned int)tree,\
									      (unsigned int)tree->key,\
									      (unsigned int)tree->path[_path_left],\
									      (unsigned int)tree->path[_path_right],\
									      (unsigned int)tree->balance); 

			if(tree == *root) printf(" << root\r\n");
			else if(tree->path[_path_left] && tree->path[_path_right]) printf(" << subtree\r\n");
			else if(tree->path[_path_left]) printf(" << left branch\r\n");
			else if(tree->path[_path_right]) printf(" << right branch\r\n");
			else printf(" << leaf\r\n");

			show_tree(root, tree->path[_path_right]);
		}
	}

	char check_tree(avl_node_t **root, avl_node_t *node) {
		if(node == 0) return 0;
		else {
			if(check_tree(root, node->path[_path_left])==1) return 1;
	
			if(node->balance == _bal_left_heavy) {
				if(node->path[_path_left] == 0) {
					printf(" @ Node(%d), Left-heavy without left child!\r\n", node->key);
					return 1;
				}
			} 
			if((node->path[_path_left]) && (node->key < node->path[_path_left]->key)) {
				printf(" @ Node(%d), Smaller than Left Child!\r\n", node->key);
				return 1;
			}	

			if(node->balance == _bal_right_heavy) {
				if(node->path[_path_right] == 0) {
					printf(" @ Node(%d), Right-heavy without right child!\r\n", node->key);
					return 1;
				}
			}	
			if((node->path[_path_right]) && (node->key > node->path[_path_right]->key)) {
				printf(" @ Node(%d), Greater than Right Child!\r\n", node->key);
				return 1;
			}	
	
			if(check_tree(root, node->path[_path_right])==1) return 1;

			return 0;
		}
	}

	unsigned int count_nodes(avl_node_t *tree) {
		if(tree == 0) return n_count;
		else {
			count_nodes(tree->path[_path_left]);

			n_count ++;

			count_nodes(tree->path[_path_right]);
			return n_count;
		}
	}

	int main(void) {
		avl_node_t *test_tree = 0;
		unsigned int ac = 0;

		for(unsigned int tc = 0x1000; tc > 0; tc--) {
			if(insert_node(&test_tree, create_node_with_key(rand()%0x1000))!=0) ac++;
			if(check_tree(&test_tree, test_tree)!=0) {
				printf("fail.\r\n");
				break;
			}
		}/*
		show_tree(&test_tree, test_tree);

		unsigned int a = 0;

///		printf("NC: %d\r\n", count_nodes(test_tree));

		while(test_tree) {
			a++;
			n_count = 0;
			delete_node(&test_tree, find_node(&test_tree, test_tree->key));
			//show_tree(&test_tree, test_tree);
			printf(" @ Checking Tree Integrity ... ");
			if(check_tree(&test_tree, test_tree)!=0) {
				printf("failure.\r\n");
			} else printf("OK!\r\n");

			printf(" $ Checking Count (n = %d) ... ", (ac- a));
			if(count_nodes(test_tree) == (ac-a)) printf("OK.\r\n");
			else {
				printf("FAIL!\r\n");
				break;
			}
		}
*/
//		if(check_tree(&test_tree, test_tree)!=0) { printf(" >>> Tree failure.\r\n");
	//	} else printf(" >>> Tree good!\r\n");

		system("PAUSE");

		return 0;
	}