#define CODE_ERROR_LOCATIONS
#include <gtest/gtest.h>

#include <NBCore/Arrays.hpp>

template<unsigned int N>
using Coord = typename nb::NDArray<uint8_t, N>::Coordinate;

template<typename A, typename B>
void ASSERT_ARRAY_EQ(A a, B b) {
	unsigned int n = a.size();
	for (int i = 0; i < n; ++i) {
		ASSERT_EQ(a[i], b[i]);
	}
}

TEST(TestNDArray, TestBasic) {
	float* data_f = new float[9] {0,1,2,3,4,5,6,7,8};
	nb::NDArray<float, 2> arr({3, 3}, data_f);
	Coord<2> x{1,1};
	ASSERT_ARRAY_EQ(arr.dimensions, Coord<2>{3,3});
	ASSERT_FLOAT_EQ(data_f[4], arr.at(x));
	arr.at(x) = 0.5;
	ASSERT_FLOAT_EQ(0.5, arr.at(x));
	delete[] data_f;

	Coord<3> y{1,0,1};
	unsigned int* data_ui = new unsigned int[8] {0,1,2,3,4,5,6,7};
	nb::NDArray<unsigned int, 3> arr3d({2,2,2}, data_ui);
	ASSERT_ARRAY_EQ(arr3d.dimensions, Coord<3>{2,2,2});
	ASSERT_EQ(data_ui[5], arr3d.at(y));	
	arr3d[{0,1,1}] = 23;
	ASSERT_EQ(23, arr3d.at({0,1,1}));
	delete[] data_ui;

	nb::NDArray<uint8_t, 4> arr4d({4,4,4,4});
	arr4d[{0,1,2,3}] = 57;
	ASSERT_ARRAY_EQ(arr4d.dimensions, Coord<4>{4,4,4,4});
	ASSERT_EQ(57, arr4d.at({0,1,2,3}));
}

TEST(TestNDArray, TestMove) {
	unsigned int data_ui[8] = {0,1,2,3,4,5,6,7};
	nb::NDArray<unsigned int, 2> arr_1({4, 2}, data_ui);
	nb::NDArray<unsigned int, 2> arr_2(std::move(arr_1));
	Coord<2> x{2, 1};
	ASSERT_EQ(arr_1.data(), nullptr);
	ASSERT_ARRAY_EQ(arr_1.dimensions, arr_2.dimensions);
	ASSERT_EQ(arr_2.at(x), 6);
}

TEST(TestNDArray, TestConversion) {
	nb::NDArray<unsigned int, 2> arr1({3,3});
	nb::NDArray<unsigned int, 3> arr2(std::move(arr1));
	ASSERT_EQ(arr1.data(), nullptr);
	ASSERT_ARRAY_EQ(arr2.dimensions, Coord<3>{3,3,1});
	nb::NDArray<unsigned int, 2> arr3(std::move(arr2));
	ASSERT_EQ(arr2.data(), nullptr);
	ASSERT_ARRAY_EQ(arr3.dimensions, Coord<2>{3,3});
}

TEST(TestNDArray, TestFromMemory) {
	auto data = new float[9]{1,2,3,4,5,6,7,8,9};
	auto arr1 = nb::NDArray<float, 2>::MemoryToNDArray({3,3}, data);
	ASSERT_EQ(arr1.data(), data);
	ASSERT_FLOAT_EQ(arr1.at({1,2}), data[5]);
	auto data2 = new float[9]{9,8,7,6,5,4,3,2,1};
	arr1.data(data2);
	ASSERT_EQ(arr1.data(), data);
	ASSERT_FLOAT_EQ(arr1.at({3,0}), data2[6]);
	delete[] data2;
}