#include <forstio/test/suite.hpp>

#include "../c++/remote.hpp"

#include <random>

namespace {
namespace schema {
using namespace saw::schema;

using TestMixedArray = Array<
	MixedPrecision<Float64,Float32>
>;

using MixedFoo = Interface<
	Member<Function<TestMixedArray, Void>, "foo">
>;

using TestDoubleArray = Array<
	Float64
>;

using DoubleFoo = Interface<
	Member<Function<TestDoubleArray, Void>, "foo">
>;

using TestFloatArray = Array<
	Float32
>;

using FloatFoo = Interface<
	Member<Function<TestFloatArray, Void>, "foo">
>;
}

constexpr uint64_t test_size = 1024ul;

SAW_TEST("SYCL Mixed Test"){
	using namespace saw;

	std::random_device r;
	std::default_random_engine e1{r()};
	std::uniform_real_distribution<> dis{-1.0,1.0};

	data<schema::TestMixedArray> host_data;
	host_data = {test_size};
	for(uint64_t i = 0; i < test_size; ++i){
		host_data.at(i) = static_cast<double>(dis(e1));
	}

	saw::event_loop loop;
	saw::wait_scope wait{loop};

	remote<rmt::Sycl> rmt;

	own<remote_address<rmt::Sycl>> rmt_addr{};

	rmt.resolve_address().then([&](auto addr){
		rmt_addr = std::move(addr);
	}).detach();

	wait.poll();
	SAW_EXPECT(rmt_addr, "Remote address hasn't been filled");
	
	data<schema::TestMixedArray, encode::Native, rmt::Sycl> device_data{host_data};
	
	cl::sycl::event ev;

	interface<schema::MixedFoo, encode::Native,rmt::Sycl, cl::sycl::queue*> cl_iface {
[&](data<schema::TestMixedArray, encode::Native, rmt::Sycl>& in, cl::sycl::queue* cmd) -> error_or<void> {
	
			ev = cmd->submit([&](cl::sycl::handler& h){

				auto acc_buff = in.template access<cl::sycl::access::mode::read_write>(h);

				h.parallel_for(cl::sycl::range<1>(test_size), [=] (cl::sycl::id<1> it){
					acc_buff[0u].at(it[0u]) = acc_buff[0u].at(it[0u]) * data<schema::Float64>{2.0};
				});
			});
			return saw::void_t{};
		}
	};
	auto& device = rmt_addr->get_device();
	
	cl_iface.template call <"foo">(device_data, &(device.get_handle()));
	device.get_handle().wait();
	
	{
		auto end = ev.get_profiling_info<cl::sycl::info::event_profiling::command_end>();
		auto start = ev.get_profiling_info<cl::sycl::info::event_profiling::command_start>();

		std::cout<<"Elapsed kernel time: "<< (end-start) / 1.0e9 << " seconds"<<std::endl;
	}
}

SAW_TEST("SYCL Float Test"){
	using namespace saw;

	std::random_device r;
	std::default_random_engine e1{r()};
	std::uniform_real_distribution<> dis{-1.0,1.0};

	data<schema::TestFloatArray> host_data;
	host_data = {test_size};
	for(uint64_t i = 0; i < test_size; ++i){
		host_data.at(i) = static_cast<double>(dis(e1));
	}

	saw::event_loop loop;
	saw::wait_scope wait{loop};

	remote<rmt::Sycl> rmt;

	own<remote_address<rmt::Sycl>> rmt_addr{};

	rmt.resolve_address().then([&](auto addr){
		rmt_addr = std::move(addr);
	}).detach();

	wait.poll();
	SAW_EXPECT(rmt_addr, "Remote address hasn't been filled");
	
	data<schema::TestFloatArray, encode::Native, rmt::Sycl> device_data{host_data};

	cl::sycl::event ev;

	interface<schema::FloatFoo, encode::Native,rmt::Sycl, cl::sycl::queue*> cl_iface {
[&](data<schema::TestFloatArray, encode::Native, rmt::Sycl>& in, cl::sycl::queue* cmd) -> error_or<void> {
	
			ev = cmd->submit([&](cl::sycl::handler& h){

				auto acc_buff = in.template access<cl::sycl::access::mode::read_write>(h);

				h.parallel_for(cl::sycl::range<1>(test_size), [=] (cl::sycl::id<1> it){
					acc_buff[0u].at(it[0u]) = acc_buff[0u].at(it[0u]) * data<schema::Float32>{2.0};
				});
			});
			return saw::void_t{};
		}
	};
	auto& device = rmt_addr->get_device();
	
	cl_iface.template call <"foo">(device_data, &(device.get_handle()));
	device.get_handle().wait();

	{
		auto end = ev.get_profiling_info<cl::sycl::info::event_profiling::command_end>();
		auto start = ev.get_profiling_info<cl::sycl::info::event_profiling::command_start>();

		std::cout<<"Elapsed kernel time: "<< (end-start) / 1.0e9 << " seconds"<<std::endl;
	}
}

SAW_TEST("SYCL Double Test"){
	using namespace saw;

	std::random_device r;
	std::default_random_engine e1{r()};
	std::uniform_real_distribution<> dis{-1.0,1.0};

	data<schema::TestDoubleArray> host_data;
	host_data = {test_size};
	for(uint64_t i = 0; i < test_size; ++i){
		host_data.at(i) = static_cast<double>(dis(e1));
	}

	saw::event_loop loop;
	saw::wait_scope wait{loop};

	remote<rmt::Sycl> rmt;

	own<remote_address<rmt::Sycl>> rmt_addr{};

	rmt.resolve_address().then([&](auto addr){
		rmt_addr = std::move(addr);
	}).detach();

	wait.poll();
	SAW_EXPECT(rmt_addr, "Remote address hasn't been filled");
	
	data<schema::TestDoubleArray, encode::Native, rmt::Sycl> device_data{host_data};
	
	cl::sycl::event ev;

	interface<schema::DoubleFoo, encode::Native,rmt::Sycl, cl::sycl::queue*> cl_iface {
[&](data<schema::TestDoubleArray, encode::Native, rmt::Sycl>& in, cl::sycl::queue* cmd) -> error_or<void> {
	
			ev = cmd->submit([&](cl::sycl::handler& h){

				auto acc_buff = in.template access<cl::sycl::access::mode::read_write>(h);

				h.parallel_for(cl::sycl::range<1>(test_size), [=] (cl::sycl::id<1> it){
					acc_buff[0u].at(it[0u]) = acc_buff[0u].at(it[0u]) * data<schema::Float64>{2.0};
				});
			});
			return saw::void_t{};
		}
	};
	auto& device = rmt_addr->get_device();
	
	cl_iface.template call <"foo">(device_data, &(device.get_handle()));
	device.get_handle().wait();
	
	{
		auto end = ev.get_profiling_info<cl::sycl::info::event_profiling::command_end>();
		auto start = ev.get_profiling_info<cl::sycl::info::event_profiling::command_start>();

		std::cout<<"Elapsed kernel time: "<< (end-start) / 1.0e9 << " seconds"<<std::endl;
	}
}

}