#include "cnn1986_treap_cache.h"

// ===================================================================
// Types
// ===================================================================

typedef struct treap_node_s {
	uint64_t key;
	uint32_t weight;
	struct treap_node_s *left, *right, *parent;
	/*
	 * colour_stack_t colours[25];
	 * data_t celldat[25];
	 * uint8_t white_count, black_count;
	 */
	float result;
} * TreapNode;

// ===================================================================
// Variables
// ===================================================================

uint32_t cnn1986_num_cached;
static TreapNode root;

// ===================================================================
// Helper declarations
// ===================================================================

void recurse_tree(const TreapNode n);
TreapNode new_treap_node(const float in_result, const uint64_t key);
void bubble_up(const TreapNode n);

// ===================================================================
// Exported functions
// ===================================================================

int cnn1986_cache_init(void) {
	root = NULL;
	cnn1986_num_cached = 0;
	return EXIT_SUCCESS;
}

void cnn1986_cache_free(void) {
	recurse_tree(root);
	return;
}

int cnn1986_cache_seek(const uint64_t key, float *out_result) {
	if (root == NULL) return EXIT_FAILURE;
	TreapNode n = root;

	while (n != NULL && n->key != key) {
		if (n->key > key) n = n->right;
		else n = n->left;
	}
	if (n == NULL) return EXIT_FAILURE;
	*out_result = n->result;
	return EXIT_SUCCESS;
}

int cnn1986_cache_insert(const uint64_t key, const float in_result) {
	if (root == NULL) {
		root = new_treap_node(in_result, key);
		cnn1986_num_cached = 1;
		return EXIT_SUCCESS;
	}
	TreapNode s = root, n = root,
		m = new_treap_node(in_result, key);
	// Find the correct position by doing a BST traversal
	while (n!=NULL) {
		s = n;
		if (n->key >= key) n = n->right;
		else n = n->left;
	}
	// Make it a leaf
	if (s->key > key) s->right = m;
	else s->left = m;
	m->parent = s;
	// Now bubble upward to satisfy the heap property
	bubble_up(m);
	cnn1986_num_cached++;
	return EXIT_SUCCESS;
}

// ===================================================================
// Helper implementations
// ===================================================================

void recurse_tree(TreapNode n) {
	if (n==NULL) return;
	if (n->left != NULL) recurse_tree(n->left);
	if (n->right != NULL) recurse_tree(n->right);
	free(n);
}

TreapNode new_treap_node(const float in_result, const uint64_t key) {
	TreapNode n = malloc(sizeof(struct treap_node_s));
	// TODO: trap
	n->key = key;
	n->left = NULL;
	n->right = NULL;
	n->parent = NULL;
	n->weight = RANDOM32;
	XORSHIFT;
	// Set value
	n->result = in_result;
	return n;
}

void rotate_left(TreapNode n) {
	TreapNode a = n->parent, b = a->left, c = n->left;
	/*
		We are the right child, so do this
					 a           n
					/ \         / \
				 b   n  -->  a   d
						/ \     / \
					 c   d   b   c
	*/
	n->parent = a->parent;
	// We may have to repair one level up as well
	if (a->parent!=NULL) {
		if (a->parent->left == a) a->parent->left = n;
		else a->parent->right = n;
	}
	a->parent = n;
	n->left = a; a->parent = n;
	a->left = b; if (b!=NULL) b->parent = a;
	a->right = c; if (c!=NULL) c->parent = a;
}

void rotate_right(TreapNode n) {
	TreapNode a = n->parent, b = a->right, d = n->right;
	 /*
		 We are the left child, so do this
				 a           n
				/ \         / \
			 n   b  -->  c   a
			/ \             / \
		 c   d           d   b
	 */
	n->parent = a->parent;
	// We may have to repair one level up as well
	if (a->parent!=NULL) {
		if (a->parent->left == a) a->parent->left = n;
		else a->parent->right = n;
	}
	a->parent = n;
	n->right = a; a->parent = n;
	a->left = d; if (d!=NULL) d->parent = a;
	a->right = b; if (b!=NULL) b->parent = a;
}

//This preserves the BST quality of the treap
void bubble_up(TreapNode n) {
	// Nothing to be done in this case
	if (n==NULL || n->parent == NULL) return;
	// Bubble until the treap invariants are satisfied
	while (n->parent != NULL && n->weight < n->parent->weight) {
		if (n->parent->left == n) rotate_right(n);
		else rotate_left(n);
	}
	if (n->parent == NULL) root = n;
}