#include <oak/basic_tree.h>
#include <oak/oak.h>

class RangeTreeTests : public CxxTest::TestSuite
{
	struct key_t
	{
		key_t (ssize_t pos = 0) : position(pos)   { }
		key_t operator== (key_t const& rhs) const { return position == rhs.position; }
		key_t operator+ (key_t const& rhs) const { return key_t(position + rhs.position); }
		key_t operator- (key_t const& rhs) const { return key_t(position - rhs.position); }

		ssize_t position;
	};

	static int position_comp (ssize_t key, key_t const& offset, key_t const& node) { return key < offset.position + node.position ? -1 : (key == offset.position + node.position ? 0 : +1); }
	typedef oak::basic_tree_t<key_t> tree_t;

	static void set (tree_t& tree, ssize_t pos)
	{
		auto it = tree.upper_bound(pos, &position_comp);
		if(it != tree.begin())
		{
			auto tmp = it;
			--tmp;
			pos -= tmp->offset.position + tmp->key.position;
		}

		if(it != tree.end())
		{
			it->key.position -= pos;
			tree.update_key(it);
		}

		tree.insert(it, pos);
	}

	static void unset (tree_t& tree, ssize_t pos)
	{
		auto it = tree.find(pos, &position_comp);
		if(it == tree.end())
			return;

		auto tmp = it;
		if(++tmp != tree.end())
		{
			tmp->key.position += it->key.position;
			tree.update_key(tmp);
		}

		tree.erase(it);
	}

	void adjust (tree_t& tree, ssize_t pos, size_t distance)
	{
		auto it = tree.lower_bound(pos, &position_comp);
		if(it == tree.end())
			return;

		it->key.position += distance;
		tree.update_key(it);
	}

public:
	void test_adjustment ()
	{
		tree_t tree;
	
		set(tree, 20);
		set(tree, 30);
	
		adjust(tree, 15, 5); // 20 → 25, 30 → 35
		adjust(tree, 25, 5); // 25 → 30, 35 → 40
		adjust(tree, 35, 5); // 30 → 30, 40 → 45
		adjust(tree, 45, 5); // 30 → 30, 45 → 50
		adjust(tree, 55, 5); // 30 → 30, 50 → 50
	
		TS_ASSERT_EQUALS(tree.size(), 2);
		TS_ASSERT_DIFFERS(tree.find(30, &position_comp), tree.end());
		TS_ASSERT_DIFFERS(tree.find(50, &position_comp), tree.end());
	
		adjust(tree, 60, -5); // 30 → 30, 50 → 50
		adjust(tree, 50, -5); // 30 → 30, 50 → 45
		adjust(tree, 40, -5); // 30 → 30, 45 → 40
		adjust(tree, 30, -5); // 30 → 25, 40 → 35
		adjust(tree, 20, -5); // 25 → 20, 35 → 30
	
		TS_ASSERT_EQUALS(tree.size(), 2);
		TS_ASSERT_DIFFERS(tree.find(20, &position_comp), tree.end());
		TS_ASSERT_DIFFERS(tree.find(30, &position_comp), tree.end());
	}

	void test_range_tree ()
	{
		std::set<ssize_t> tmp;
		while(tmp.size() < 32)
			tmp.insert((random() & 0xFFFFFF) - 0x7FFFFF);

		std::vector<ssize_t> keys(tmp.begin(), tmp.end());
		std::random_shuffle(keys.begin(), keys.end());

		tree_t tree;
		iterate(key, keys)
			set(tree, *key);

		std::vector<ssize_t> fromTree;
		iterate(it, tree)
			fromTree.push_back(it->offset.position + it->key.position);

		std::sort(keys.begin(), keys.end());
		TS_ASSERT_EQUALS(fromTree, keys);

		// ==================
		// = Test adjusting =
		// ==================

		std::random_shuffle(keys.begin(), keys.end());
		ssize_t pos = keys[0], distance = 200;
		adjust(tree, pos, distance);
		iterate(key, keys)
		{
			if(pos <= *key)
				*key += distance;
		}

		iterate(key, keys)
		{
			TS_ASSERT_DIFFERS(tree.find(*key, &position_comp), tree.end());
		}

		// =======================
		// = Erase half the keys =
		// =======================

		std::random_shuffle(keys.begin(), keys.end());
		for(size_t i = keys.size() >> 1; i < keys.size(); ++i)
		{
			auto it = tree.find(keys[i], &position_comp);
			TS_ASSERT_DIFFERS(it, tree.end());
			unset(tree, it->offset.position + it->key.position);
			tmp.erase(keys[i]);
		}

		keys.resize(keys.size() >> 1);
		TS_ASSERT_EQUALS(tree.size(), keys.size());

		std::random_shuffle(keys.begin(), keys.end());
		iterate(key, keys)
			unset(tree, *key);
		TS_ASSERT(tree.empty());
	}
};
