#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* root) {
	Node newNode = makeNode(num);

	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) {
  while (n != NULL) {
    (n->rightChild == child)?n->balance++:n->balance--;
    if (n->balance == 0) return;
    if (n->balance == -1 || n->balance == 1) {
	  child = n;
          n = n->parent;
    } else if (n->balance == -2) {
      if (child->balance == -1) {
	// single rotation on left child
	Node 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;
	child->balance++;
	n->balance = - child->balance;
	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->leftChild->balance = - max(temp->balance, 0);
	temp->rightChild->balance = - min(temp->balance, 0);
	temp->balance = 0;
	return;
      }
    } else {
      if (child->balance == 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;
	child->balance--;
	n->balance = - child->balance;
	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->leftChild->balance = - max(temp->balance, 0);
	temp->rightChild->balance = - min(temp->balance, 0);
	temp->balance = 0;
	return;
      }
    }
  }
}

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 n) {
  int lheight = 0;
  int rheight = 0;
  if (n->leftChild != NULL) lheight = height(n->leftChild);
  if (n->rightChild != NULL) rheight = height(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
// >>>>> Need to free memory <<<<<<
int removeNode(uint32_t num, Node* root) {
  if (*root == NULL) {
    return FALSE;
  }
  Node n = *root;
  int offset = 0;
  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 root is to be removed and has two children, 
      // the following if statement will be called twice.
      if (n->parent != NULL) {
	if (n->parent->leftChild == n) {
	  offset = 1;
	} else {
	  offset = -1;
	}
      }
      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
	  n->leftChild->parent = n->parent;
	  if (n->parent != NULL) {
	    if (n->parent->leftChild == n) {
	      n->parent->leftChild = n->leftChild;
	    } else {
	      n->parent->rightChild = n->leftChild;
	    }
	  } else {
	    *root = n->leftChild;
	  }
	  break;
	}
      } else {
	if (n->rightChild != NULL) {
	  // right child
	  n->rightChild->parent = n->parent;
	  if (n->parent != NULL) {
	    if (n->parent->leftChild == n) {
	      n->parent->leftChild = n->rightChild;
	    } else {
	      n->parent->rightChild = n->rightChild;
	    }
	  } else {
	    *root = n->rightChild;
	  }
	  break;
	} else {
	  // leaf
	  if (n->parent != NULL) {
	    if (n->parent->leftChild == n) {
	      n->parent->leftChild = NULL;
	    } else {
	      n->parent->rightChild = NULL;
	    }
	  } else {
	    *root = NULL;
	  }
	  break;
	}
      }
    }
  }
  if (n->parent != NULL) {
    n->parent->balance += offset; // check
    balance2(root, n->parent);
  }
  return TRUE;
}


void balance2(Node* root, Node n) {
  printf("balance on: %d\n", n->value);
  Node temp = n;
  while (temp != NULL) {
    printf("parent: %d\n", temp->value);
    temp = temp->parent;
  }

  Node child = NULL;
  while (n != NULL) {
    printf("now: %d  balance:%d \n", n->value, n->balance);
    if (n->balance == -2) {
      // >>>>>>>>>> fix the following if by comparing againt 1 and flipping the rotations
      if (n->leftChild->balance == -1 || n->leftChild->balance == 0) {
	// single rotation on left child
	// cased fixed
	Node child = n->leftChild;
	int childBalance = child->balance;
	Node temp = child->rightChild;
	Node p = n->parent;
	child->rightChild = n;
	n->leftChild = temp;
	if (p != NULL) {
	  if (p->leftChild == n) {
	    p->leftChild = child;
	    if (child->balance == -1) {
	      p->balance++;
	    }
	  } else {
	    p->rightChild = child;
	    if (child->balance == -1) {
	      p->balance--;
	    }
	  }	
	} else {
	  *root = child;
	}
	child->parent = p;
	if (temp != NULL) temp->parent = n;
	n->parent = child;
	child->balance++;
	n->balance = - child->balance;
	if (childBalance == 0) return;
	n = p;
      } else {
	// double rotation on left child
	Node child = n->leftChild;
	Node temp = child->rightChild;
	Node p = n->parent;
	child->rightChild = temp->leftChild;
	n->leftChild = temp->rightChild;
	temp->leftChild = child;
	temp->rightChild = n;
	if (p != NULL) {
	  if (p->leftChild == n) {
	    p->leftChild = temp;
	    p->balance++;
	  } else {
	    p->rightChild = temp;
	    p->balance--;
	  }	
	} else {
	  *root = temp;
	}
	temp->parent = p;
	child->parent = temp;
	n->parent = temp;
	if (child->rightChild != NULL) child->rightChild->parent = child;
	if (n->leftChild != NULL) n->leftChild->parent = n;
	temp->leftChild->balance = - max(temp->balance, 0); 
	temp->rightChild->balance = - min(temp->balance, 0);
	temp->balance = 0;
	n = p;
      }
    } else if (n->balance == 2) {
      printf("hello\n");

      if (n->leftChild == NULL) printf("no left child\n");
      if (n->rightChild == NULL) printf("no right child\n");

      printf("bye\n");
      printf("leftChild: %d\n", n->leftChild->value);
      printf("rightChild: %d\n", n->rightChild->value);


      // >>>>>>>>>>>>> simplify if by moving this case and checking for -1
      if (n->rightChild->balance == 1 || n->rightChild->balance == 0) {
	// single rotation on right child
	Node child = n->rightChild;
	int childBalance = child->balance;
	Node temp = child->leftChild;
	Node p = n->parent;
	child->leftChild = n;
	n->rightChild = temp;
	if (p != NULL) {
	  if (p->leftChild == n) {
	    p->leftChild = child;
	    if (child->balance == 1) {
	      p->balance++;
	    }
	  } else {
	    p->rightChild = child;
	    if (child->balance == 1) {
	      p->balance--;
	    }
	  }	
	} else {
	  *root = child;
	}
	child->parent = n->parent;
	if (temp != NULL) temp->parent = n;
	n->parent = child;
	child->balance--;
	n->balance = - child->balance;
	if (childBalance == 0) return;
	n = p;
      } else {
	// double rotation on right child
	Node child = n->rightChild;
	Node temp = child->leftChild;
	Node p = n->parent;
	child->leftChild = temp->rightChild;
	n->rightChild = temp->leftChild;
	temp->rightChild = child;
	temp->leftChild = n;
	if (p != NULL) {
	  if (p->leftChild == n) {
	    p->leftChild = temp;
	    p->balance++;
	  } else {
	    p->rightChild = temp;
	    p->balance--;
	  }	
	} else {
	  *root = temp;
	}
	temp->parent = p;
	child->parent = temp;
	n->parent = temp;
	if (child->leftChild != NULL) child->leftChild->parent = child;
	if (n->rightChild != NULL) n->rightChild->parent = n;
	temp->leftChild->balance = - max(temp->balance, 0);
	temp->rightChild->balance = - min(temp->balance, 0);
	temp->balance = 0;
	n = p;
      }
    } else {
      // No rotation took place on this run through the loop but one or more took place earlier
      child = n;
      n = n->parent;
      if (n != NULL) {
	if (n->rightChild != NULL && n->rightChild == child) {
	  n->balance--;
	} else {
	  n->balance++;
	}
      }
    }
  }
}




void printTree(Node n) {
	if (n == NULL) return;
	printf("%d %d %d\n",n->value, n->balance, (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->balance = 0;
	return n;
}