/*
 # This file is part of the Astrometry.net suite.
 # Licensed under a 3-clause BSD style license - see LICENSE

 The AVL tree portion of this code was adapted from GNU libavl,
 licensed under the GPL.
 */

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

