/*
  This file is part of the Astrometry.net suite.
  Copyright 2006, 2007 Dustin Lang, Keir Mierle and Sam Roweis.

  The Astrometry.net suite is free software; you can redistribute
  it and/or modify it under the terms of the GNU General Public License
  as published by the Free Software Foundation, version 2.

  The Astrometry.net suite is distributed in the hope that it will be
  useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with the Astrometry.net suite ; if not, write to the Free Software
  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301 USA
*/

#include <stdio.h>
#include <errno.h>
#include <string.h>
#include <assert.h>
#include <stdlib.h>

#include "bt.h"

/*
  The AVL tree portion of this code was adapted from GNU libavl.
*/
#define AVL_MAX_HEIGHT 32

// adapt compare_func to compare_func_2
static int compare_helper(const void* v1, const void* v2, void* token) {
	compare_func f = token;
	return f(v1, v2);
}

// data follows the bt_datablock*.
static Const void* NODE_DATA(bt_leaf* leaf) {
	return leaf+1;
}
static Const char* NODE_CHARDATA(bt_leaf* leaf) {
	return (char*)(leaf+1);
}

static Pure anbool isleaf(bt_node* node) {
	return node->leaf.isleaf;
}

static Pure int node_N(bt_node* node) {
	if (isleaf(node))
		return node->leaf.N;
	return node->branch.N;
}

static Pure int sum_childN(bt_branch* branch) {
	return node_N(branch->children[0]) + node_N(branch->children[1]);
}

static Pure bt_node* getchild(bt_node* node, int i) {
	if (isleaf(node)) return NULL;
	return node->branch.children[i];
}

static Pure bt_node* getleftchild(bt_node* node) {
	if (isleaf(node)) return NULL;
	return node->branch.children[0];
}

static Pure bt_node* getrightchild(bt_node* node) {
	if (isleaf(node)) return NULL;
	return node->branch.children[1];
}

static Pure bt_leaf* firstleaf(bt_node* node) {
	if (isleaf(node)) return &node->leaf;
	return node->branch.firstleaf;
}

static int height_slow(bt_node* node) {
	int hl, hr;
	if (isleaf(node)) return 1;
	hl = height_slow(node->branch.children[0]);
	hr = height_slow(node->branch.children[1]);
	return 1 + (hl > hr ? hl : hr);
}

#define CHECK(expr) {  \
int truthval = (expr); \
assert(truthval); \
if (!truthval) return -1; \
}

static int bt_check_node(bt* tree, bt_node* node) {
	int hl, hr;
	bt_node* leftmost;
	if (isleaf(node)) {
		CHECK(node->leaf.N <= tree->blocksize);
		return 0;
	}

	CHECK(sum_childN(&node->branch) == node->branch.N);

	leftmost = node;
	while (!isleaf(leftmost))
		leftmost = leftmost->branch.children[0];
	CHECK(&leftmost->leaf == node->branch.firstleaf);

	hl = height_slow(node->branch.children[0]);
	hr = height_slow(node->branch.children[1]);
	CHECK(node->branch.balance == (hr - hl));
	CHECK((node->branch.balance == 0) ||
          (node->branch.balance == 1) ||
          (node->branch.balance == -1));

    if (bt_check_node(tree, node->branch.children[0]) ||
        bt_check_node(tree, node->branch.children[1]))
        return -1;
    return 0;
}

int bt_check(bt* tree) {
	if (tree->root) {
		CHECK(node_N(tree->root) == tree->N);
        return bt_check_node(tree, tree->root);
	}
    return 0;
}

bt* bt_new(int datasize, int blocksize) {
	bt* tree = calloc(1, sizeof(bt));
	if (!tree) {
		fprintf(stderr, "Failed to allocate a new bt struct: %s\n", strerror(errno));
		return NULL;
	}
	tree->datasize = datasize;
	tree->blocksize = blocksize;
	return tree;
}

int bt_size(bt* tree) {
	return tree->N;
}

int bt_height(bt* tree) {
	bt_node* n;
	int h;
	n = tree->root;
	if (!n) return 0;
	if (isleaf(n)) return 1;
	for (h=1; !isleaf(n); h++) {
		if (n->branch.balance > 0)
			n = getrightchild(n);
		else
			n = getleftchild(n);
	}
	return h;
}

static int bt_count_leaves_rec(bt_node* node) {
	if (isleaf(node)) return 1;
	else return
			 bt_count_leaves_rec(node->branch.children[0]) +
			 bt_count_leaves_rec(node->branch.children[1]);
}

int bt_count_leaves(bt* tree) {
	return bt_count_leaves_rec(tree->root);
}

static void bt_free_node(bt_node* node) {
	if (!isleaf(node)) {
		bt_free_node(getleftchild(node));
		bt_free_node(getrightchild(node));
	}
	free(node);
}

void bt_free(bt* tree) {
	if (tree->root)
		bt_free_node(tree->root);
	free(tree);
}

static Pure void* get_element(bt* tree, bt_leaf* leaf, int index) {
	return NODE_CHARDATA(leaf) + index * tree->datasize;
}

static Pure void* first_element(bt_node* n) {
	if (isleaf(n)) return NODE_DATA(&n->leaf);
	else return NODE_DATA(n->branch.firstleaf);
}

static Malloc bt_node* bt_new_branch(bt* tree) {
	bt_node* n = calloc(1, sizeof(bt_node));
	if (!n) {
		fprintf(stderr, "Failed to allocate a new bt_node: %s\n", strerror(errno));
		return NULL;
	}
	return n;
}

static Malloc bt_node* bt_new_leaf(bt* tree) {
	bt_node* n = malloc(sizeof(bt_leaf) + tree->datasize * tree->blocksize);
	if (!n) {
		fprintf(stderr, "Failed to allocate a new bt_node: %s\n", strerror(errno));
		return NULL;
	}
	n->leaf.isleaf = 1;
	n->leaf.N = 0;
	return n;
}

static anbool bt_leaf_insert(bt* tree, bt_leaf* leaf, void* data, anbool unique,
						   compare_func_2 compare, void* token, void* overflow) {
	int lower, upper;
	int nshift;
	int index;

	// binary search...
	lower = -1;
	upper = leaf->N;
	while (lower < (upper-1)) {
		int mid;
		int cmp;
		mid = (upper + lower) / 2;
		cmp = compare(data, get_element(tree, leaf, mid), token);
		if (!cmp && unique) return FALSE;
		if (cmp >= 0)
			lower = mid;
		else
			upper = mid;
	}
	// index to insert at:
	index = lower + 1;

	// duplicate value?
	if (unique && (index > 0))
		if (compare(data, get_element(tree, leaf, index-1), token) == 0)
			return FALSE;

	// shift...
	nshift = leaf->N - index;
	if (leaf->N == tree->blocksize) {
		// this node is full.  insert the element and put the overflowing
		// element in "overflow".
		if (nshift) {
			memcpy(overflow, get_element(tree, leaf, leaf->N-1), tree->datasize);
			nshift--;
		} else {
			memcpy(overflow, data, tree->datasize);
			return TRUE;
		}
	} else {
		leaf->N++;
		tree->N++;
	}
	memmove(get_element(tree, leaf, index+1),
			get_element(tree, leaf, index),
			nshift * tree->datasize);
	// insert...
	memcpy(get_element(tree, leaf, index), data, tree->datasize);
	return TRUE;
}

static bt_node* next_node(bt_node** ancestors, int nancestors,
						  bt_node* child,
						  bt_node** nextancestors, int* nnextancestors) {
	// -first, find the first ancestor of whom we are a left
	//  (grand^n)-child.
	bt_node* parent = NULL;
	int i, j;
	for (i=nancestors-1; i>=0; i--) {
		parent = ancestors[i];
		if (parent->branch.children[0] == child)
			break;
		child = parent;
	}
	if (i < 0) {
		// no next node.
		return NULL;
	}

	// we share ancestors from the root to "parent".
	for (j=i; j>=0; j--)
		nextancestors[j] = ancestors[j];
	*nnextancestors = i+1;

	// -next, find the leftmost leaf of the parent's right branch.
	child = parent->branch.children[1];
	while (!isleaf(child)) {
		nextancestors[(*nnextancestors)++] = child;
	    child = child->branch.children[0];
	}
	return child;
}

// Pure?
static anbool bt_leaf_contains(bt* tree, bt_leaf* leaf, void* data,
							 compare_func_2 compare, void* token) {
	int lower, upper;
	lower = -1;
	upper = leaf->N;
	while (lower < (upper-1)) {
		int mid;
		int cmp;
		mid = (upper + lower) / 2;
		cmp = compare(data, get_element(tree, leaf, mid), token);
		if (cmp == 0) return TRUE;
		if (cmp > 0)
			lower = mid;
		else
			upper = mid;
	}
	// duplicate value?
	if (lower >= 0)
		if (compare(data, get_element(tree, leaf, lower), token) == 0)
			return TRUE;
	return FALSE;
}

// Pure?
anbool bt_contains(bt* tree, void* data, compare_func compare) {
	return bt_contains2(tree, data, compare_helper, compare);
}

anbool bt_contains2(bt* tree, void* data, compare_func_2 compare, void* token) {
	bt_node *n;
	int cmp;
	int dir;

	if (!tree->root)
		return FALSE;

	dir = 0;
	for (n = tree->root; !isleaf(n); n = getchild(n, dir)) {
		cmp = compare(data, first_element(n->branch.children[1]), token);
		if (cmp == 0)
			return TRUE;
		dir = (cmp > 0);
	}
	return bt_leaf_contains(tree, &n->leaf, data, compare, token);
}

static void update_firstleaf(bt_node** ancestors, int nancestors,
							 bt_node* child, bt_leaf* leaf) {
	int i;
	for (i=nancestors-1; i>=0; i--) {
		if (ancestors[i]->branch.children[0] != child)
			break;
		ancestors[i]->branch.firstleaf = leaf;
		child = ancestors[i];
	}
}

anbool bt_insert(bt* tree, void* data, anbool unique, compare_func compare) {
	return bt_insert2(tree, data, unique, compare_helper, compare);
}

anbool bt_insert2(bt* tree, void* data, anbool unique, compare_func_2 compare, void* token) {
	bt_node *y, *z; /* Top node to update balance factor, and parent. */
	bt_node *p, *q; /* Iterator, and parent. */
	bt_node *n;     /* Newly inserted node. */
	bt_node *w;     /* New root of rebalanced subtree. */
	int dir;                /* Direction to descend. */
	bt_node* np;

	bt_node* ancestors[AVL_MAX_HEIGHT];
	int nancestors = 0;
	unsigned char da[AVL_MAX_HEIGHT]; /* Cached comparison results. */
	int k = 0;              /* Number of cached results. */
	unsigned char overflow[tree->datasize];
	anbool rtn;
	anbool willfit;
	int cmp;
	int i;

	if (!tree->root) {
		// inserting the first element...
		n = bt_new_leaf(tree);
		tree->root = n;
		bt_leaf_insert(tree, &n->leaf, data, unique, compare, token, NULL);
		return TRUE;
	}

	z = y = tree->root;
	dir = 0;

	for (q = z, p = y;
		 !isleaf(p);
		 q = p, p = p->branch.children[dir]) {
		cmp = compare(data, first_element(p->branch.children[1]), token);
		if (unique && (cmp == 0))
			return FALSE;
		if (p->branch.balance != 0) {
			z = q;
			y = p;
			k = 0;
		}
		ancestors[nancestors++] = p;
		da[k++] = dir = (cmp > 0);
	}
	cmp = compare(data, first_element(p), token);

	// will this element fit in the current node?
	willfit = (p->leaf.N < tree->blocksize);
	if (willfit) {
		rtn = bt_leaf_insert(tree, &p->leaf, data, unique, compare, token, overflow);
		// duplicate value?
		if (!rtn)
			return rtn;
		for (i=0; i<nancestors; i++)
			ancestors[i]->branch.N++;
		return TRUE;
	}

	da[k++] = (cmp > 0);

	if (cmp > 0) {
		// insert the new element into this node and shuffle the
		// overflowing element (which may be the new element)
		// into the next node, if it exists, or a new node.
		bt_node* nextnode;
		bt_node* nextancestors[AVL_MAX_HEIGHT];
		int nnextancestors;

		/*
		  HACK - should we traverse the tree looking for the next node,
		  or just take the right sibling if we're the left child of a
		  balanced parent?
		*/

		rtn = bt_leaf_insert(tree, &p->leaf, data, unique, compare, token, overflow);
		if (!rtn)
			// duplicate value.
			return rtn;
		nextnode = next_node(ancestors, nancestors, p, nextancestors, &nnextancestors);
		assert(!nextnode || isleaf(nextnode));
		if (nextnode && (nextnode->leaf.N < tree->blocksize)) {
			// there's room; insert the element!
			rtn = bt_leaf_insert(tree, &nextnode->leaf, overflow, unique, compare, token, NULL);
			for (i=0; i<nnextancestors; i++)
				nextancestors[i]->branch.N++;
			return rtn;
		}

		// no room (or no next node); add a new node to the right to hold
		// the overflowed data.
		dir = 1;
		data = overflow;
	} else {
		// add a new node to the left.
		dir = 0;
	}

	// we have "q", the parent node, and "p", the existing leaf node which is
	// full.  we create a new branch node, "np", to take "p"'s place.
	// we move "p" down to be the child (1-dir) of "np", and create a new
	// leaf node, "n", to be child (dir) of "np".

	// create "n", a new leaf node to hold this element.
	n = bt_new_leaf(tree);
	if (!n)
		return FALSE;
	rtn = bt_leaf_insert(tree, &n->leaf, data, unique, compare, token, NULL);
	if (!rtn) {
		free(n);
		return FALSE;
	}

	// create "np", a new branch node to take p's place.
	np = bt_new_branch(tree);
	if (!np)
		return FALSE;
	np->branch.children[dir] = n;
	np->branch.children[1-dir] = p;
	np->branch.N = p->leaf.N + 1;
	np->branch.firstleaf = &(np->branch.children[0]->leaf);
    if (!isleaf(q)) {
        if (q->branch.children[0] == p) {
            q->branch.children[0] = np;
            if (!dir)
                update_firstleaf(ancestors, nancestors, np, np->branch.firstleaf);
        } else if (q->branch.children[1] == p) // need this because it could be that p = q = root.
            q->branch.children[1] = np;
    }

	if (p == tree->root)
		tree->root = np;

	for (i=0; i<nancestors; i++)
		ancestors[i]->branch.N++;

	if (!y || isleaf(y))
		return TRUE;

	for (p = y, k = 0;
		 p != np;
		 p = p->branch.children[da[k]], k++)
		if (da[k] == 0)
			p->branch.balance--;
		else
			p->branch.balance++;
	
	if (y->branch.balance == -2) {
		bt_node *x = y->branch.children[0];
		if (x->branch.balance == -1) {
			/*
			   y
			   |- x
			   |  |- x0
			   |  |- x1
			   |- y1

			   becomes

			   x
			   |- x0
			   |- y
			   .  |- x1
			   .  |- y1
			*/
			w = x;
			y->branch.children[0] = x->branch.children[1];
			x->branch.children[1] = y;
			x->branch.balance = y->branch.balance = 0;

			y->branch.firstleaf = firstleaf(y->branch.children[0]);
			y->branch.N = sum_childN(&y->branch);
			x->branch.N = sum_childN(&x->branch);

        } else {
			/*
			  y
			  |- x
			  |  |- x0
			  |  |- x1 (= w)
			  |     |- x10
			  |     |- x11
			  |- y1

			  becomes

			  x1
			  |- x
			  |  |- x0
			  |  |- x10
			  |- y
			  .  |- x11
			  .  |- y1
			*/

			assert (x->branch.balance == 1);
			w = x->branch.children[1];
			x->branch.children[1] = w->branch.children[0];
			w->branch.children[0] = x;
			y->branch.children[0] = w->branch.children[1];
			w->branch.children[1] = y;

			y->branch.firstleaf = firstleaf(y->branch.children[0]);
			w->branch.firstleaf = x->branch.firstleaf;

			y->branch.N = sum_childN(&y->branch);
			x->branch.N = sum_childN(&x->branch);
			w->branch.N = sum_childN(&w->branch);

			if (w->branch.balance == -1) {
				x->branch.balance = 0;
				y->branch.balance = 1;
			} else if (w->branch.balance == 0)
				x->branch.balance = y->branch.balance = 0;
			else {
				x->branch.balance = -1;
				y->branch.balance = 0;
			}
			w->branch.balance = 0;

        }
    } else if (y->branch.balance == 2) {
		bt_node *x = y->branch.children[1];
		if (x->branch.balance == 1) {
			/*
			  y
			  |- y0
			  |- x
			  .   |- x0
			  .   |- x1

			  becomes

			  x
			  |- y
			  |  |- y0
			  |  |- x0
			  |- x1
			*/
			w = x;
			y->branch.children[1] = x->branch.children[0];
			x->branch.children[0] = y;
			x->branch.balance = y->branch.balance = 0;

			x->branch.firstleaf = y->branch.firstleaf;

			y->branch.N = sum_childN(&y->branch);
			x->branch.N = sum_childN(&x->branch);


        } else {
			/*
			  y
			  |- y0
			  |- x
			  .  |- x0
			  .  |  |- x00
			  .  |  |- x01
			  .  |- x1

			  becomes

			  x0
			  |- y
			  |  |- y0
			  |  |- x00
			  |- x
			  .  |- x01
			  .  |- x1
			*/
			assert (x->branch.balance == -1);
			w = x->branch.children[0];
			x->branch.children[0] = w->branch.children[1];
			w->branch.children[1] = x;
			y->branch.children[1] = w->branch.children[0];
			w->branch.children[0] = y;

			y->branch.N = sum_childN(&y->branch);
			x->branch.N = sum_childN(&x->branch);
			w->branch.N = sum_childN(&w->branch);

			x->branch.firstleaf = firstleaf(x->branch.children[0]);
			w->branch.firstleaf = y->branch.firstleaf;

			if (w->branch.balance == 1) {
				x->branch.balance = 0;
				y->branch.balance = -1;
			} else if (w->branch.balance == 0)
				x->branch.balance = y->branch.balance = 0;
			else {
				x->branch.balance = 1;
				y->branch.balance = 0;
			}
			w->branch.balance = 0;

        }
    } else
		goto finished;

	if (y == tree->root)
		tree->root = w;
	else {
		if (y == z->branch.children[0]) {
			z->branch.children[0] = w;
			z->branch.firstleaf = w->branch.firstleaf;
		} else
			z->branch.children[1] = w;
	}

 finished:
	return TRUE;
}

// Pure?
void* bt_access(bt* tree, int index) {
	bt_node* n = tree->root;
	int offset = index;
	while (!isleaf(n)) {
		int Nleft = node_N(n->branch.children[0]);
		if (offset >= Nleft) {
			n = n->branch.children[1];
			offset -= Nleft;
		} else
			n = n->branch.children[0];
	}
	return get_element(tree, &n->leaf, offset);
}

static void bt_print_node(bt* tree, bt_node* node, char* indent,
						  void (*print_element)(void* val)) {
	printf("%s", indent);
	printf("(%p) ", node);
	printf("N=%i", node_N(node));

	if (!isleaf(node)) {
		char* subind;
		char* addind = "  ";
		printf(", leftmost (%p)", node->branch.firstleaf);
		printf(", Nleft=%i, Nright=%i, balance %i.\n",
			   node_N(node->branch.children[0]),
			   node_N(node->branch.children[1]),
			   node->branch.balance);
		subind = malloc(strlen(indent) + strlen(addind) + 1);
		sprintf(subind, "%s%s", indent, addind);
		printf("%sLeft child:\n", indent);
		bt_print_node(tree, node->branch.children[0], subind, print_element);
		printf("%sRight child:\n", indent);
		bt_print_node(tree, node->branch.children[1], subind, print_element);
		free(subind);
	} else {
		int i;
		printf(".  Leaf.");
		if (print_element) {
			printf("  [ ");
			for (i=0; i<node_N(node); i++)
				print_element(get_element(tree, &node->leaf, i));
		}
		printf("]\n");
	}
}

void bt_print(bt* tree, void (*print_element)(void* val)) {
	printf("  bt: N=%i, datasize=%i, blocksize=%i.\n", tree->N,
		   tree->datasize, tree->blocksize);
	printf("      root=(%p)\n", tree->root);
	if (!tree->root) {
		printf("  Empty tree.\n");
		return;
	}
	bt_print_node(tree, tree->root, "  ", print_element);
}

static void bt_print_struct_node(bt* tree, bt_node* node, char* indent,
								 void (*print_element)(void* val)) {

	printf("%s", indent);
	if (!isleaf(node)) {
		char* subind;
		char* addind = "|--";
		printf("(bal %i)\n", node->branch.balance);
		subind = malloc(strlen(indent) + strlen(addind) + 1);
		sprintf(subind, "%s%s", indent, addind);
		bt_print_struct_node(tree, node->branch.children[0], subind, print_element);
		bt_print_struct_node(tree, node->branch.children[1], subind, print_element);
	} else {
		int i;
		printf("(leaf)");
		if (print_element) {
			printf(" [ ");
			for (i=0; i<node_N(node); i++)
				print_element(get_element(tree, &node->leaf, i));
			printf("]");
		}
		printf("\n");
	}
}

void bt_print_structure(bt* tree, void (*print_element)(void* val)) {
	bt_print_struct_node(tree, tree->root, "   ", print_element);
}

