// balance and height functions based on code from
// http://www.sanfoundry.com/cpp-program-implement-avl-trees/
// Finding elements based on code from
// http://www.cplusplus.com/forum/beginner/54835/

#include <cassert>
#include <iostream>
#include <string>
#include <queue>

#include "json.hpp"
#include "AVLCommands.h"

BSTNode::BSTNode(int key) :
	key_(key),
	parent_(std::weak_ptr<BSTNode>()),
	left_(nullptr),
	right_(nullptr),
	height_(0),
	balance_(0) {} // add a height_

BSTNode::BSTNode(int key, std::weak_ptr<BSTNode> parent) :
	key_(key),
	parent_(parent),
	left_(nullptr),
	right_(nullptr),
	height_(0),
	balance_(0) {}

bool BSTNode::IsLeaf() const {
	return left_ == nullptr && right_ == nullptr;
}

bool BSTNode::HasLeftChild() const {
	return left_ != nullptr;
}

bool BSTNode::HasRightChild() const {
	return right_ != nullptr;
}

void BSTNode::DeleteChild(std::shared_ptr<BSTNode> v) {
	if (left_ == v) {
		left_ = nullptr;
	} else if (right_ == v) {
		right_ = nullptr;
	} else {
		std::cerr << "BSTNode::DeleteChild Error: non-child passed as argument\n";
		exit(EXIT_FAILURE);
	}
}

void BSTNode::ReplaceChild(std::shared_ptr<BSTNode> v, std::shared_ptr<BSTNode> u) {
	if (left_ == u || right_ == u) {
		std::cerr << "BSTNode::ReplaceChild Error: child passed as replacement\n";
	}
	if (left_ == v) {
		left_ = u;
		u->parent_ = v->parent_;
	} else if (right_ == v) {
		right_ = u;
		u->parent_ = v->parent_;
	} else {
		std::cerr << "BSTNode::ReplaceChild Error: non-child passed as argument\n";
		exit(EXIT_FAILURE);
	}
}

AVLCommands::AVLCommands() : root_(nullptr), size_(0) {} // add height here?

int AVLCommands::max(int a, int b) {
	return (a > b)? a : b;
}

int AVLCommands::height(std::shared_ptr<BSTNode> currentNode) {
	if (root_ == nullptr) {
		return 0;
	}
	int leftHeight, rightHeight, heightMax;

	if (currentNode != nullptr) {
		if (currentNode->left_ != nullptr) {
			leftHeight = height(currentNode->left_);
		} else {
			leftHeight = -1;
		}

		if (currentNode->right_ != nullptr) {
			rightHeight = height(currentNode->right_);
		} else {
			rightHeight = -1;
		}

		heightMax = 1 + max(leftHeight, rightHeight);
		currentNode->height_ = heightMax;
		// update balance for each node
		currentNode->balance_ = leftHeight - rightHeight;
	}
  return currentNode->height_;
}

int AVLCommands::getBalance(std::shared_ptr<BSTNode> currentNode) {
	if (root_ == nullptr) {
		return 0;
	}

	int leftHeight, rightHeight, balanceFactor;
	if (currentNode != nullptr) {
		if (currentNode->left_ != nullptr) {
			leftHeight = height(currentNode->left_);
		} else {
			leftHeight = -1;
		}

		if (currentNode->right_ != nullptr) {
			rightHeight = height(currentNode->right_);
		} else {
			rightHeight = -1;
		}

		balanceFactor = leftHeight - rightHeight;
	}
  return balanceFactor;
}

std::shared_ptr<BSTNode> AVLCommands::rightRotate(std::shared_ptr<BSTNode> currentNode){
	std::cout << "Right Rotate!" << std::endl;
	std::shared_ptr<BSTNode> x = currentNode->left_;
	std::shared_ptr<BSTNode> T2 = x->right_;
	// Perform rotation
	x->right_ = currentNode;
	currentNode->left_ = T2;

	// Update heights
	currentNode->height_ = max(height(currentNode->left_), height(currentNode->right_))+1;
	x->height_ = max(height(x->left_), height(x->right_))+1;

	// Return new root
	return x;

	return currentNode;
}

std::shared_ptr<BSTNode> AVLCommands::leftRotate(std::shared_ptr<BSTNode> currentNode){
	std::cout << "Left Rotate!" << std::endl;
	std::shared_ptr<BSTNode> y = currentNode->right_;
  std::shared_ptr<BSTNode> T2 = y->left_;

  // Perform rotation
  y->left_ = currentNode;
  currentNode->right_ = T2;

  //  Update heights
  currentNode->height_ = max(height(currentNode->left_), height(currentNode->right_))+1;
  y->height_ = max(height(y->left_), height(y->right_))+1;

  // Return new root
  return y;
}

void AVLCommands::printTree(){
	printTree(root_);
}

void AVLCommands::printTree(std::shared_ptr<BSTNode> currentNode){
	if (root_ == nullptr) {
		std::cout << "tree is empty" << std::endl;
		return;
	} else {
		std::cout << "key: " << currentNode->key_ << std::endl;
		std::cout << "height: " << currentNode->height_ << std::endl;
		std::cout << "balance: " << currentNode->balance_ << std::endl;
		if (currentNode->left_ != nullptr) {
			std::cout << "left" << std::endl;
			printTree(currentNode->left_);
		}
		if (currentNode->right_ != nullptr) {
			std::cout << "right" << std::endl;
			printTree(currentNode->right_);
		}
	}
}

void AVLCommands::Insert(int key) {
	std::cout << "Inserting new key: " << key << std::endl;
	// BST insertion
	if (root_ == nullptr) {
		root_ = std::make_shared<BSTNode>(key);
		size_++;
		return;
	}
	std::shared_ptr<BSTNode> currentNode = root_, lastNode = nullptr;
	while (currentNode != nullptr) {
		lastNode = currentNode;
		currentNode = (key < currentNode->key_) ?
			currentNode->left_ : currentNode->right_;
	}
	if (key < lastNode->key_) {
		lastNode->left_ = std::make_shared<BSTNode>(key, lastNode);
	} else {
		lastNode->right_ = std::make_shared<BSTNode>(key, lastNode);
	}
	size_++;

	// update height and balance of the tree
	height(root_);

	std::cout << "Rotation to be preformed:" << std::endl;

	// Check if unbalanced
	// Left Left Case
  if (lastNode->balance_ > 1 && key < lastNode->left_->key_) {
		std::cout << "Left-Left" << std::endl;
		rightRotate(lastNode);
	}

  // Right Right Case
  if (lastNode->balance_ < -1 && key > lastNode->right_->key_) {
		std::cout << "Right-Right" << std::endl;
    leftRotate(lastNode);
	}

  // Left Right Case
  if (lastNode->balance_ > 1 && key > lastNode->left_->key_) {
		std::cout << "Left-Right" << std::endl;
    lastNode->left_ =  leftRotate(lastNode->left_);
    rightRotate(lastNode);
  }

  // Right Left Case
  if (lastNode->balance_ < -1 && key < lastNode->right_->key_) {
		std::cout << "Right-Left" << std::endl;
    lastNode->right_ = rightRotate(lastNode->right_);
    leftRotate(lastNode);
  }

}

bool AVLCommands::Delete(int key) {
	// base BST delete
	std::shared_ptr<BSTNode> currentNode = root_;
	while (currentNode != nullptr) {
		if (currentNode->key_ == key) {
			if (currentNode->IsLeaf()) {
				DeleteLeaf(currentNode);
			} else if (currentNode->left_ == nullptr) {
				assert(currentNode->right_ != nullptr);
				std::shared_ptr<BSTNode> parent = currentNode->parent_.lock();
				parent->ReplaceChild(currentNode, currentNode->right_);
				size_--; assert(size_ >= 0);
			} else if (currentNode->right_ == nullptr) {
				assert(currentNode->left_ != nullptr);
				std::shared_ptr<BSTNode> parent = currentNode->parent_.lock();
				parent->ReplaceChild(currentNode, currentNode->left_);
				size_--; assert(size_ >= 0);
			} else {
				currentNode->key_ = DeleteMin(currentNode);
			}
		}
		currentNode = (key < currentNode->key_) ?
			currentNode->left_ : currentNode->right_;
	}

	std::cout << "current node value: " << currentNode->key_ << std::endl;
	// update height of current node
	std::cout << "delete height update attempt..." << std::endl;
	currentNode->height_ = 1 + max(height(currentNode->left_), height(currentNode->right_));
	std::cout << "New height: " << currentNode->height_ << std::endl;
	std::cout << "delete height update attempt success" << std::endl;

	// Update balance factor of this ancestor's node
	std::cout << "delete balance update attempt..." << std::endl;
	int currentNodeBalance = getBalance(currentNode);
	std::cout << "delete balance update attempt success!" << std::endl;

	// Check if unbalanced
	// Left Left Case
  if (currentNodeBalance > 1 && getBalance(currentNode->left_) >= 0)
      rightRotate(currentNode);

  // Right Right Case
  if (currentNodeBalance < -1 && getBalance(currentNode->right_) <= 0)
      leftRotate(currentNode);

  // Left Right Case
  if (currentNodeBalance > 1 && getBalance(currentNode->left_) < 0) {
      currentNode->left_ =  leftRotate(currentNode->left_);
      rightRotate(currentNode);
  }

  // Right Left Case
  if (currentNodeBalance < -1 && getBalance(currentNode->right_) > 0) {
      currentNode->right_ = rightRotate(currentNode->right_);
      leftRotate(currentNode);
  }
	return false;
}


int AVLCommands::DeleteMin() {
	return DeleteMin(root_);
}


void AVLCommands::DeleteLeaf(std::shared_ptr<BSTNode> currentNode) {
	std::shared_ptr<BSTNode> parent = currentNode->parent_.lock();
	if (parent == nullptr) {
		// Delete root
		root_ = nullptr;
		size_--; assert(size_ == 0);
	} else {
		if (parent->right_ == currentNode) {
			parent->right_ = nullptr;
		} else if (parent->left_ == currentNode) {
			parent->left_ = nullptr;
		} else {
			std::cerr << "BST::DeleteLeaf Error: inconsistent state\n";
		}
		size_--; assert(size_ >= 0);
	}
}

int AVLCommands::DeleteMin(std::shared_ptr<BSTNode> currentNode) {
	std::shared_ptr<BSTNode> lastNode = nullptr;
	while (currentNode != nullptr) {
		lastNode = currentNode;
		currentNode = currentNode->left_;
	}
	int result = lastNode->key_;
	std::shared_ptr<BSTNode> parent = lastNode->parent_.lock();
	if (parent == nullptr) {
		// lastNode is root
		if (lastNode->right_ != nullptr) {
			root_ = lastNode->right_;
			lastNode->right_->parent_.reset();
		} else {
			root_ = nullptr;
		}
	} else {
		// lastNode under the root
		if (lastNode->right_ != nullptr) {
			parent->left_ = lastNode->right_;
			lastNode->right_->parent_ = parent;
		} else {
			parent->left_ = nullptr;
		}
  }
	size_--; assert(size_ >= 0);
	return result;
}

size_t AVLCommands::size() const {
	return size_;
}

bool AVLCommands::empty() const {
	return size_ == 0;
}

bool AVLCommands::Find(int key) const {
	std::shared_ptr<BSTNode> currentNode = root_;
	while (currentNode != nullptr) {
		if (currentNode->key_ == key) {
			return true;
		}
		currentNode = (key < currentNode->key_) ?
			currentNode->left_ : currentNode->right_;
	}
	return false;
}

std::string AVLCommands::JSON() const {
	nlohmann::json result;
	std::queue< std::shared_ptr<BSTNode> > nodes;
	if (root_ != nullptr) {
		result["root"] = root_->key_;
		nodes.push(root_);
		while (!nodes.empty()) {
			auto v = nodes.front();
			nodes.pop();
			std::string key = std::to_string(v->key_);
			//std::string height = std::to_string(v->height_);
			//result[key]["height"] = height;
			result[key]["left"] = v->left_->key_;
			if (v->left_ != nullptr) {
				result[key]["left"] = v->left_->key_;
				nodes.push(v->left_);
			}
			if (v->right_ != nullptr) {
				result[key]["right"] = v->right_->key_;
				nodes.push(v->right_);
			}
			if (v->parent_.lock() != nullptr) {
				result[key]["parent"] = v->parent_.lock()->key_;
			} else {
				result[key]["root"] = true;
			}
		}
	}
	result["size"] = size_;
	return result.dump(2) + "\n";
}