/*
	Copyright (C) 2012  fmaj7b5.info

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 2 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef FM7b5_DEVICE_MEMORY_CUH
#define FM7b5_DEVICE_MEMORY_CUH

#include <stdexcept>
#include <type_traits>
#include <cuda.h>
#include <cuda_runtime.h>

#include "MDView.cuh"

namespace FM7b5
{
	namespace memory
	{
		//! Linear 1-D memory
		template <class T>
		class Linear
		{
		public:
			typedef T value_type;
			typedef typename std::add_pointer<T>::type pointer_type;
			typedef typename std::add_const<pointer_type>::type const_pointer_type;
			typedef typename std::add_reference<T>::type reference_type;

			Linear(const size_t count = 0)
				: m_count(count)
			{
				m_extent = sizeof(value_type);
				md_data = static_cast<pointer_type>(allocate(m_extent * count));
			}

			~Linear()
			{
				deallocate(md_data);
			}

			void* allocate(const size_t byte)
			{
				void* p(nullptr);
				cudaError_t status;

				status = cudaMalloc(&p, byte);
				if (status != cudaSuccess) {
					throw std::bad_alloc(cudaGetErrorString(status));
				}

				return p;
			}

			void deallocate(void* p)
			{
				cudaFree(p);
			}

			void copy_from(const void* const src, const size_t byte)
			{
				cudaMemcpy(md_data, src, byte, cudaMemcpyHostToDevice);
			}

			void copy_to(void* const dst, const size_t byte)
			{
				cudaMemcpy(dst, md_data, byte, cudaMemcpyDeviceToHost);
			}

			size_t size() const { return m_extent * m_count; }
			size_t count() const { return m_count; }
			pointer_type data() const { return md_data; }

			MDView<value_type, 1> ref() const
			{
				return MDView<value_type, 1>(md_data, &m_extent);
			}

		protected:
			size_t m_extent;
			size_t m_count;
			pointer_type md_data;
		};

		//! Linear 2-D memory
		template <class T>
		class LinearPitch
		{
		public:
			typedef T value_type;
			typedef typename std::add_pointer<T>::type pointer_type;
			typedef typename std::add_const<pointer_type>::type const_pointer_type;
			typedef typename std::add_reference<T>::type reference_type;

			LinearPitch(const size_t width, const size_t height)
				: m_width(width), m_height(height)
			{
				md_data = static_cast<pointer_type>(allocate(width, height, &m_extent[0]));
				m_extent[1] = sizeof(value_type);
			}

			~LinearPitch()
			{
				deallocate(md_data);
			}

			void* allocate(const size_t width, const size_t height, size_t* pitch)
			{
				void* p(nullptr);
				cudaError_t status;

				status = cudaMallocPitch(&p, pitch, sizeof(value_type) * width, height);
				if (status != cudaSuccess) {
					throw std::bad_alloc(cudaGetErrorString(status));
				}

				return p;
			}

			void deallocate(void* p)
			{
				cudaFree(p);
			}

			void copy_from(const void* const src, const size_t width, const size_t height, size_t spitch = 0)
			{
				if (spitch < 1) {
					spitch = width;
				}
				cudaMemcpy2D(md_data, pitch(), src, spitch, width, height, cudaMemcpyHostToDevice);
			}

			void copy_to(void* const dst, const size_t width, const size_t height, size_t dpitch = 0)
			{
				if (dpitch < 1) {
					dpitch = width;
				}
				cudaMemcpy2D(dst, dpitch, md_data, pitch(), width, height, cudaMemcpyDeviceToHost);
			}

			size_t size() const { return m_extent[1] * m_height; }
			size_t width() const { return m_width; }
			size_t height() const { return m_height; }
			size_t pitch() const { return m_extent[0]; }
			pointer_type data() const { return md_data; }

			MDView<value_type, 2> ref() const
			{
				return MDView<value_type, 2>(md_data, m_extent);
			}

		protected:
			size_t m_width;
			size_t m_height;
			size_t m_extent[2];
			pointer_type md_data;
		};
	}
}

#endif /* FM7b5_DEVICE_MEMORY_CUH */