/*!
 * \file kdtree.c
 *
 * \brief binary search tree
 *
 * Dynamic balanced k-d tree implementation
 *
 * (C) 2014 by the GRASS Development Team
 *
 * This program is free software under the GNU General Public License
 * (>=v2).  Read the file COPYING that comes with GRASS for details.
 *
 * \author Markus Metz
 */

#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <grass/gis.h>
#include <grass/glocale.h>
#include "kdtree.h"

#define KD_BTOL 7

#ifdef KD_DEBUG
#undef KD_DEBUG
#endif

static int rcalls = 0;
static int rcallsmax = 0;

static struct kdnode *kdtree_insert2(struct kdtree *, struct kdnode *,
                                     struct kdnode *, int, int);
static int kdtree_replace(struct kdtree *, struct kdnode *);
static int kdtree_balance(struct kdtree *, struct kdnode *, int);
static int kdtree_first(struct kdtrav *, double *, int *);
static int kdtree_next(struct kdtrav *, double *, int *);

static int cmp(struct kdnode *a, struct kdnode *b, int p)
{
    if (a->c[p] < b->c[p])
        return -1;
    if (a->c[p] > b->c[p])
        return 1;

    return (a->uid < b->uid ? -1 : a->uid > b->uid);
}

static int cmpc(struct kdnode *a, struct kdnode *b, struct kdtree *t)
{
    int i;

    for (i = 0; i < t->ndims; i++) {
        if (a->c[i] != b->c[i]) {
            return 1;
        }
    }

    return 0;
}

static struct kdnode *kdtree_newnode(struct kdtree *t)
{
    struct kdnode *n = G_malloc(sizeof(struct kdnode));

    n->c = G_malloc(t->ndims * sizeof(double));
    n->dim = 0;
    n->depth = 0;
    n->balance = 0;
    n->uid = 0;
    n->child[0] = NULL;
    n->child[1] = NULL;

    return n;
}

static void kdtree_free_node(struct kdnode *n)
{
    G_free(n->c);
    G_free(n);
}

static void kdtree_update_node(struct kdtree *t, struct kdnode *n)
{
    int ld, rd, btol;

    ld = (!n->child[0] ? -1 : n->child[0]->depth);
    rd = (!n->child[1] ? -1 : n->child[1]->depth);
    n->depth = MAX(ld, rd) + 1;

    n->balance = 0;
    /* set balance flag if any of the node's subtrees needs balancing
     * or if the node itself needs balancing */
    if ((n->child[0] && n->child[0]->balance) ||
        (n->child[1] && n->child[1]->balance)) {
        n->balance = 1;

        return;
    }

    btol = t->btol;
    if (!n->child[0] || !n->child[1])
        btol = 2;

    if (ld > rd + btol || rd > ld + btol)
        n->balance = 1;
}

/* create a new k-d tree with ndims dimensions,
 * optionally set balancing tolerance */
struct kdtree *kdtree_create(char ndims, int *btol)
{
    int i;
    struct kdtree *t;

    t = G_malloc(sizeof(struct kdtree));

    t->ndims = ndims;
    t->csize = ndims * sizeof(double);
    t->btol = KD_BTOL;
    if (btol) {
        t->btol = *btol;
        if (t->btol < 2)
            t->btol = 2;
    }

    t->nextdim = G_malloc(ndims * sizeof(char));
    for (i = 0; i < ndims - 1; i++)
        t->nextdim[i] = i + 1;
    t->nextdim[ndims - 1] = 0;

    t->count = 0;
    t->root = NULL;

    return t;
}

/* clear the tree, removing all entries */
void kdtree_clear(struct kdtree *t)
{
    struct kdnode *it;
    struct kdnode *save = t->root;

    /*
       Rotate away the left links so that
       we can treat this like the destruction
       of a linked list
     */
    while ((it = save) != NULL) {
        if (it->child[0] == NULL) {
            /* No left links, just kill the node and move on */
            save = it->child[1];
            kdtree_free_node(it);
            it = NULL;
        }
        else {
            /* Rotate away the left link and check again */
            save = it->child[0];
            it->child[0] = save->child[1];
            save->child[1] = it;
        }
    }
    t->root = NULL;
}

/* destroy the tree */
void kdtree_destroy(struct kdtree *t)
{
    /* remove all entries */
    kdtree_clear(t);
    G_free(t->nextdim);

    G_free(t);
    t = NULL;
}

/* insert an item (coordinates c and uid) into the k-d tree
 * dc == 1: allow duplicate coordinates */
int kdtree_insert(struct kdtree *t, double *c, int uid, int dc)
{
    struct kdnode *nnew;
    size_t count = t->count;

    nnew = kdtree_newnode(t);
    memcpy(nnew->c, c, t->csize);
    nnew->uid = uid;

    t->root = kdtree_insert2(t, t->root, nnew, 1, dc);

    /* print depth of recursion
     * recursively called fns are insert2, balance, and replace */
    /*
       if (rcallsmax > 1)
       fprintf(stdout, "%d\n", rcallsmax);
     */

    return count < t->count;
}

/* remove an item from the k-d tree
 * coordinates c and uid must match */
int kdtree_remove(struct kdtree *t, double *c, int uid)
{
    struct kdnode sn, *n;
    struct kdstack {
        struct kdnode *n;
        int dir;
    } s[256];
    int top;
    int dir, found;
    int balance, bmode;

    sn.c = c;
    sn.uid = uid;

    /* find sn node */
    top = 0;
    s[top].n = t->root;
    dir = 1;
    found = 0;
    while (!found) {
        n = s[top].n;
        found = (!cmpc(&sn, n, t) && sn.uid == n->uid);
        if (!found) {
            dir = cmp(&sn, n, n->dim) > 0;
            s[top].dir = dir;
            top++;
            s[top].n = n->child[dir];

            if (!s[top].n) {
                G_warning("Node does not exist");

                return 0;
            }
        }
    }

    if (s[top].n->depth == 0) {
        kdtree_free_node(s[top].n);
        s[top].n = NULL;
        if (top) {
            top--;
            n = s[top].n;
            dir = s[top].dir;
            n->child[dir] = NULL;

            /* update node */
            kdtree_update_node(t, n);
        }
        else {
            t->root = NULL;

            return 1;
        }
    }
    else
        kdtree_replace(t, s[top].n);

    while (top) {
        top--;
        n = s[top].n;

        /* update node */
        kdtree_update_node(t, n);
    }

    balance = 1;
    bmode = 1;
    if (balance) {
        struct kdnode *r;
        int iter, bmode2;

        /* fix any inconsistencies in the (sub-)tree */
        iter = 0;
        bmode2 = 0;
        top = 0;
        r = t->root;
        s[top].n = r;
        while (top >= 0) {

            n = s[top].n;

            /* top-down balancing
             * slower but more compact */
            if (!bmode2) {
                while (kdtree_balance(t, n, bmode))
                    ;
            }

            /* go down */
            if (n->child[0] && n->child[0]->balance) {
                dir = 0;
                top++;
                s[top].n = n->child[dir];
            }
            else if (n->child[1] && n->child[1]->balance) {
                dir = 1;
                top++;
                s[top].n = n->child[dir];
            }
            /* go back up */
            else {

                /* bottom-up balancing
                 * faster but less compact */
                kdtree_update_node(t, n);
                if (bmode2) {
                    while (kdtree_balance(t, n, bmode))
                        ;
                }
                top--;
                if (top >= 0) {
                    kdtree_update_node(t, s[top].n);
                }
                if (!bmode2 && top == 0) {
                    iter++;
                    if (iter == 2) {
                        /* the top node has been visited twice,
                         * switch from top-down to bottom-up balancing */
                        iter = 0;
                        bmode2 = 1;
                    }
                }
            }
        }
    }

    return 1;
}

/* k-d tree optimization, only useful if the tree will be used heavily
 * (more searches than items in the tree)
 * level 0 = a bit, 1 = more, 2 = a lot */
void kdtree_optimize(struct kdtree *t, int level)
{
    struct kdnode *n, *n2;
    struct kdstack {
        struct kdnode *n;
        int dir;
        char v;
    } s[256];
    int dir;
    int top;
    int ld, rd;
    int diffl, diffr;
    int nbal;

    if (!t->root)
        return;

    G_debug(1, "k-d tree optimization for %zd items, tree depth %d", t->count,
            t->root->depth);

    nbal = 0;
    top = 0;
    s[top].n = t->root;
    while (s[top].n) {
        n = s[top].n;

        ld = (!n->child[0] ? -1 : n->child[0]->depth);
        rd = (!n->child[1] ? -1 : n->child[1]->depth);

        if (ld < rd)
            while (kdtree_balance(t, n->child[0], level))
                ;
        else if (ld > rd)
            while (kdtree_balance(t, n->child[1], level))
                ;

        ld = (!n->child[0] ? -1 : n->child[0]->depth);
        rd = (!n->child[1] ? -1 : n->child[1]->depth);
        n->depth = MAX(ld, rd) + 1;

        dir = (rd > ld);

        top++;
        s[top].n = n->child[dir];
    }

    while (top) {
        top--;
        n = s[top].n;

        /* balance node */
        while (kdtree_balance(t, n, level)) {
            nbal++;
        }
        while (kdtree_balance(t, n->child[0], level))
            ;
        while (kdtree_balance(t, n->child[1], level))
            ;

        ld = (!n->child[0] ? -1 : n->child[0]->depth);
        rd = (!n->child[1] ? -1 : n->child[1]->depth);
        n->depth = MAX(ld, rd) + 1;

        while (kdtree_balance(t, n, level)) {
            nbal++;
        }
    }

    while (s[top].n) {
        n = s[top].n;

        /* balance node */
        while (kdtree_balance(t, n, level)) {
            nbal++;
        }
        while (kdtree_balance(t, n->child[0], level))
            ;
        while (kdtree_balance(t, n->child[1], level))
            ;

        ld = (!n->child[0] ? -1 : n->child[0]->depth);
        rd = (!n->child[1] ? -1 : n->child[1]->depth);
        n->depth = MAX(ld, rd) + 1;

        while (kdtree_balance(t, n, level)) {
            nbal++;
        }

        ld = (!n->child[0] ? -1 : n->child[0]->depth);
        rd = (!n->child[1] ? -1 : n->child[1]->depth);

        dir = (rd > ld);

        top++;
        s[top].n = n->child[dir];
    }

    while (top) {
        top--;
        n = s[top].n;

        /* update node depth */
        ld = (!n->child[0] ? -1 : n->child[0]->depth);
        rd = (!n->child[1] ? -1 : n->child[1]->depth);
        n->depth = MAX(ld, rd) + 1;
    }

    if (level) {
        top = 0;
        s[top].n = t->root;
        while (s[top].n) {
            n = s[top].n;

            /* balance node */
            while (kdtree_balance(t, n, level)) {
                nbal++;
            }
            while (kdtree_balance(t, n->child[0], level))
                ;
            while (kdtree_balance(t, n->child[1], level))
                ;

            ld = (!n->child[0] ? -1 : n->child[0]->depth);
            rd = (!n->child[1] ? -1 : n->child[1]->depth);
            n->depth = MAX(ld, rd) + 1;

            while (kdtree_balance(t, n, level)) {
                nbal++;
            }

            diffl = diffr = -1;
            if (n->child[0]) {
                n2 = n->child[0];
                ld = (!n2->child[0] ? -1 : n2->child[0]->depth);
                rd = (!n2->child[1] ? -1 : n2->child[1]->depth);

                diffl = ld - rd;
                if (diffl < 0)
                    diffl = -diffl;
            }
            if (n->child[1]) {
                n2 = n->child[1];
                ld = (!n2->child[0] ? -1 : n2->child[0]->depth);
                rd = (!n2->child[1] ? -1 : n2->child[1]->depth);

                diffr = ld - rd;
                if (diffr < 0)
                    diffr = -diffr;
            }

            dir = (diffr > diffl);

            top++;
            s[top].n = n->child[dir];
        }

        while (top) {
            top--;
            n = s[top].n;

            /* update node depth */
            ld = (!n->child[0] ? -1 : n->child[0]->depth);
            rd = (!n->child[1] ? -1 : n->child[1]->depth);
            n->depth = MAX(ld, rd) + 1;
        }
    }

    G_debug(1, "k-d tree optimization: %d times balanced, new depth %d", nbal,
            t->root->depth);

    return;
}

/* find k nearest neighbors
 * results are stored in uid (uids) and d (squared distances)
 * optionally an uid to be skipped can be given
 * useful when searching for the nearest neighbors of an item
 * that is also in the tree */
int kdtree_knn(struct kdtree *t, double *c, int *uid, double *d, int k,
               int *skip)
{
    int i, found;
    double diff, dist, maxdist;
    struct kdnode sn, *n;
    struct kdstack {
        struct kdnode *n;
        int dir;
        char v;
    } s[256];
    int dir;
    int top;

    if (!t->root)
        return 0;

    sn.c = c;
    sn.uid = (int)0x80000000;
    if (skip)
        sn.uid = *skip;

    maxdist = INFINITY;
    found = 0;

    /* go down */
    top = 0;
    s[top].n = t->root;
    while (s[top].n) {
        n = s[top].n;
        dir = cmp(&sn, n, n->dim) > 0;
        s[top].dir = dir;
        s[top].v = 0;
        top++;
        s[top].n = n->child[dir];
    }

    /* go back up */
    while (top) {
        top--;

        if (!s[top].v) {
            s[top].v = 1;
            n = s[top].n;

            if (n->uid != sn.uid) {
                if (found < k) {
                    dist = 0.0;
                    i = t->ndims - 1;
                    do {
                        diff = sn.c[i] - n->c[i];
                        dist += diff * diff;

                    } while (i--);

                    i = found;
                    while (i > 0 && d[i - 1] > dist) {
                        d[i] = d[i - 1];
                        uid[i] = uid[i - 1];
                        i--;
                    }
                    if (i < found && d[i] == dist && uid[i] == n->uid)
                        G_fatal_error("knn: inserting duplicate");
                    d[i] = dist;
                    uid[i] = n->uid;
                    maxdist = d[found];
                    found++;
                }
                else {
                    dist = 0.0;
                    i = t->ndims - 1;
                    do {
                        diff = sn.c[i] - n->c[i];
                        dist += diff * diff;

                    } while (i-- && dist <= maxdist);

                    if (dist < maxdist) {
                        i = k - 1;
                        while (i > 0 && d[i - 1] > dist) {
                            d[i] = d[i - 1];
                            uid[i] = uid[i - 1];
                            i--;
                        }
                        if (d[i] == dist && uid[i] == n->uid)
                            G_fatal_error("knn: inserting duplicate");
                        d[i] = dist;
                        uid[i] = n->uid;

                        maxdist = d[k - 1];
                    }
                }
                if (found == k && maxdist == 0.0)
                    break;
            }

            /* look on the other side ? */
            dir = s[top].dir;
            diff = sn.c[(int)n->dim] - n->c[(int)n->dim];
            dist = diff * diff;

            if (dist <= maxdist) {
                /* go down the other side */
                top++;
                s[top].n = n->child[!dir];
                while (s[top].n) {
                    n = s[top].n;
                    dir = cmp(&sn, n, n->dim) > 0;
                    s[top].dir = dir;
                    s[top].v = 0;
                    top++;
                    s[top].n = n->child[dir];
                }
            }
        }
    }

    return found;
}

/* find all nearest neighbors within distance aka radius search
 * results are stored in puid (uids) and pd (squared distances)
 * memory is allocated as needed, the calling fn must free the memory
 * optionally an uid to be skipped can be given */
int kdtree_dnn(struct kdtree *t, double *c, int **puid, double **pd,
               double maxdist, int *skip)
{
    int i, k, found;
    double diff, dist;
    struct kdnode sn, *n;
    struct kdstack {
        struct kdnode *n;
        int dir;
        char v;
    } s[256];
    int dir;
    int top;
    int *uid;
    double *d, maxdistsq;

    if (!t->root)
        return 0;

    sn.c = c;
    sn.uid = (int)0x80000000;
    if (skip)
        sn.uid = *skip;

    *pd = NULL;
    *puid = NULL;

    k = 0;
    uid = NULL;
    d = NULL;

    found = 0;
    maxdistsq = maxdist * maxdist;

    /* go down */
    top = 0;
    s[top].n = t->root;
    while (s[top].n) {
        n = s[top].n;
        dir = cmp(&sn, n, n->dim) > 0;
        s[top].dir = dir;
        s[top].v = 0;
        top++;
        s[top].n = n->child[dir];
    }

    /* go back up */
    while (top) {
        top--;

        if (!s[top].v) {
            s[top].v = 1;
            n = s[top].n;

            if (n->uid != sn.uid) {
                dist = 0;
                i = t->ndims - 1;
                do {
                    diff = sn.c[i] - n->c[i];
                    dist += diff * diff;

                } while (i-- && dist <= maxdistsq);

                if (dist <= maxdistsq) {
                    if (found + 1 >= k) {
                        k = found + 10;
                        uid = G_realloc(uid, k * sizeof(int));
                        d = G_realloc(d, k * sizeof(double));
                    }
                    i = found;
                    while (i > 0 && d[i - 1] > dist) {
                        d[i] = d[i - 1];
                        uid[i] = uid[i - 1];
                        i--;
                    }
                    if (i < found && d[i] == dist && uid[i] == n->uid)
                        G_fatal_error("dnn: inserting duplicate");
                    d[i] = dist;
                    uid[i] = n->uid;
                    found++;
                }
            }

            /* look on the other side ? */
            dir = s[top].dir;

            diff = fabs(sn.c[(int)n->dim] - n->c[(int)n->dim]);
            if (diff <= maxdist) {
                /* go down the other side */
                top++;
                s[top].n = n->child[!dir];
                while (s[top].n) {
                    n = s[top].n;
                    dir = cmp(&sn, n, n->dim) > 0;
                    s[top].dir = dir;
                    s[top].v = 0;
                    top++;
                    s[top].n = n->child[dir];
                }
            }
        }
    }

    *pd = d;
    *puid = uid;

    return found;
}

/* find all nearest neighbors within range aka box search
 * the range is specified with min and max for each dimension as
 * (min1, min2, ..., minn, max1, max2, ..., maxn)
 * results are stored in puid (uids)
 * memory is allocated as needed, the calling fn must free the memory
 * optionally an uid to be skipped can be given */
int kdtree_rnn(struct kdtree *t, double *c, int **puid, int *skip)
{
    int i, k, found, inside;
    struct kdnode sn, *n;
    struct kdstack {
        struct kdnode *n;
        int dir;
        char v;
    } s[256];
    int dir;
    int top;
    int *uid;

    if (!t->root)
        return 0;

    sn.c = c;
    sn.uid = (int)0x80000000;
    if (skip)
        sn.uid = *skip;

    *puid = NULL;

    k = 0;
    uid = NULL;

    found = 0;

    /* go down */
    top = 0;
    s[top].n = t->root;
    while (s[top].n) {
        n = s[top].n;
        dir = cmp(&sn, n, n->dim) > 0;
        s[top].dir = dir;
        s[top].v = 0;
        top++;
        s[top].n = n->child[dir];
    }

    /* go back up */
    while (top) {
        top--;

        if (!s[top].v) {
            s[top].v = 1;
            n = s[top].n;

            if (n->uid != sn.uid) {
                inside = 1;
                for (i = 0; i < t->ndims; i++) {
                    if (n->c[i] < sn.c[i] || n->c[i] > sn.c[i + t->ndims]) {
                        inside = 0;
                        break;
                    }
                }

                if (inside) {
                    if (found + 1 >= k) {
                        k = found + 10;
                        uid = G_realloc(uid, k * sizeof(int));
                    }
                    i = found;
                    uid[i] = n->uid;
                    found++;
                }
            }

            /* look on the other side ? */
            dir = s[top].dir;
            if (n->c[(int)n->dim] >= sn.c[(int)n->dim] &&
                n->c[(int)n->dim] <= sn.c[(int)n->dim + t->ndims]) {
                /* go down the other side */
                top++;
                s[top].n = n->child[!dir];
                while (s[top].n) {
                    n = s[top].n;
                    dir = cmp(&sn, n, n->dim) > 0;
                    s[top].dir = dir;
                    s[top].v = 0;
                    top++;
                    s[top].n = n->child[dir];
                }
            }
        }
    }

    *puid = uid;

    return found;
}

/* initialize tree traversal
 * (re-)sets trav structure
 * returns 0
 */
int kdtree_init_trav(struct kdtrav *trav, struct kdtree *tree)
{
    trav->tree = tree;
    trav->curr_node = tree->root;
    trav->first = 1;
    trav->top = 0;

    return 0;
}

/* traverse the tree
 * useful to get all items in the tree non-recursively
 * struct kdtrav *trav needs to be initialized first
 * returns 1, 0 when finished
 */
int kdtree_traverse(struct kdtrav *trav, double *c, int *uid)
{
    if (trav->curr_node == NULL) {
        if (trav->first)
            G_debug(1, "k-d tree: empty tree");
        else
            G_debug(1, "k-d tree: finished traversing");

        return 0;
    }

    if (trav->first) {
        trav->first = 0;
        return kdtree_first(trav, c, uid);
    }

    return kdtree_next(trav, c, uid);
}

/**********************************************/
/*            internal functions              */

/**********************************************/

static int kdtree_replace(struct kdtree *t, struct kdnode *r)
{
    double mindist;
    int rdir, ordir, dir;
    int ld, rd;
    struct kdnode *n, *rn, * or ;
    struct kdstack {
        struct kdnode *n;
        int dir;
        char v;
    } s[256];
    int top, top2;
    int is_leaf;
    int nr;

    if (!r)
        return 0;
    if (!r->child[0] && !r->child[1])
        return 0;

    /* do not call kdtree_balance in this fn, this can cause
     * stack overflow due to too many recursive calls */

    /* find replacement for r
     * overwrite r, delete replacement */
    nr = 0;

    /* pick a subtree */
    rdir = 1;

    or = r;
    ld = (! or->child[0] ? -1 : or->child[0]->depth);
    rd = (! or->child[1] ? -1 : or->child[1]->depth);

    if (ld > rd) {
        rdir = 0;
    }

    /* replace old root, make replacement the new root
     * repeat until replacement is leaf */
    ordir = rdir;
    is_leaf = 0;
    s[0].n = or ;
    s[0].dir = ordir;
    top2 = 1;
    mindist = -1;
    while (!is_leaf) {
        rn = NULL;

        /* find replacement for old root */
        top = top2;
        s[top].n = or->child[ordir];

        n = s[top].n;
        rn = n;
        mindist = or->c[(int) or->dim] - n->c[(int) or->dim];
        if (ordir)
            mindist = -mindist;

        /* go down */
        while (s[top].n) {
            n = s[top].n;
            dir = !ordir;
            if (n->dim != or->dim)
                dir = cmp(or, n, n->dim) > 0;
            s[top].dir = dir;
            s[top].v = 0;
            top++;
            s[top].n = n->child[dir];
        }

        /* go back up */
        while (top > top2) {
            top--;

            if (!s[top].v) {
                s[top].v = 1;
                n = s[top].n;
                if ((cmp(rn, n, or->dim) > 0) == ordir) {
                    rn = n;
                    mindist = or->c[(int) or->dim] - n->c[(int) or->dim];
                    if (ordir)
                        mindist = -mindist;
                }

                /* look on the other side ? */
                dir = s[top].dir;
                if (n->dim != or->dim && mindist >= fabs(n->c[(int)n->dim] -
                                                         n->c[(int)n->dim])) {
                    /* go down the other side */
                    top++;
                    s[top].n = n->child[!dir];
                    while (s[top].n) {
                        n = s[top].n;
                        dir = !ordir;
                        if (n->dim != or->dim)
                            dir = cmp(or, n, n->dim) > 0;
                        s[top].dir = dir;
                        s[top].v = 0;
                        top++;
                        s[top].n = n->child[dir];
                    }
                }
            }
        }

#ifdef KD_DEBUG
        if (!rn)
            G_fatal_error("No replacement");
        if (ordir && or->c[(int) or->dim] > rn->c[(int) or->dim])
            G_fatal_error("rn is smaller");

        if (!ordir && or->c[(int) or->dim] < rn->c[(int) or->dim])
            G_fatal_error("rn is larger");

        if (or->child[1]) {
            dir = cmp(or->child[1], rn, or->dim);
            if (dir < 0) {
                int i;

                for (i = 0; i < t->ndims; i++)
                    G_message("rn c %g, or child c %g", rn->c[i],
                              or->child[1]->c[i]);
                G_fatal_error(
                    "Right child of old root is smaller than rn, dir is %d",
                    ordir);
            }
        }
        if (or->child[0]) {
            dir = cmp(or->child[0], rn, or->dim);
            if (dir > 0) {
                int i;

                for (i = 0; i < t->ndims; i++)
                    G_message("rn c %g, or child c %g", rn->c[i],
                              or->child[0]->c[i]);
                G_fatal_error(
                    "Left child of old root is larger than rn, dir is %d",
                    ordir);
            }
        }
#endif

        is_leaf = (rn->child[0] == NULL && rn->child[1] == NULL);

#ifdef KD_DEBUG
        if (is_leaf && rn->depth != 0)
            G_fatal_error("rn is leaf but depth is %d", (int)rn->depth);
        if (!is_leaf && rn->depth <= 0)
            G_fatal_error("rn is not leaf but depth is %d", (int)rn->depth);
#endif

        nr++;

        /* go to replacement from or->child[ordir] */
        top = top2;
        dir = 1;
        while (dir) {
            n = s[top].n;
            dir = cmp(rn, n, n->dim);
            if (dir) {
                s[top].dir = dir > 0;
                top++;
                s[top].n = n->child[dir > 0];

                if (!s[top].n) {
                    G_fatal_error("(Last) replacement disappeared %d", nr);
                }
            }
        }

#ifdef KD_DEBUG
        if (s[top].n != rn)
            G_fatal_error("rn is unreachable from or");
#endif

        top2 = top;
        s[top2 + 1].n = NULL;

        /* copy replacement to old root */
        memcpy(or->c, rn->c, t->csize);
        or->uid = rn->uid;

        if (!is_leaf) {
            /* make replacement the old root */
            or = rn;

            /* pick a subtree */
            ordir = 1;
            ld = (! or->child[0] ? -1 : or->child[0]->depth);
            rd = (! or->child[1] ? -1 : or->child[1]->depth);
            if (ld > rd) {
                ordir = 0;
            }
            s[top2].dir = ordir;
            top2++;
        }
    }

    if (!rn)
        G_fatal_error("No replacement at all");

    /* delete last replacement */
    if (s[top2].n != rn) {
        G_fatal_error("Wrong top2 for last replacement");
    }
    top = top2 - 1;
    n = s[top].n;
    dir = s[top].dir;
    if (n->child[dir] != rn) {
        G_fatal_error("Last replacement disappeared");
    }
    kdtree_free_node(rn);
    n->child[dir] = NULL;
    t->count--;

    kdtree_update_node(t, n);
    top++;

    /* go back up */
    while (top) {
        top--;
        n = s[top].n;

#ifdef KD_DEBUG
        /* debug directions */
        if (n->child[0]) {
            if (cmp(n->child[0], n, n->dim) > 0)
                G_warning("Left child is larger");
        }
        if (n->child[1]) {
            if (cmp(n->child[1], n, n->dim) < 1)
                G_warning("Right child is not larger");
        }
#endif

        /* update node */
        kdtree_update_node(t, n);
    }

    return nr;
}

static int kdtree_balance(struct kdtree *t, struct kdnode *r, int bmode)
{
    struct kdnode * or ;
    int dir;
    int rd, ld;
    int old_depth;
    int btol;

    if (!r) {
        return 0;
    }

    ld = (!r->child[0] ? -1 : r->child[0]->depth);
    rd = (!r->child[1] ? -1 : r->child[1]->depth);
    old_depth = MAX(ld, rd) + 1;

    if (old_depth != r->depth) {
        G_warning("balancing: depth is wrong: %d != %d", r->depth, old_depth);
        kdtree_update_node(t, r);
    }

    /* subtree difference */
    btol = t->btol;
    if (!r->child[0] || !r->child[1])
        btol = 2;
    dir = -1;
    ld = (!r->child[0] ? -1 : r->child[0]->depth);
    rd = (!r->child[1] ? -1 : r->child[1]->depth);
    if (ld > rd + btol) {
        dir = 0;
    }
    else if (rd > ld + btol) {
        dir = 1;
    }
    else {
        return 0;
    }

    or = kdtree_newnode(t);
    memcpy(or->c, r->c, t->csize);
    or->uid = r->uid;
    or->dim = t->nextdim[r->dim];

    if (!kdtree_replace(t, r))
        G_fatal_error("kdtree_balance: nothing replaced");

#ifdef KD_DEBUG
    if (!cmp(r, or, r->dim)) {
        G_warning("kdtree_balance: replacement failed");
        kdtree_free_node(or);

        return 0;
    }
#endif

    r->child[!dir] =
        kdtree_insert2(t, r->child[!dir], or, bmode, 1); /* bmode */

    /* update node */
    kdtree_update_node(t, r);

    if (r->depth == old_depth) {
        G_debug(4, "balancing had no effect");
        return 1;
    }

    if (r->depth > old_depth)
        G_fatal_error("balancing failed");

    return 1;
}

static struct kdnode *kdtree_insert2(struct kdtree *t, struct kdnode *r,
                                     struct kdnode *nnew, int balance, int dc)
{
    struct kdnode *n;
    struct kdstack {
        struct kdnode *n;
        int dir;
    } s[256];
    int top;
    int dir;
    int bmode;

    if (!r) {
        r = nnew;
        t->count++;

        return r;
    }

    /* level of recursion */
    rcalls++;
    if (rcallsmax < rcalls)
        rcallsmax = rcalls;

    /* balancing modes
     * bmode = 0: no recursion (only insert -> balance -> insert)
     *            slower, higher tree depth
     * bmode = 1: recursion (insert -> balance -> insert -> balance ...)
     *            faster, more compact tree
     *  */
    bmode = 1;

    /* find node with free child */
    top = 0;
    s[top].n = r;
    while (s[top].n) {

        n = s[top].n;

        if (!cmpc(nnew, n, t) && (!dc || nnew->uid == n->uid)) {

            G_debug(1, "KD node exists already, nothing to do");
            kdtree_free_node(nnew);

            if (!balance) {
                rcalls--;
                return r;
            }

            break;
        }
        dir = cmp(nnew, n, n->dim) > 0;
        s[top].dir = dir;

        top++;
        if (top > 255)
            G_fatal_error("depth too large: %d", top);
        s[top].n = n->child[dir];
    }

    if (!s[top].n) {
        /* insert to child pointer of parent */
        top--;
        n = s[top].n;
        dir = s[top].dir;
        n->child[dir] = nnew;
        nnew->dim = t->nextdim[n->dim];

        t->count++;
        top++;
    }

    /* go back up */
    while (top) {
        top--;
        n = s[top].n;

        /* update node */
        kdtree_update_node(t, n);

        /* do not balance on the way back up */

#ifdef KD_DEBUG
        /* debug directions */
        if (n->child[0]) {
            if (cmp(n->child[0], n, n->dim) > 0)
                G_warning("Insert2: Left child is larger");
        }
        if (n->child[1]) {
            if (cmp(n->child[1], n, n->dim) < 1)
                G_warning("Insert2: Right child is not larger");
        }
#endif
    }

    if (balance) {
        int iter, bmode2;

        /* fix any inconsistencies in the (sub-)tree */
        iter = 0;
        bmode2 = 0;
        top = 0;
        s[top].n = r;
        while (top >= 0) {

            n = s[top].n;

            /* top-down balancing
             * slower but more compact */
            if (!bmode2) {
                while (kdtree_balance(t, n, bmode))
                    ;
            }

            /* go down */
            if (n->child[0] && n->child[0]->balance) {
                dir = 0;
                top++;
                s[top].n = n->child[dir];
            }
            else if (n->child[1] && n->child[1]->balance) {
                dir = 1;
                top++;
                s[top].n = n->child[dir];
            }
            /* go back up */
            else {

                /* bottom-up balancing
                 * faster but less compact */
                if (bmode2) {
                    while (kdtree_balance(t, n, bmode))
                        ;
                }
                top--;
                if (top >= 0) {
                    kdtree_update_node(t, s[top].n);
                }
                if (!bmode2 && top == 0) {
                    iter++;
                    if (iter == 2) {
                        /* the top node has been visited twice,
                         * switch from top-down to bottom-up balancing */
                        iter = 0;
                        bmode2 = 1;
                    }
                }
            }
        }
    }

    rcalls--;

    return r;
}

/* start traversing the tree
 * returns pointer to first item
 */
static int kdtree_first(struct kdtrav *trav, double *c, int *uid)
{
    /* get smallest item */
    while (trav->curr_node->child[0] != NULL) {
        trav->up[trav->top++] = trav->curr_node;
        trav->curr_node = trav->curr_node->child[0];
    }

    memcpy(c, trav->curr_node->c, trav->tree->csize);
    *uid = trav->curr_node->uid;

    return 1;
}

/* continue traversing the tree in ascending order
 * returns pointer to data item, NULL when finished
 */
static int kdtree_next(struct kdtrav *trav, double *c, int *uid)
{
    if (trav->curr_node->child[1] != NULL) {
        /* something on the right side: larger item */
        trav->up[trav->top++] = trav->curr_node;
        trav->curr_node = trav->curr_node->child[1];

        /* go down, find smallest item in this branch */
        while (trav->curr_node->child[0] != NULL) {
            trav->up[trav->top++] = trav->curr_node;
            trav->curr_node = trav->curr_node->child[0];
        }
    }
    else {
        /* at smallest item in this branch, go back up */
        struct kdnode *last;

        do {
            if (trav->top == 0) {
                trav->curr_node = NULL;
                break;
            }
            last = trav->curr_node;
            trav->curr_node = trav->up[--trav->top];
        } while (last == trav->curr_node->child[1]);
    }

    if (trav->curr_node != NULL) {
        memcpy(c, trav->curr_node->c, trav->tree->csize);
        *uid = trav->curr_node->uid;

        return 1;
    }

    return 0; /* finished traversing */
}
