// Copyright (c) 2019, Christoph Peters
// All rights reserved.
// 
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//     * Neither the name of the Karlsruhe Institute of Technology nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
// 
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
// DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


#pragma once
#include <vector>
#include <array>
#include <istream>
#include <ostream>
#include <memory>
#include <thread>
#include <atomic>
#include <utility>
#include <iostream>

/*! Handles a uniformly spaced grid in 3D space. The grid covers an axis-
	aligned cuboid using equidistant samples in the centers of cuboid cells.
	Each sample holds a 3D vector as data. For example, such a grid may map XYZ
	to Fourier or vice versa.*/
class Grid3D {
public:
	/*! Initializes this grid to cover the given cuboid using the given
		resolution. The initial data is undefined.*/
	Grid3D(std::array<std::size_t, 3> resolution, std::array<float, 3> cuboidMin, std::array<float, 3> cuboidMax);
	/*! Loads this grid from a binary input stream that was previously written
		by store().*/
	Grid3D(std::istream& rInputStream);

	/*! If you pass the given parameters to the constructor, neither rCuboidMin 
		nor rCuboidMax will get a sample in the grid because samples are at 
		cell centers. This behavior is not always desirable and therefore this 
		function modifies rCuboidMin and rCuboidMax minimally such that the 
		given points will be sample locations.*/
	static void expandGrid(std::array<std::size_t, 3> resolution, std::array<float, 3>& rCuboidMin, std::array<float, 3>& rCuboidMax);

	/*! \return The resolution along each axis.*/
	inline const std::array<std::size_t, 3>& getResolution() const {
		return m_resolution;
	}
	/*! \return The minimal coordinates within the covered cuboid.*/
	inline const std::array<float, 3>& getCuboidMin() const {
		return m_cuboidMin;
	}
	/*! \return The minimal coordinates within the covered cuboid.*/
	inline const std::array<float, 3>& getCuboidMax() const {
		return m_cuboidMax;
	}
	/*! \return The total number of voxels/cells/data samples in this grid.*/
	inline std::size_t getCellCount() const {
		return m_resolution[0] * m_resolution[1] * m_resolution[2];
	}
	/*! \return The index triple that corresponds to the given flat index.*/
	inline std::array<std::size_t, 3> unravelIndex(std::size_t flatIndex) const {
		std::size_t z = flatIndex / m_sliceSize;
		flatIndex -= z * m_sliceSize;
		return {
			flatIndex % m_resolution[0],
			flatIndex / m_resolution[0],
			z
		};
	}
	/*! \return The flat index that corresponds to the given index triple.*/
	inline std::size_t ravelMultiIndex(const std::array<std::size_t, 3>& rMultiIndex) const {
		return rMultiIndex[0] + rMultiIndex[1] * m_resolution[0] + rMultiIndex[2] * m_sliceSize;
	}

	/*! \return The coordinate inside the cuboid associated with the given
				multi index.*/
	inline std::array<float, 3> getDataCoordinate(const std::array<std::size_t, 3>& rMultiIndex) const {
		return {
			rMultiIndex[0] * m_indexToCuboidFactor[0] + m_indexToCuboidSummand[0],
			rMultiIndex[1] * m_indexToCuboidFactor[1] + m_indexToCuboidSummand[1],
			rMultiIndex[2] * m_indexToCuboidFactor[2] + m_indexToCuboidSummand[2]
		};
	}
	/*! \return The data at the given flat index.*/
	inline std::array<float, 3>& operator[](std::size_t flatIndex) {
		return m_data[flatIndex];
	}
	/*! \return The data at the given flat index. It may be modified.*/
	inline const std::array<float, 3>& operator[](std::size_t flatIndex) const {
		return m_data[flatIndex];
	}

	/*! Invokes the given std::array<float,3>(std::array<float,3>) functor once
		for each cell passing the result of getDataCoordinate(). The result
		overwrites the data for this cell.
		\param nThread If you pass a non-zero value, the given number of
				threads is going to be spawned to complete the work on the
				cells in parallel. The function does not return before all
				threads have finished. In this case it is assumed that the 
				operation is lengthy, and frequent progress updates are output 
				to std::cout.*/
	template<class CoordinateToDataVectorFunctor>
	void fill(unsigned int nThread, const CoordinateToDataVectorFunctor& rCoordinateToDataVectorFunctor) {
		std::size_t nCell = getCellCount();
		if(nThread == 0) {
			for(std::size_t i = 0; i != nCell; ++i) {
				m_data[i] = rCoordinateToDataVectorFunctor(getDataCoordinate(unravelIndex(i)));
			}
		}
		else {
			const std::size_t nCellPerThread = (nCell + nThread - 1) / nThread;
			std::vector<std::unique_ptr<std::thread>> threadPool;
			std::atomic<std::size_t> nProcessedSlice(0);
			for(std::size_t i = 0; i != nThread; ++i) {
				std::size_t cellBegin = i * nCellPerThread;
				std::size_t cellEnd = std::min(nCell, (i + 1) * nCellPerThread);
				threadPool.push_back(std::unique_ptr<std::thread>(new std::thread(
					[this, &nProcessedSlice, cellBegin, cellEnd, &rCoordinateToDataVectorFunctor]()
				{
					for(std::size_t i = cellBegin; i != cellEnd; ++i) {
						m_data[i] = rCoordinateToDataVectorFunctor(getDataCoordinate(unravelIndex(i)));
						if((i + 1) % m_sliceSize == 0) {
							++nProcessedSlice;
						}
					}
				}
				)));
			}
			// Give frequent progress updates
			std::size_t percentageComplete = 0;
			std::size_t nCurrentProcessedSlice;
			while((nCurrentProcessedSlice = nProcessedSlice) < m_resolution[2]) {
				if((percentageComplete + 1) * 0.01f * m_resolution[2] < nCurrentProcessedSlice) {
					std::cout << "Grid " << percentageComplete << "% filled." << std::endl;
					++percentageComplete;
				}
				std::this_thread::sleep_for(std::chrono::milliseconds(10));
			}
			// Wait for all threads to finish
			for(auto& rpThread : threadPool) {
				rpThread->join();
			}
		}
	}

	/*! Provides read-only access to the raw data of the cells. Indices for the
		vector are flat indices.*/
	const std::vector<std::array<float, 3>>& getData() const {
		return m_data;
	}

	/*! Stores this 3D grid to a simple binary file format. Refer to the
		implmentation of this function for details on the file format.
		\note If you pass an std::ofstream, do not forget to initialize it with
			std::ios_base::binary.*/
	void store(std::ostream& rOutputStream) const;

private:
	/*! Updates redundant attributes.*/
	void update();

	/*! The number of voxels along each of the three axes.*/
	std::array<std::size_t, 3> m_resolution;
	/*! The product m_resolution[0]*m_resolution[1].*/
	std::size_t m_sliceSize;
	/*! The minimal and maximal coordinates of the cuboid.*/
	std::array<float, 3> m_cuboidMin, m_cuboidMax;
	/*! Multiplying integer indices by this factor element-wise and adding the
		summand yields the data coordinates.*/
	std::array<float, 3> m_indexToCuboidFactor, m_indexToCuboidSummand;
	/*! The data of all samples in this grid. The data for index x,y,z is found
		at m_data[ravelMultiIndex({x,y,z})].*/
	std::vector<std::array<float, 3>> m_data;
};
