#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/init.h>
#include <linux/list.h>
#include <linux/slab.h>

long long last_saved_time;

long long current_timestamp(void) {
	struct timespec ts;
	getnstimeofday(&ts); // get current time
	long long milliseconds = ts.tv_sec*1000LL + ts.tv_nsec/1000000; // calculate milliseconds
	return milliseconds;
}

/*
 * This structure represents a single node in a BST.
 */
struct bst_node {
	int val;
	struct bst_node* left;
	struct bst_node* right;
};

/*
 * This structure represents an entire BST.  Note that we only need a
 * reference to the root node of the tree.
 */
struct bst {
	struct bst_node* root;
};


struct bst* bst_create(void) {
	struct bst* bst = kmalloc(sizeof(struct bst), GFP_KERNEL);
	bst->root = NULL;
	return bst;
}


int bst_isempty(struct bst* bst) {
	return bst->root == NULL;
}

int _bst_subtree_min_val(struct bst_node* n) {
	/*
	 * The minimum value in any subtree is just the leftmost value.  Keep going
	 * left till we get there.
	 */
	while (n->left != NULL) {
		n = n->left;
	}
	return n->val;
}

/*
 * Helper function to remove a given value from a subtree of a BST rooted at
 * a specified node.  Operates recursively by figuring out whether val is in
 * the left or the right subtree of the specified node and performing the
 * remove operation on that subtree.
 *
 * Returns the potentially new root of the given subtree, modified to have
 * the specified value removed.
 */
struct bst_node* _bst_subtree_remove(int val, struct bst_node* n) {

	if (n == NULL) {

		/*
		 * If n is NULL, that means we've reached a leaf node without finding
		 * the value we wanted to remove.  The tree won't be modified.
		 */
		return NULL;

	} else if (val < n->val) {

		/*
		 * If val is less than n, remove val from n's left subtree and update
		 * n->left to point to the modified subtree (with val removed).  Return n,
		 * whose subtree itself has now been modified.
		 */
		n->left = _bst_subtree_remove(val, n->left);
		return n;

	} else if (val > n->val) {

		/*
		 * If val is greater than n, remove val from n's right subtree and update
		 * n->right to point to the modified subtree (with val removed).  Return n,
		 * whose subtree itself has now been modified.
		 */
		n->right = _bst_subtree_remove(val, n->right);
		return n;

	} else {

		/*
		 * If we've reached this point, we've found a node with value val.  We
		 * need to remove this node from the tree, and the way we do that will
		 * differ based on whether the node has 0, 1, or 2 children.
		 */
		if (n->left != NULL && n->right != NULL) {

			/*
			 * If n has 2 children, we replace the value at n with the value at n's
			 * in-order successor node, which is the minimum value in n's right
			 * subtree.  Then we recursively remove n's in-order successor node from
			 * the tree (specifically from n's right subtree).
			 */
			n->val = _bst_subtree_min_val(n->right);
			n->right = _bst_subtree_remove(n->val, n->right);
			return n;

		} else if (n->left != NULL) {

			/*
			 * If n has only a left child, we simply delete n by freeing it and
			 * returning the left child node so that it becomes the new child of
			 * n's parent via the recursion.
			 */
			struct bst_node* left_child = n->left;
			kfree(n);
			return left_child;

		} else if (n->right != NULL) {

			/*
			 * If n has only a right child, we simply delete n by freeing it and
			 * returning the right child node so that it becomes the new child of
			 * n's parent via the recursion.
			 */
			struct bst_node* right_child = n->right;
			kfree(n);
			return right_child;

		} else {

			/*
			 * Otherwise, n has no children, and we can simply free it and return
			 * NULL so that n's parent will lose n as a child via the recursion.
			 */
			kfree(n);
			return NULL;

		}

	}

}
void bst_remove(int val, struct bst* bst) {
	/*
	 * We remove val by using our subtree removal function starting with the
	 * subtree rooted at bst->root (i.e. the whole tree).
	 */
	bst->root = _bst_subtree_remove(val, bst->root);
}

void bst_free(struct bst* bst) {

	/*
	 * Assume that bst_remove() frees each node it removes and use it to free
	 * all of the nodes in the tree.
	 */
	while (!bst_isempty(bst)) {
		bst_remove(bst->root->val, bst);
	}
	kfree(bst);
	printk("BST IS FREE\n");
}


/*
 * Helper function to generate a single BST node containing a given value.
 */
struct bst_node* _bst_node_create(int val) {
	struct bst_node* n = kmalloc(sizeof(struct bst_node), GFP_KERNEL);
	n->val = val;
	n->left = n->right = NULL;
	return n;
}


/*
 * Helper function to insert a given value into a subtree of a BST rooted at
 * a given node.  Operates recursively by determining into which subtree (left
 * or right) under the given node the value should be inserted and performing
 * the insertion on that subtree.
 *
 * Returns the root of the given subtree, modified to contain a new node with
 * the specified value.
 */
struct bst_node* _bst_subtree_insert(int val, struct bst_node* n) {

	if (n == NULL) {

		/*
		 * If n is NULL, we know we've reached a place to insert val, so we
		 * create a new node holding val and return it.
		 */
		return _bst_node_create(val);

	} else if (val < n->val) {

		/*
		 * If val is less than the value at n, we insert val in n's left subtree
		 * (somewhere) and update n->left to point to the modified subtree (with
		 * val inserted).
		 */
		n->left = _bst_subtree_insert(val, n->left);

	} else {

		/*
		 * If val is greater than or equal to the value at n, we insert val in n's
		 * right subtree (somewhere) and update n->right to point to the modified
		 * subtree (with val inserted).
		 */
		n->right = _bst_subtree_insert(val, n->right);

	}

	/*
	 * For the else if and else conditions, the subtree rooted at n has already
	 * been modified (by setting n->left or n->right above), so we can just
	 * return n here.
	 */
	return n;

}


void bst_insert(int val, struct bst* bst) {
	/*
	 * We insert val by using our subtree insertion function starting with the
	 * subtree rooted at bst->root (i.e. the whole tree).
	 */
	bst->root = _bst_subtree_insert(val, bst->root);

}


void bst_search(int val, struct bst_node* node){
	if (!node)
		return 0;
	else if (val < node->val)
		bst_search(val, node->left);
	else if (val > node->val)
		bst_search(val, node->right);
	else
		return 0;
}

int bst_insert_search_delete(int n) {
	//	int num_of_elems[3] = {1000, 10000, 100000};
	int i;
	for (i=0; i < 1; i++){
		// initialize bst
		struct bst* bst = bst_create();

		// add items into bst
		last_saved_time = current_timestamp();
		int j;
		for (j=0; j < n; j++){
			bst_insert(j, bst);
		}
		printk("time in ms to insert %d elements: %lld\n", n, current_timestamp()-last_saved_time);

		// search items in bst
		last_saved_time = current_timestamp();
		int k;
		for (k=0; k < n; k++){
			bst_search(k, bst->root);
		}
		printk("time in ms to search %d elements: %lld\n", n, current_timestamp()-last_saved_time);

		// delete items from bst
		last_saved_time = current_timestamp();
		for (k=n; k >= 0; k--){
			bst_remove(k, bst);
		}
		printk("time in ms to delete %d elements: %lld\n", n, current_timestamp()-last_saved_time);
		bst_free(bst);
		printk("removed bst\n");
	}
}

int __init bst_module_init(void){
	printk("bst module init\n");
	bst_insert_search_delete(383);
	return 0;
}

void __exit bst_module_cleanup(void){
	printk("bst module exit\n");
}

module_init(bst_module_init);
module_exit(bst_module_cleanup);
MODULE_LICENSE("GPL");