/*
 *  gretl -- Gnu Regression, Econometrics and Time-series Library
 *  Copyright (C) 2001 Allin Cottrell and Riccardo "Jack" Lucchetti
 *
 *  This program 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, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

/* The code here is a simplified version of GLib's GTree, modified
   such that we can use doubles for key and value rather than
   pointers; this gives us a means of setting up fast lookup for
   the hansl replace() function as applied to series or matrices.
   See glib/gtree.c at https://github.com/GNOME/glib for the
   original code, which is under the GNU Lesser General Public
   License, version 2.1 or higher.
*/

#include <glib.h>
#include <math.h>
#include "gretl_btree.h"

#define MAX_BTREE_HEIGHT 40

typedef struct _BTreeNode  BTreeNode;

struct _BTree {
    BTreeNode *root;
    guint nnodes;
};

struct _BTreeNode {
    gdouble key;        /* key for this node */
    gdouble value;      /* value stored at this node */
    BTreeNode *left;    /* left subtree */
    BTreeNode *right;   /* right subtree */
    gint8 balance;      /* height (right) - height (left) */
    guint8 left_child;
    guint8 right_child;
};

static BTreeNode *btree_node_rotate_left (BTreeNode *node);
static BTreeNode *btree_node_rotate_right (BTreeNode *node);

static BTreeNode *btree_node_new (gdouble key,
				  gdouble value)
{
    BTreeNode *node = g_slice_new(BTreeNode);

    node->balance = 0;
    node->left = NULL;
    node->right = NULL;
    node->left_child = FALSE;
    node->right_child = FALSE;
    node->key = key;
    node->value = value;

    return node;
}

static gint key_compare (gdouble a, gdouble b)
{
    /* Handle NaNs in matrices? Let NaN be "bigger" than
       anything else, and in this context let it compare
       equal to itself.
    */
    if (isnan(a)) {
	return isnan(b) ? 0 : 1;
    } else if (isnan(b)) {
	return -1;
    }

    return a - b > 0 ? 1 : a == b ? 0 : -1;
}

BTree *gretl_btree_new (void)
{
    BTree *tree = g_slice_new(BTree);

    tree->root = NULL;
    tree->nnodes = 0;

    return tree;
}

static inline BTreeNode *btree_first_node (BTree *tree)
{
    BTreeNode *tmp;

    if (!tree->root) {
	return NULL;
    }

    tmp = tree->root;

    while (tmp->left_child) {
	tmp = tmp->left;
    }

    return tmp;
}

static inline BTreeNode *btree_node_previous (BTreeNode *node)
{
    BTreeNode *tmp = node->left;

    if (node->left_child) {
	while (tmp->right_child) {
	    tmp = tmp->right;
	}
    }

    return tmp;
}

static inline BTreeNode *btree_node_next (BTreeNode *node)
{
    BTreeNode *tmp = node->right;

    if (node->right_child) {
	while (tmp->left_child) {
	    tmp = tmp->left;
	}
    }

    return tmp;
}

static void btree_remove_all (BTree *tree)
{
    BTreeNode *node, *next;

    g_return_if_fail(tree != NULL);

    node = btree_first_node(tree);

    while (node) {
	next = btree_node_next(node);
	g_slice_free(BTreeNode, node);
	node = next;
    }

    tree->root = NULL;
    tree->nnodes = 0;
}

void gretl_btree_destroy (BTree *tree)
{
    g_return_if_fail(tree != NULL);

    btree_remove_all(tree);
    g_slice_free(BTree, tree);
}

static BTreeNode *btree_node_balance (BTreeNode *node)
{
    if (node->balance < -1) {
	if (node->left->balance > 0) {
	    node->left = btree_node_rotate_left(node->left);
	}
	node = btree_node_rotate_right(node);
    } else if (node->balance > 1) {
	if (node->right->balance < 0) {
	    node->right = btree_node_rotate_right(node->right);
	}
	node = btree_node_rotate_left(node);
    }

    return node;
}

void gretl_btree_insert (BTree *tree,
			 gdouble key,
			 gdouble value)
{
    BTreeNode *node;
    BTreeNode *path[MAX_BTREE_HEIGHT];
    int idx;

    g_return_if_fail(tree != NULL);

    if (!tree->root) {
	tree->root = btree_node_new(key, value);
	tree->nnodes++;
	return;
    }

    idx = 0;
    path[idx++] = NULL;
    node = tree->root;

    while (1) {
	int cmp = key_compare(key, node->key);

	if (cmp == 0) {
	    node->value = value;
	    return;
	} else if (cmp < 0) {
	    if (node->left_child) {
		path[idx++] = node;
		node = node->left;
	    } else {
		BTreeNode *child = btree_node_new(key, value);

		child->left = node->left;
		child->right = node;
		node->left = child;
		node->left_child = TRUE;
		node->balance -= 1;

		tree->nnodes++;
		break;
	    }
	} else {
	    if (node->right_child) {
		path[idx++] = node;
		node = node->right;
	    } else {
		BTreeNode *child = btree_node_new(key, value);

		child->right = node->right;
		child->left = node;
		node->right = child;
		node->right_child = TRUE;
		node->balance += 1;

		tree->nnodes++;
		break;
	    }
	}
    }

    /* restore balance */
    while (1) {
	BTreeNode *bparent = path[--idx];
	gboolean left_node = (bparent && node == bparent->left);

	g_assert(!bparent || bparent->left == node || bparent->right == node);

	if (node->balance < -1 || node->balance > 1) {
	    node = btree_node_balance (node);
	    if (bparent == NULL) {
		tree->root = node;
	    } else if (left_node) {
		bparent->left = node;
	    } else {
		bparent->right = node;
	    }
	}

	if (node->balance == 0 || bparent == NULL) {
	    break;
	}

	if (left_node) {
	    bparent->balance -= 1;
	} else {
	    bparent->balance += 1;
	}

	node = bparent;
    }
}

static BTreeNode *btree_find_node (BTree *tree,
				   gdouble key)
{
    BTreeNode *node = tree->root;
    gint cmp;

    if (!node) {
	return NULL;
    }

    while (1) {
	cmp = key_compare(key, node->key);
	if (cmp == 0) {
	    return node;
	} else if (cmp < 0) {
	    if (!node->left_child) {
		return NULL;
	    }
	    node = node->left;
        } else {
	    if (!node->right_child) {
		return NULL;
	    }
	    node = node->right;
	}
    }
}

gdouble gretl_btree_lookup (BTree *tree,
			    gdouble key)
{
    BTreeNode *node;

    g_return_val_if_fail(tree != NULL, key);

    node = btree_find_node(tree, key);

    return node ? node->value : key;
}

void gretl_btree_minmax (BTree *tree,
			 gdouble *keymin,
			 gdouble *keymax)
{
    BTreeNode *L, *R;

    g_return_if_fail(tree != NULL);

    L = R = tree->root;

    while (L) {
	*keymin = L->key;
	L = L->left;
    }

    while (R) {
	*keymax = R->key;
	R = R->right;
    }
}

static BTreeNode *btree_node_rotate_left (BTreeNode *node)
{
    BTreeNode *right;
    gint a_bal, b_bal;

    right = node->right;

    if (right->left_child) {
	node->right = right->left;
    } else {
	node->right_child = FALSE;
	right->left_child = TRUE;
    }
    right->left = node;

    a_bal = node->balance;
    b_bal = right->balance;

    if (b_bal <= 0) {
	if (a_bal >= 1) {
	    right->balance = b_bal - 1;
	} else {
	    right->balance = a_bal + b_bal - 2;
	}
	node->balance = a_bal - 1;
    } else {
	if (a_bal <= b_bal) {
	    right->balance = a_bal - 2;
	} else {
	    right->balance = b_bal - 1;
	}
	node->balance = a_bal - b_bal - 1;
    }

    return right;
}

static BTreeNode *btree_node_rotate_right (BTreeNode *node)
{
    BTreeNode *left;
    gint a_bal, b_bal;

    left = node->left;

    if (left->right_child) {
	node->left = left->right;
    } else {
	node->left_child = FALSE;
	left->right_child = TRUE;
    }
    left->right = node;

    a_bal = node->balance;
    b_bal = left->balance;

    if (b_bal <= 0) {
	if (b_bal > a_bal) {
	    left->balance = b_bal + 1;
	} else {
	    left->balance = a_bal + 2;
	}
	node->balance = a_bal - b_bal + 1;
    } else {
	if (a_bal <= -1) {
	    left->balance = b_bal + 1;
	} else {
	    left->balance = a_bal + b_bal + 2;
	}
	node->balance = a_bal + 1;
    }

    return left;
}
