#pragma once

#include "common.hpp"
#include <kel/lbm/lbm.hpp>

namespace kel {
namespace lbm {
namespace encode {
template<typename Encode>
struct Sycl {
};
}

/*
namespace impl {
template<typename Schema>
struct struct_has_only_equal_dimension_array{};
}
*/
}
}

namespace saw {
template<typename Sch, uint64_t... Dims, typename Encode>
class data<schema::FixedArray<Sch,Dims...>, kel::lbm::encode::Sycl<Encode>> final {
public:
	using Schema = schema::FixedArray<Sch,Dims...>;
private:
	acpp::sycl::queue* q_;
	data<Sch>* values_;

	SAW_FORBID_COPY(data);
public:
	data(acpp::sycl::queue& q__):
		q_{&q__},
		values_{nullptr}
	{
		values_ = acpp::sycl::malloc<data<Sch>>(ct_multiply<uint64_t,Dims...>::value,q_);
	}

	~data(){
		if(not values_){
			return;
		}

		acpp::sycl::free(values_,q_);
	}

	static constexpr data<schema::FixedArray<schema::UInt64, sizeof...(Dims)>> get_dims() {
		return {std::array<uint64_t, sizeof...(Dims)>{Dims...}};
	}

	data<Sch>& at(data<schema::FixedArray<schema::UInt64,sizeof...(Dims)>>& index){
		return values_[kel::lbm::flatten_index<schema::UInt64,sizeof...(Dims)>::apply(index,get_dims()).get()];
	}
	
	constexpr data<Sch,encode::Native>* flat_data() {
		return values_;
	}
};

template<typename Sch, uint64_t Ghost, uint64_t... Sides, typename Encode>
class data<kel::lbm::sch::Chunk<Sch,Ghost,Sides...>,kel::lbm::encode::Sycl<Encode>> final {
public:
	using Schema = kel::lbm::sch::Chunk<Sch,Ghost,Sides...>;
private:
	using InnerSchema = typename Schema::InnerSchema;
	using ValueSchema = typename InnerSchema::ValueType;

	data<InnerSchema, kel::lbm::encode::Sycl<Encode>> values_;
public:
	data(acpp::sycl::queue& q__):
		values_{q__}
	{}

	data<ValueSchema, kel::lbm::encode::Sycl<Encode>>& ghost_at(const data<schema::FixedArray<schema::UInt64,sizeof...(Sides)>>& index){
		return values_.at(index);
	}

	const data<ValueSchema, kel::lbm::encode::Sycl<Encode>>& ghost_at(const data<schema::FixedArray<schema::UInt64,sizeof...(Sides)>>& index) const {
		return values_.at(index);
	}

	static constexpr auto get_ghost_dims() {
		return data<InnerSchema,kel::lbm::encode::Sycl<Encode>>::get_dims();
	}

	data<ValueSchema, kel::lbm::encode::Sycl<Encode>>& at(const data<schema::FixedArray<schema::UInt64,sizeof...(Sides)>>& index){
		std::decay_t<decltype(index)> ind;
		for(uint64_t i = 0u; i < sizeof...(Sides); ++i){
			ind.at({i}) = index.at({i}) + Ghost;
		}
		return values_.at(ind);
	}

	const data<ValueSchema, kel::lbm::encode::Sycl<Encode>>& at(const data<schema::FixedArray<schema::UInt64,sizeof...(Sides)>>& index) const {
		std::decay_t<decltype(index)> ind;
		for(uint64_t i = 0u; i < sizeof...(Sides); ++i){
			ind.at({i}) = index.at({i}) + Ghost;
		}
		return values_.at(ind);
	}

	static constexpr auto get_dims(){
		return data<schema::FixedArray<schema::UInt64, sizeof...(Sides)>,kel::lbm::encode::Sycl<Encode>>{{Sides...}};
	}

	auto flat_data(){
		return values_.flat_data();
	}

	static constexpr auto flat_size() {
		return data<InnerSchema,kel::lbm::encode::Sycl<Encode>>::flat_size();
	}
};

template<uint64_t Ghost, uint64_t... Meta, typename... Sch, string_literal... Keys, typename Encode>
class data<schema::Struct<schema::Member<kel::lbm::sch::Chunk<Sch,Ghost,Meta...>, Keys>...>, kel::lbm::encode::Sycl<Encode> > final {
public:
	static constexpr data<schema::FixedArray<schema::UInt64,sizeof...(Meta)>> meta = {{Meta...}};
	using StorageT = std::tuple<data<kel::lbm::sch::Chunk<Sch,Ghost,Meta...>,kel::lbm::encode::Sycl<Encode>>...>;
	using Schema = schema::Struct<schema::Member<kel::lbm::sch::Chunk<Sch,Ghost,Meta...>, Keys>...>;
private:

	/**
	 * @todo Check by static assert that the members all have the same dimensions. Alternatively
	 * Do it here by specializing.
	 */
	StorageT members_;
public:
	data(acpp::sycl::queue& q__):
		members_{{data<kel::lbm::sch::Chunk<Sch,Ghost,Meta...>,kel::lbm::encode::Sycl<Encode>>{q__}...}}
	{}

	template<saw::string_literal K>
	auto& get(){
		return std::get<parameter_key_pack_index<K, Keys...>::value>(members_);
	}
};


}

namespace kel {
namespace lbm {
namespace impl {
template<typename Sch, typename Encode>
struct sycl_malloc_struct_helper;
	
template<typename... Members, typename Encode>
struct sycl_malloc_struct_helper<sch::Struct<Members...>, Encode> final {
	using Schema = sch::Struct<Members...>;

	template<uint64_t i>
	static saw::error_or<void> allocate_on_device_member(typename saw::data<Schema,encode::Sycl<Encode>>::StorageT& storage, sycl::queue& q){
		if constexpr (i < sizeof...(Members)){
			using M = typename saw::parameter_pack_type<i,Members...>::type;
			auto& ptr = std::get<i>(storage);

			ptr = sycl::malloc_device<M::ValueType::InnerSchema>(1u,q);

			return allocate_on_device_member<i+1u>(storage,q);
		}
		
		return saw::make_void();
	}
	
	static saw::error_or<void> allocate_on_device(saw::data<Schema,encode::Sycl<Encode>>& sycl_data, sycl::queue& q){
		typename saw::data<Schema,encode::Sycl<Encode>>::StorageT storage;
 
		auto eov = allocate_on_device_member<0u>(storage,q);
		sycl_data = {storage, q};

		return eov;
	}
};

template<typename Sch, typename Encode>
struct sycl_copy_helper;

template<typename... Members, typename Encode>
struct sycl_copy_helper<sch::Struct<Members...>, Encode> final {
	using Schema = sch::Struct<Members...>;

	template<uint64_t i>
	static saw::error_or<void> copy_to_device_member(saw::data<Schema,Encode>& host_data, saw::data<Schema,encode::Sycl<Encode>>& sycl_data, sycl::queue& q){
		if constexpr (i < sizeof...(Members)){
			using M = typename saw::parameter_pack_type<i,Members...>::type;
			auto& host_member_data = host_data.template get<i>();
			auto& sycl_member_data = sycl_data.template get<i>();

			auto host_ptr = host_member_data.flat_data();
			auto sycl_ptr = sycl_member_data.flat_data();

			q.memcpy(host_ptr, sycl_ptr, sizeof(std::decay_t<decltype(host_ptr)>) * host_member_data.flat_size() );

			return copy_to_device_member<i+1u>(host_data,sycl_data,q);
		}
		
		return saw::make_void();
	}

	static saw::error_or<void> copy_to_device(saw::data<Schema,Encode>& host_data, saw::data<Schema, encode::Sycl<Encode>>& sycl_data, sycl::queue& q){

		return copy_to_device_member<0u>(host_data, sycl_data, q);
	}

	template<uint64_t i>
	static saw::error_or<void> copy_to_host_member(saw::data<Schema,encode::Sycl<Encode>>& sycl_data, saw::data<Schema,Encode>& host_data, sycl::queue& q){
		if constexpr (i < sizeof...(Members)){
			using M = typename saw::parameter_pack_type<i,Members...>::type;
			auto& host_member_data = host_data.template get<i>();
			auto& sycl_member_data = sycl_data.template get<i>();

			auto host_ptr = host_member_data.flat_data();
			auto sycl_ptr = sycl_member_data.flat_data();

			q.memcpy(sycl_ptr, host_ptr, sizeof(std::decay_t<decltype(host_ptr)>) * host_member_data.flat_size() );

			return copy_to_host_member<i+1u>(sycl_data,host_data,q);
		}
		
		return saw::make_void();
	}
	

	static saw::error_or<void> copy_to_host(saw::data<Schema,Encode>& sycl_data, saw::data<Schema,encode::Sycl<Encode>>& host_data, sycl::queue& q){
		return copy_to_host_member<0u>(sycl_data, host_data, q);
	}
};
}
class device final {
private:
	sycl::queue q_;

	SAW_FORBID_COPY(device);
	SAW_FORBID_MOVE(device);
public:
	device() = default;
	~device() = default;

	template<typename Sch, typename Encode>
	saw::error_or<void> allocate_on_device(saw::data<Sch,encode::Sycl<Encode>>& sycl_data){
		auto eov = impl::sycl_malloc_struct_helper<Sch,Encode>::allocate_on_device(sycl_data, q_);
		if(eov.is_error()){
			return eov;
		}
		return saw::make_void();
	}

	template<typename Sch, typename Encode>
	saw::error_or<void> copy_to_device(saw::data<Sch,Encode>& host_data, saw::data<Sch,encode::Sycl<Encode>>& sycl_data){
		return impl::sycl_copy_helper<Sch,Encode>::copy_to_device(host_data, sycl_data, q_);
	}

	template<typename Sch, typename Encode>
	saw::error_or<void> copy_to_host(saw::data<Sch,encode::Sycl<Encode>>& sycl_data, saw::data<Sch,Encode>& host_data){
		return impl::sycl_copy_helper<Sch,Encode>::copy_to_host(sycl_data, host_data, q_);
	}

	auto& get_handle(){
		return q_;
	}
};
}
}