#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.h>
#include "avl.h"

#define TRUE 1
#define FALSE 0

// Insert element into tree, return FALSE, if node is already in tree
int insert(uint32_t num, Node newNode, Node* root) {
	if (*root == NULL) {
		*root = newNode;
		return TRUE;
	}
	
	int returnValue = FALSE;
	Node n = *root;
	while (TRUE) {
		if (n->value > num) {
			if (n->leftChild == NULL) {
				n->leftChild = newNode;
				newNode->parent = n;
				returnValue = TRUE;
				break;
			} else {
				n = n->leftChild;
			}
		} else if (n->value < num) {
			if (n->rightChild == NULL) {
				n->rightChild = newNode;
				newNode->parent = n;
				returnValue = TRUE;
				break;
			} else {
				n = n->rightChild;
			}
                } else {
		  break;
                }
	}
       	if (returnValue) balance(root, n, newNode);
	return returnValue;
}

void balance(Node* root, Node n, Node child) {
  int nBalance;
  int childBalance;
  Node temp;
  while (n != NULL) {
    nBalance = ((n->rightChild != NULL)?n->rightChild->height:-1) -
               ((n->leftChild != NULL)?n->leftChild->height:-1);
    if (nBalance == -2) {
      childBalance = ((child->rightChild != NULL)?child->rightChild->height:-1) -
	             ((child->leftChild != NULL)?child->leftChild->height:-1);
      if (childBalance == -1) {
	// single rotation on left child
	temp = child->rightChild;
	child->rightChild = n;
	n->leftChild = temp;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = child;
	  } else {
	    n->parent->rightChild = child;
	  }	
	} else {
	  *root = child;
	}
	child->parent = n->parent;
	n->parent = child;
	if (temp != NULL) temp->parent = n;
	n->height--;
	return;
      } else {
	// double rotation on left child
	Node temp = child->rightChild;
	n->leftChild = temp->rightChild;
	child->rightChild = temp->leftChild;
	temp->leftChild = child;
	temp->rightChild = n;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = temp;
	  } else {
	    n->parent->rightChild = temp;
	  }	
	} else {
	  *root = temp;
	}
	temp->parent = n->parent;
	child->parent = temp;
	n->parent = temp;
	if (child->rightChild != NULL) child->rightChild->parent = child;
	if (n->leftChild != NULL) n->leftChild->parent = n;
	temp->height++; 
	child->height--;
	n->height--;
	return;
      }
    } else if (nBalance == 2) {
      childBalance = ((child->rightChild != NULL)?child->rightChild->height:-1) -
	             ((child->leftChild != NULL)?child->leftChild->height:-1);
      if (childBalance == 1) {
	// single rotation on right child
	Node temp = child->leftChild;
	child->leftChild = n;
	n->rightChild = temp;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = child;
	  } else {
	    n->parent->rightChild = child;
	  }	
	} else {
	  *root = child;
	}
	child->parent = n->parent;
	n->parent = child;
	if (temp != NULL) temp->parent = n;
	n->height--;
	return;
      } else {
	// double rotation on right child
	Node temp = child->leftChild;
	n->rightChild = temp->leftChild;
	child->leftChild = temp->rightChild;
	temp->rightChild = child;
	temp->leftChild = n;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = temp;
	  } else {
	    n->parent->rightChild = temp;
	  }	
	} else {
	  *root = temp;
	}
	temp->parent = n->parent;
	child->parent = temp;
	n->parent = temp;
	if (child->leftChild != NULL) child->leftChild->parent = child;
	if (n->rightChild != NULL) n->rightChild->parent = n;
	temp->height++; 
	child->height--;
	n->height--;
	return;
      }
    } else {
      n->height = max((n->rightChild != NULL)?n->rightChild->height:-1,
		      (n->leftChild != NULL)?n->leftChild->height:-1) + 1;
      child = n;
      n = n->parent;
    }
  }
}

inline int max(int x, int y) {
  return (x>y)?x:y;
}

inline int min(int x, int y) {
  return (x<y)?x:y;
}

// Finds a node, returns node or FALSE, if not found.
int get(uint32_t num, Node root) {
	if (root == NULL) return FALSE;

	Node n = root;
	while (TRUE) {
		if (n->value > num) {
			if (n->leftChild == NULL) {
				return FALSE;
			} else {
				n = n->leftChild;
			}
		} else if (n->value < num) {
			if (n->rightChild == NULL) {
				return FALSE;
			} else {
				n = n->rightChild;
			}
		} else {
			return num;
		}
	}
}


int height(Node root) {
  if (root == NULL) return -1;
  return heightHelper(root);
}
		   

int heightHelper(Node n) {
  int lheight = 0;
  int rheight = 0;
  if (n->leftChild != NULL) lheight = heightHelper(n->leftChild);
  if (n->rightChild != NULL) rheight = heightHelper(n->rightChild);
  if (lheight > rheight) return lheight + 1;
  return rheight + 1;
}


// Finds max node
uint32_t findMax(Node n) {
	while (n->rightChild != NULL) {
		n = n->rightChild;
	}
	return n->value;
}



// Removes node, returns TRUE if the node was in the tree and FALSE otherwise
int removeNode(uint32_t num, Node* root) {
  if (*root == NULL) {
    return FALSE;
  }
  Node n = *root;
  Node parent = NULL;
  Node child = NULL;
  while (TRUE) {
    if (n->value > num) {
      if (n->leftChild == NULL) {
	return FALSE;
      } else {
	n = n->leftChild;
      }
    } else if (n->value < num) {
      if (n->rightChild == NULL) {
	return FALSE;
      } else {
	n = n->rightChild;
      }
    } else { // found it
      if (n->leftChild != NULL) {
	if (n->rightChild != NULL) {
	  // two children
	  uint32_t temp = findMax(n->leftChild);
	  n->value = temp;
	  num = temp;
	  n = n->leftChild;
	} else {
	  // left child
	  child = n->leftChild;
	  child->parent = n->parent;
	  if (n->parent != NULL) {
	    parent = n->parent;
	    if (n->parent->leftChild == n) {
	      n->parent->leftChild = child;
	    } else {
	      n->parent->rightChild = child;
	    }
	  } else {
	    *root = child;
	  }
	  break;
	}
      } else {
	if (n->rightChild != NULL) {
	  // right child
	  child = n->rightChild;
	  child->parent = n->parent;
	  if (n->parent != NULL) {
	    parent = n->parent;
	    if (n->parent->leftChild == n) {
	      n->parent->leftChild = child;
	    } else {
	      n->parent->rightChild = child;
	    }
	  } else {
	    *root = child;
	  }
	  break;
	} else {
	  // leaf
	  if (n->parent != NULL) {
	    parent = n->parent;
	    if (n->parent->leftChild == n) {
	      n->parent->leftChild = NULL;
	    } else {
	      n->parent->rightChild = NULL;
	    }
	  } else {
	    *root = NULL;
	  }
	  break;
	}
      }
    }
  }
  if (parent != NULL)  balance2(root, parent);
  return TRUE;
}


void balance2(Node* root, Node n) {
  int nBalance;
  int childBalance;
  Node temp;
  Node child;
  while (n != NULL) {
    nBalance = ((n->rightChild != NULL)?n->rightChild->height:-1) -
               ((n->leftChild != NULL)?n->leftChild->height:-1);
    if (nBalance == -2) {
      child = n->leftChild;
      childBalance = ((child->rightChild != NULL)?child->rightChild->height:-1) -
	             ((child->leftChild != NULL)?child->leftChild->height:-1);
      if (childBalance == 1) {
	// double rotation on left child
	Node temp = child->rightChild;
	n->leftChild = temp->rightChild;
	child->rightChild = temp->leftChild;
	temp->leftChild = child;
	temp->rightChild = n;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = temp;
	  } else {
	    n->parent->rightChild = temp;
	  }	
	} else {
	  *root = temp;
	}
	temp->parent = n->parent;
	child->parent = temp;
	n->parent = temp;
	if (child->rightChild != NULL) child->rightChild->parent = child;
	if (n->leftChild != NULL) n->leftChild->parent = n;
	temp->height++; 
	child->height--;
	n->height--;
      } else {
	// single rotation on left child
	temp = child->rightChild;
	child->rightChild = n;
	n->leftChild = temp;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = child;
	  } else {
	    n->parent->rightChild = child;
	  }	
	} else {
	  *root = child;
	}
	child->parent = n->parent;
	n->parent = child;
	if (temp != NULL) temp->parent = n;
	n->height--;
      }
    } else if (nBalance == 2) {
      child = n->rightChild;
      childBalance = ((child->rightChild != NULL)?child->rightChild->height:-1) -
	             ((child->leftChild != NULL)?child->leftChild->height:-1);
      if (childBalance == -1) {
	// double rotation on right child
	Node temp = child->leftChild;
	n->rightChild = temp->leftChild;
	child->leftChild = temp->rightChild;
	temp->rightChild = child;
	temp->leftChild = n;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = temp;
	  } else {
	    n->parent->rightChild = temp;
	  }	
	} else {
	  *root = temp;
	}
	temp->parent = n->parent;
	child->parent = temp;
	n->parent = temp;
	if (child->leftChild != NULL) child->leftChild->parent = child;
	if (n->rightChild != NULL) n->rightChild->parent = n;
	temp->height++; 
	child->height--;
	n->height--;
      } else {
	// single rotation on right child
	Node temp = child->leftChild;
	child->leftChild = n;
	n->rightChild = temp;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = child;
	  } else {
	    n->parent->rightChild = child;
	  }	
	} else {
	  *root = child;
	}
	child->parent = n->parent;
	n->parent = child;
	if (temp != NULL) temp->parent = n;
	n->height--;
      }
    } else {
      n->height = max((n->rightChild != NULL)?n->rightChild->height:-1,
		      (n->leftChild != NULL)?n->leftChild->height:-1) + 1;
      n = n->parent;
    }
  }
}


void printTree(Node n) {
	if (n == NULL) return;
	printf("%d %d %d\n",n->value, n->height, (n->parent != NULL)?n->parent->value:-1);
	printTree(n->leftChild);
	printTree(n->rightChild);
}

uint32_t treeSize(Node n) { 
	if (n == NULL) return 0;
	return treeSize(n->leftChild) + treeSize(n->rightChild) + 1;
}

/*
 * Create a new node
 */
Node makeNode(uint32_t num) {
	Node n = (Node)malloc(sizeof(struct _Node));
	if (n == NULL) {
		printf("Error allocating memory for node with num %d\n", num);
		exit(1);
	}
	n->value = num;
	n->leftChild = NULL;
	n->rightChild = NULL;
	n->parent = NULL;
	n->height = 0;
	return n;
}