#pragma once

#include "IndexedNode.h"
#include "IndexedData.h"
#include <thread>

namespace gelib_impl
{

	template <const size_t N>
	class IndexSet
	{

		///////////////////////////////////////////////////////////////////////////////
		// The default number of threads to iterate over a list of independent
		// indexed calculations
		///////////////////////////////////////////////////////////////////////////////
		static size_t default_num_threads;

		///////////////////////////////////////////////////////////////////////////////
		// The indices of each output to some form of n-dimensional storage,
		// The outer vector has one element for each varied dimension
		///////////////////////////////////////////////////////////////////////////////
		std::vector<std::vector<std::unique_ptr<IndexedObject<N>>>> m_indices;

		std::array<size_t, N> m_ordering;

	public:

		static void prepare_single_dimension(
			size_t varied_dimension,
			d_vector const& grid, std::vector<d_vector>& grid_data,
			IndexedData<N> const& input_data, std::vector<const IndexedData<N>*>& v_input_data,
			IndexedData<N>& output_data, std::vector<IndexedData<N>*>& v_output_data)
		{
			grid_data.resize(N);

			for (size_t i = 0; i < N; ++i)
				grid_data[i] = d_vector(input_data.get_dimension_size(i));

			grid_data[varied_dimension] = grid;

			v_input_data = { &input_data };
			v_output_data = { &output_data };
		}

		static void prepare_single_dimension(
			size_t varied_dimension,
			d_vector const& grid, std::vector<d_vector>& grid_data,
			IndexedData<N> const& input_data, std::vector<const IndexedData<N>*>& v_input_data,
			IndexedData<N>& output_data, std::vector<IndexedData<N>*>& v_output_data,
			std::array<size_t, N>& ordering)
		{
			for (size_t i = 0; i < N; ++i)
				ordering[i] = i;

			std::swap(ordering[0], ordering[varied_dimension]);

			prepare_single_dimension(varied_dimension, grid, grid_data, input_data, v_input_data, output_data, v_output_data);
		}

		static void prepare_single_dimension(
			size_t varied_dimension,
			d_vector const& grid, std::vector<d_vector>& grid_data,
			std::vector<const IndexedData<N>*>& v_input_data)
		{
			grid_data.resize(N);

			for (size_t i = 0; i < N; ++i)
				grid_data[i] = d_vector(v_input_data[0]->get_dimension_size(i));

			grid_data[varied_dimension] = grid;
		}


		static void prepare_single_dimension(
			size_t varied_dimension,
			d_vector const& grid, std::vector<d_vector>& grid_data,
			std::vector<const IndexedData<N>*>& v_input_data,
			std::array<size_t, N>& ordering)
		{
			for (size_t i = 0; i < N; ++i)
				ordering[i] = i;

			std::swap(ordering[0], ordering[varied_dimension]);

			prepare_single_dimension(
				varied_dimension,
				grid, grid_data, v_input_data);
		}


		void distribute_work(size_t thread_number, size_t which_indices, size_t* start_ptr, size_t* end_ptr)
		{
			distribute_work(default_num_threads, thread_number, which_indices, start_ptr, end_ptr);
		}

		void distribute_work(size_t num_threads, size_t thread_number, size_t which_indices, size_t* start_ptr, size_t* end_ptr)
		{
			size_t start, end, chunk_size;

			// if chunks not evenly divided make sure last thread picks up remainder
			if (num_threads == (thread_number + 1))
			{
				chunk_size = m_indices[which_indices].size() / num_threads + m_indices[which_indices].size() % num_threads;
				start = thread_number * (m_indices[which_indices].size() / num_threads);
			}
			else
			{
				chunk_size = m_indices[which_indices].size() / num_threads;
				start = thread_number * chunk_size;
			}

			end = start + chunk_size;

			*start_ptr = start;
			*end_ptr = end;
		}

		void construct_indices(
			std::vector<d_vector>& grid_data,
			std::vector<const IndexedData<N>*>& input_data,
			std::vector<IndexedData<N>*>& output_data,
			std::array<size_t, 3> ordering)
		{
			m_ordering = ordering;

			size_t tot_no_fits = 0;
			size_t no_fits = 1;

			m_indices.resize(N);

			std::vector<size_t> order_copy;
			order_copy.resize(N);

			std::copy(ordering.begin(), ordering.end(), order_copy.begin());

			for (size_t i = 0; i < output_data.size(); ++i)
			{
				for (size_t j = 0; j < N; ++j)
				{
					if (j != order_copy[i])
						no_fits *= output_data[i]->get_dimension_size(j);
				}

				m_indices[i].reserve(no_fits);
				tot_no_fits += no_fits;
				no_fits = 1;
			}

			// indices before sorting by the output/dimension they correspond to
			std::vector<std::unique_ptr<IndexedObject<N>>> indices;
			indices.reserve(tot_no_fits);

			// No need to have a reference to created object, as indices required 
			// are loaded into the vector
			DataIndexerRoot<N>(
				grid_data,
				input_data, output_data,
				indices, ordering);

			for (auto& index : indices)
				m_indices[index->get_fixed_depth()].push_back(move(index));

		}

		const size_t num_indices(size_t which) const
		{
			return m_indices[which].size();
		}

		const std::unique_ptr<IndexedObject<N>>& get_indexed_data(size_t which, size_t i) const
		{
			return m_indices[which][i];
		}

		const std::array<int, N>& get_varied_indices(size_t which, size_t i) const
		{
			return m_indices[which][i]->get_indices();
		}

		std::array<size_t, N> const& get_ordering() const
		{
			return m_ordering;
		}

		static size_t get_num_threads();

		static void set_num_threads(size_t threads);

	};

	// default to serial
	template<const size_t N>
	size_t IndexSet<N>::default_num_threads = 1;

	template<size_t N>
	size_t IndexSet<N>::get_num_threads()
	{
		return default_num_threads;
	}

	template<size_t N>
	void IndexSet<N>::set_num_threads(size_t threads)
	{
		default_num_threads = std::min<size_t>((size_t)std::max<unsigned int>(std::thread::hardware_concurrency(), (unsigned int)1), threads);
	}

}