/*
  Copyright (C) 2018 David Robillard <d@drobilla.net>

  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 2 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 <https://www.gnu.org/licenses/>.
*/

#undef NDEBUG

#include "test_utils.hpp"

#include "chilbert/BigBitVec.hpp"
#include "chilbert/FixBitVec.hpp"
#include "chilbert/Hilbert.hpp"

#include <gmpxx.h>

/// Return a `D`-dimensional point with `M` bits of precision per dimension
template <size_t M, size_t D>
std::array<uint64_t, D>
make_random_point(Context& ctx)
{
	std::array<uint64_t, D> p;
	for (size_t i = 0; i < D; ++i) {
		p[i] = rand_between(ctx, 0, (1UL << M) - 1);
	}
	return p;
}

/// Return a `D`-dimensional point within `ms` per-dimension precision
template <size_t D>
std::array<uint64_t, D>
make_random_point(Context& ctx, const std::array<size_t, D>& ms)
{
	std::array<uint64_t, D> p;
	for (size_t i = 0; i < D; ++i) {
		p[i] = rand_between(ctx, 0, (1UL << ms[i]) - 1);
	}
	return p;
}

/// Return the squared distance from point `a` to point `b`
template <class T, size_t D>
T
squared_distance(const std::array<T, D>& a, const std::array<T, D>& b)
{
	T sdist = 0;
	for (size_t i = 0; i < D; ++i) {
		const T diff = a[i] > b[i] ? a[i] - b[i] : b[i] - a[i];
		sdist += diff * diff;
	}
	return sdist;
}

/// Convert bit vector `vec` to a big integer
template <class T>
mpz_class
to_big_int(const T& vec)
{
	mpz_t ia;
	mpz_init(ia);
	mpz_import(
	  ia, vec.num_racks(), -1, sizeof(chilbert::FBV_UINT), 0, 0, vec.racks());
	const mpz_class num(ia);
	mpz_clear(ia);
	return num;
}

/// Convert big integer `num` to a bit vector
template <class T, size_t M>
T
from_big_int(const mpz_class& num)
{
	T      vec   = make_zero_bitvec<T, M>();
	size_t count = 0;
	mpz_export(vec.racks(),
	           &count,
	           -1,
	           sizeof(chilbert::FBV_UINT),
	           0,
	           0,
	           num.get_mpz_t());
	assert(count <= static_cast<size_t>(vec.num_racks()));
	return vec;
}

template <class T, size_t M, size_t D>
void
test_standard(Context& ctx)
{
	static_assert(M < sizeof(chilbert::FBV_UINT) * CHAR_BIT, "");

	// Generate random point and its hilbert index
	const auto pa = make_random_point<M, D>(ctx);

	T ha = make_zero_bitvec<T, D * M>();
	assert(ha.size() >= D * M);
	chilbert::coordsToIndex(pa.data(), M, D, ha);

	{
		// Ensure unmapping results in the original point
		auto pa_out = make_random_point<M, D>(ctx);
		chilbert::indexToCoords(pa_out.data(), M, D, ha);
		assert(pa_out == pa);
	}

	// Convert hilbert indices to a big integer for manipulation/comparison
	const auto ia = to_big_int(ha);

	// Generate the next hilbert index
	const auto ib = ia + 1;
	const auto hb = from_big_int<T, D * M>(ib);

	// Unmap next hilbert index to a point
	auto pb = make_random_point<M, D>(ctx);
	chilbert::indexToCoords(pb.data(), M, D, hb);

	// Ensure next point is 1 unit of distance away from first
	assert(squared_distance(pa, pb) == 1);
}

template <class T, size_t M, size_t D>
void
test_compact(Context& ctx)
{
	static_assert(M < sizeof(chilbert::FBV_UINT) * CHAR_BIT, "");

	// Generate random point and its hilbert index
	const auto ms = make_random_precisions<D * M, D>(ctx);
	const auto pa = make_random_point<D>(ctx, ms);

	T ha = make_zero_bitvec<T, D * M>();
	assert(ha.size() >= D * M);
	chilbert::coordsToCompactIndex(pa.data(), ms.data(), D, ha);

	{
		// Ensure unmapping results in the original point
		auto pa_out = make_random_point<M, D>(ctx);
		chilbert::compactIndexToCoords(pa_out.data(), ms.data(), D, ha);
		assert(pa_out == pa);
	}

	// Convert hilbert indices to a big integer for manipulation/comparison
	const auto ia = to_big_int(ha);

	// Generate the next hilbert index
	const auto ib = ia + 1;
	const auto hb = from_big_int<T, D * M>(ib);

	// Unmap next hilbert index to a point
	auto pb = make_random_point<M, D>(ctx);
	chilbert::compactIndexToCoords(pb.data(), ms.data(), D, hb);

	// Ensure next point is 1 unit of distance away from first
	assert(squared_distance(pa, pb) == 1);
}

int
main()
{
	Context ctx;

	test_standard<chilbert::CFixBitVec, 4, 2>(ctx);
	test_standard<chilbert::CFixBitVec, 32, 2>(ctx);
	test_standard<chilbert::CFixBitVec, 16, 4>(ctx);
	test_standard<chilbert::CFixBitVec, 8, 8>(ctx);
	test_standard<chilbert::CFixBitVec, 4, 16>(ctx);
	test_standard<chilbert::CFixBitVec, 2, 32>(ctx);
	test_standard<chilbert::CFixBitVec, 1, 64>(ctx);

	test_standard<chilbert::CBigBitVec, 4, 65>(ctx);
	test_standard<chilbert::CBigBitVec, 32, 64>(ctx);
	test_standard<chilbert::CBigBitVec, 63, 128>(ctx);

	test_compact<chilbert::CFixBitVec, 4, 2>(ctx);
	test_compact<chilbert::CFixBitVec, 32, 2>(ctx);
	test_compact<chilbert::CFixBitVec, 16, 4>(ctx);
	test_compact<chilbert::CFixBitVec, 8, 8>(ctx);
	test_compact<chilbert::CFixBitVec, 4, 16>(ctx);
	test_compact<chilbert::CFixBitVec, 2, 32>(ctx);
	test_compact<chilbert::CFixBitVec, 1, 64>(ctx);

	test_compact<chilbert::CBigBitVec, 4, 65>(ctx);
	test_compact<chilbert::CBigBitVec, 32, 64>(ctx);
	test_compact<chilbert::CBigBitVec, 63, 128>(ctx);

	return 0;
}