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

#define TRUE 1
#define FALSE 0

#define RED 1
#define BLACK 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;
    (*root)->color = BLACK;
    return TRUE;
  }
	
  int returnValue = FALSE;
  Node n = *root;
  while (TRUE) {
    if (n->color == BLACK && n->leftChild != NULL && 
        n->leftChild->color == RED && n->rightChild != NULL && 
        n->rightChild->color == RED)
      {
	n->color = RED;
	n->leftChild->color = BLACK;
	n->rightChild->color = BLACK;
	if (n->parent != NULL && n->parent->color == RED) {
	  balance(root, n->parent, n);
	}
      }
    if (n->value > num) {
      if (n->leftChild == NULL) {
	n->leftChild = newNode;
	newNode->parent = n;
	if (n->color == RED) {
	  balance(root, n, n->leftChild);
	}
	returnValue = TRUE;
	break;
      } else {
	n = n->leftChild;
      }
    } else if (n->value < num) {
      if (n->rightChild == NULL) {
	n->rightChild = newNode;
	newNode->parent = n;
	if (n->color == RED) {
	  balance(root, n, n->rightChild);
	}
	returnValue = TRUE;
	break;
      } else {
	n = n->rightChild;
      }
    } else {
      break;
    }
  }
  (*root)->color = BLACK;
  return returnValue;
}

// p is RED and x is RED
  void balance(Node* root, Node p, Node x) {
    Node g = p->parent;
    Node ggp;
    Node temp;
    Node l;
    Node r;
    if (g == NULL) {
      printf("foobar!\n");
      exit(0);
    }

    if (g->leftChild == p) {
      if (p->leftChild == x) {
	// single rotation on left child
	temp = p->rightChild;
	g->leftChild = temp;
	p->rightChild = g;
	ggp = g->parent;
	if (g->parent != NULL) {
	  if (ggp->leftChild == g) {
	    ggp->leftChild = p;
	  } else {
	    ggp->rightChild = p;
	  }	
	} else {
	  *root = p;
	}
	p->parent = ggp;
	g->parent = p;
	if (temp != NULL) temp->parent = g;
	p->color = BLACK;
	g->color = RED;
      } else {
    	// double rotation on left child
	l = x->leftChild;
	r = x->rightChild;
	x->leftChild = p;
	x->rightChild = g;
	p->rightChild = l;
	g->leftChild = r;
	ggp = g->parent;
	if (g->parent != NULL) {
	  if (ggp->leftChild == g) {
	    ggp->leftChild = x;
	  } else {
	    ggp->rightChild = x;
	  }	
	} else {
	  *root = x;
	}
	x->parent = ggp;
	p->parent = x;
	g->parent = x;
	if (l != NULL) l->parent = p;
	if (r != NULL) r->parent = g;
	x->color = BLACK;
	g->color = RED;
      }
    } else {
      if (p->rightChild == x) {
	// single rotation on right child
	temp = p->leftChild;
	g->rightChild = temp;
	p->leftChild = g;
	ggp = g->parent;
	if (g->parent != NULL) {
	  if (ggp->leftChild == g) {
	    ggp->leftChild = p;
	  } else {
	    ggp->rightChild = p;
	  }	
	} else {
	  *root = p;
	}
	p->parent = ggp;
	g->parent = p;
	if (temp != NULL) temp->parent = g;
	p->color = BLACK;
	g->color = RED;
      } else {
	// double rotation on right child
	l = x->leftChild;
	r = x->rightChild;
	x->leftChild = g;
	x->rightChild = p;
	p->leftChild = r;
	g->rightChild = l;
	ggp = g->parent;
	if (g->parent != NULL) {
	  if (ggp->leftChild == g) {
	    ggp->leftChild = x;
	  } else {
	    ggp->rightChild = x;
	  }	
	} else {
	  *root = x;
	}
	x->parent = ggp;
	p->parent = x;
	g->parent = x;
	if (l != NULL) l->parent = g;
	if (r != NULL) r->parent = p;
	x->color = BLACK;
	g->color = RED;
      }
    }
  }


// 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;
    }
  }
}


// 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) {
  //  printf("remove: %d\n", num);
  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 == n->parent) {
      printf("remove: %d\n", n->value);
      exit(0);
    }    
  }
  if (n->parent != NULL) n->parent->balance += offset;
  balance2(root, n->parent);
  return TRUE;
}


void balance2(Node* root, Node n) {
  //  printf("balance on: %d\n", n->value);
  Node child = NULL;
  int rotation = FALSE;
  while (n != NULL) {
    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
	rotation = TRUE;
	Node child = n->leftChild;
	Node temp = child->rightChild;
	child->rightChild = n;
	n->leftChild = temp;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = child;
	    n->parent->balance++; //iffy
	  } else {
	    n->parent->rightChild = child;
	    n->parent->balance--;//iffy
	  }	
	} else {
	  *root = child;
	}
	child->parent = n->parent;
	if (temp != NULL) temp->parent = n;
	n->parent = child;
	child->balance++;
	n->balance = - child->balance;
	return;
      } else {
	// double rotation on left child
	rotation = TRUE;
	Node child = n->leftChild;
	Node temp = child->rightChild;
	child->rightChild = temp->leftChild;
	n->leftChild = temp->rightChild;
	temp->leftChild = child;
	temp->rightChild = n;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = temp;
	    n->parent->balance++;
	  } else {
	    n->parent->rightChild = temp;
	    n->parent->balance--;
	  }	
	} 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;
	child->balance = 0;
	n->balance = 0;
	temp->balance = 0;
	return;
      }
    } else if (n->balance == 2) {
      // >>>>>>>>>>>>> simplify if by moving this case and checking for -1
      if (n->rightChild->balance == 1 || n->rightChild->balance == 0) {
	// single rotation on right child
	rotation = TRUE;
	Node child = n->rightChild;
	Node temp = child->leftChild;
	child->leftChild = n;
	n->rightChild = temp;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = child;
	    n->parent->balance++;
	  } else {
	    n->parent->rightChild = child;
	    n->parent->balance--;
	  }	
	} else {
	  *root = child;
	}
	child->parent = n->parent;
	if (temp != NULL) temp->parent = n;
	n->parent = child;
	child->balance--;
	n->balance = - child->balance;
	return;
      } else {
	// double rotation on right child
	rotation = TRUE;
	Node child = n->rightChild;
	Node temp = child->leftChild;
	child->leftChild = temp->rightChild;
	n->rightChild = temp->leftChild;
	temp->rightChild = child;
	temp->leftChild = n;
	if (n->parent != NULL) {
	  if (n->parent->leftChild == n) {
	    n->parent->leftChild = temp;
	    n->parent->balance++;
	  } else {
	    n->parent->rightChild = temp;
	    n->parent->balance--;
	  }	
	} 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;
	child->balance = 0;
	n->balance = 0;
	temp->balance = 0;
	return;
      }
    }
    //    assert(n != n->parent);
    child = n;
    n = n->parent;

    // balance update needs to be fixed (in particular, if there was no rotation)
    if (n != NULL && rotation) {
      if (n->balance == 0) {
	(n->rightChild == child)?n->balance--:n->balance++;
	return;
      } else {
	(n->rightChild == child)?n->balance--:n->balance++;
      }
    }
  }
}


*/


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

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;
}

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->color = RED;
	return n;
}