#include "avlNode.h"

AvlNode::AvlNode(int key) {
  this->left = nullptr;
  this->right = nullptr;
  this->key = key;
  this->height = 1;
}

// helper function which returns the max of the two inputs
int AvlNode::max(int n1, int n2) { return (n1 > n2) ? n1 : n2; }

// caluclate height difference of a given Node by calculating difference of
// heights of left and right subtree
int AvlNode::getHeightDifference() {
  int diff;
  if (this == nullptr) {
    diff = 0;
  } else {
    diff = this->left->getHeight() - this->right->getHeight();
  }
  return diff;
}

// returns height of a given AVL Node
int AvlNode::getHeight() {
  if (this == nullptr) {
    return 0;
  } else {
    return this->height;
  }
}

// insert a new node by recursively going to the sorted position which yields
// nullptr
AvlNode *AvlNode::insert(int key) {
  if (this == nullptr) {
    AvlNode *newNode = new AvlNode(key);
    return newNode;
  } else if (key > this->key) {
    this->right = this->right->insert(key);
  } else if (key < this->key) {
    this->left = this->left->insert(key);
  } else {
    return this;
  }
  // update heights of all other nodes
  this->height =
      1 + this->max(this->left->getHeight(), this->right->getHeight());

  int diff = this->getHeightDifference();
  // check if rebalancing is necessary
  if (diff > 1) { // check which kind of rebalancing is necessary
    if (key > this->left->key) {
      this->left = this->left->leftRotation();
      return this->rightRotation();

    } else if (key < this->left->key) {
      return this->rightRotation();
    }
  } else if (diff < -1) { // check which kind of rebalancing is necessary
    if (key > this->right->key) {
      return this->leftRotation();
    } else if (key < this->right->key) {
      this->right = this->right->rightRotation();
      return this->leftRotation();
    }
  }
  return this;
}

// perform a left rotation (see lecture)
AvlNode *AvlNode::leftRotation() {
  std::cout << "Do a left rotation on node " << this->key << "\n";
  AvlNode *rightNode = this->right;
  AvlNode *leftOfRightNode = rightNode->left;

  rightNode->left = this;
  this->right = leftOfRightNode;

  // update heights
  this->height =
      this->max(this->left->getHeight(), this->right->getHeight()) + 1;
  rightNode->height =
      this->max((rightNode->left->getHeight()), rightNode->right->getHeight()) +
      1;
  return rightNode;
}

// perform a right rotation (see lecture)
AvlNode *AvlNode::rightRotation() {
  std::cout << "Do a right rotation on node " << this->key << "\n";
  AvlNode *leftNode = this->left;
  AvlNode *rightOfLeftNode = leftNode->right;

  leftNode->right = this;
  this->left = rightOfLeftNode;

  // update heights
  this->height =
      this->max(this->left->getHeight(), this->right->getHeight()) + 1;
  leftNode->height =
      this->max(leftNode->left->getHeight(), leftNode->right->getHeight()) + 1;
  return leftNode;
}

// function to print an AVL tree in pre-order: (sub)root, left (sub)tree, right
// (sub)tree
std::string AvlNode::printPreorder() {
  std::stringstream output;
  output << "node:\t" << this->key << ",\t";
  output << "height:\t" << this->getHeight() << ",\t";
  output << "Diff:\t" << this->getHeightDifference() << "\n";

  if (this->left != nullptr) {
    output << this->left->printPreorder();
  }
  if (this->right != nullptr) {
    output << this->right->printPreorder();
  }
  return output.str();
}

AvlNode *AvlNode::deleteItem(int key) {
  AvlNode *node = this;

  if (node == nullptr) {
    return this;
  } else if (key < node->key) {
    node->left = node->left->deleteItem(key);
  } else if (key > node->key) {
    node->right = node->right->deleteItem(key);
  } else {
    // AvlNode *deleteNode = node;
    if (node->left == nullptr && node->right == nullptr) {
      delete node;
      node = nullptr;
      return nullptr;
    } else if (node->left == nullptr) { // only children in right subtree
      AvlNode *temp = node;
      node = node->right;
      delete temp;
    } else if (this->right == nullptr) { // only children in left subtree
      AvlNode *temp = node;
      node = node->left;
      delete temp;
    } else { // we have to keep the BST structure, here, we look for the minimum
             // in the right subtree (see lecture)
      AvlNode *temp = node->right;
      while (temp->left != nullptr) {
        temp = temp->left;
      }
      node->key = temp->key;
      node->right = node->right->deleteItem(temp->key);
    }
  }
  this->height = 1 + max(this->left->getHeight(), this->right->getHeight());
  int diff = this->getHeightDifference();

  if (diff > 1) {
    if (this->left->getHeightDifference() >= 0) {
      // Left Left Case
      return this->rightRotation();
    } else {
      // Left Right Case
      this->left = this->left->leftRotation();
      return this->rightRotation();
    }
  }
  if (diff < -1) {
    if (this->right->getHeightDifference() <= 0) {
      // Right Right Case
      return this->leftRotation();
    } else {
      // Right Left Case
      this->right = this->right->rightRotation();
      return this->leftRotation();
    }
  }
  return node;
}