/*
 * container_gpu.h
 *
 *  Created on: Jun 9, 2011
 *      Author: ditlevsen
 */

#ifndef CONTAINER_DEVICE_H_
#define CONTAINER_DEVICE_H_

#include "container_host.h"

class mat_device {
private:
	float *pbuf;
	int ld;
public:
	const int dim_x;
	const int dim_y;
	const int dim_z;
	const int len;
	const T_mat_format format;

	mat_device(int l_y, int l_x, int l_z, bool init=true, bool pitch=false, T_mat_format storage=mat_col_major);
	mat_device(int num_elements, bool init=true);
	mat_device(const mat_device &m);
	mat_device(const mat_host &m, bool pitch=false);
	~mat_device() { if(pbuf) HANDLE_ERROR(cudaFree(pbuf)); }

	int leading_dim() const { return ld; }

	float *data_dev_ptr() { return pbuf; }
	const float *data_dev_ptr() const { return pbuf; }

	mat_device &operator=(const mat_device &m) throw(illegal_mat_gpu_assignment);
	mat_device &operator=(const mat_host &m) throw(illegal_mat_gpu_assignment);
};

class sparse_mat_device {
private:
	int *p_ind;
	int *p_ptr;
	float *p_val;
public:
	const int dim_y;
	const int dim_x;
	const int nnz;
	T_sparse_mat_format format;

	sparse_mat_device(int num_dim_y, int num_dim_x, int n_nonzero, T_sparse_mat_format storage_format, bool init=true);
	sparse_mat_device(const sparse_mat_device &m);
	sparse_mat_device(const sparse_mat_host &m);
	~sparse_mat_device();

	float *val() { return p_val; }
	const float *val() const { return p_val; }

	int *ptr() { return p_ptr; }
	const int *ptr() const { return p_ptr; }

	int *ind() { return p_ind; }
	const int *ind() const { return p_ind; }

	sparse_mat_device &operator=(const sparse_mat_device &m) throw(illegal_mat_gpu_assignment);
	sparse_mat_device &operator=(const sparse_mat_host &m) throw(illegal_mat_gpu_assignment);
};

//---------------------------------------
// class geomety_device
//---------------------------------------

class geometry_device {
private:
	int *x_emitters;
	int *y_emitters;
	int *z_emitters;
	int *x_receivers;
	int *y_receivers;
	int *z_receivers;

public:
	int *x_em_dev_ptr() { return x_emitters; }
	const int *x_em_dev_ptr() const { return x_emitters; }
	int *y_em_dev_ptr() { return y_emitters; }
	const int *y_em_dev_ptr() const { return y_emitters; }
	int *z_em_dev_ptr() { return z_emitters; }
	const int *z_em_dev_ptr() const { return z_emitters; }
	int *x_re_dev_ptr() { return x_receivers; }
	const int *x_re_dev_ptr() const { return x_receivers; }
	int *y_re_dev_ptr() { return y_receivers; }
	const int *y_re_dev_ptr() const { return y_receivers; }
	int *z_re_dev_ptr() { return z_receivers; }
	const int *z_re_dev_ptr() const { return z_receivers; }

	const int num_emitters;
	const int num_receivers;
	const int rv_x;
	const int rv_y;
	const int rv_z;
	const float scale_factor;

	geometry_device(const geometry_host &geom_host);
	geometry_device(const geometry_device &geom_dev);
	~geometry_device();
};

// ------------------------------------------------------------------------
// member functions, etc. class mat_device
// ------------------------------------------------------------------------

//member functions...
mat_device::mat_device(int l_y, int l_x, int l_z, bool init, bool pitch, T_mat_format storage): pbuf(NULL), dim_y(l_y),
		dim_z(l_z), dim_x(l_x), len(l_y*l_x*l_z), format(storage) {
	if(len > 0) {
		if(pitch) {
			size_t width = (format == mat_row_major) ? dim_x*sizeof(float) : dim_y*sizeof(float);
			size_t height = (format == mat_row_major) ? dim_y : dim_x;
			size_t p;
			HANDLE_ERROR(cudaMallocPitch(&pbuf, &p, width, height * dim_z));
			ld = p / sizeof(float);
			if(init)
				HANDLE_ERROR(cudaMemset2D(pbuf, p, 0, width, height * dim_z));
		} else {
			HANDLE_ERROR(cudaMalloc(&pbuf, len*sizeof(float)));
			ld = (format == mat_row_major) ? dim_x : dim_y;
			if(init)
				HANDLE_ERROR(cudaMemset(pbuf, 0, len*sizeof(float)));
		}
	}
};

mat_device::mat_device(int num_elements, bool init): pbuf(NULL), dim_y(num_elements), dim_x(1), dim_z(1), len(num_elements),
		format(mat_col_major) {
	if(len > 0) {
		HANDLE_ERROR(cudaMalloc(&pbuf, len*sizeof(float)));
		ld = (format == mat_row_major) ? dim_x : dim_y;
		if(init)
			HANDLE_ERROR(cudaMemset(pbuf, 0, len*sizeof(float)));
	}
};

mat_device::mat_device(const mat_device &m): pbuf(NULL), dim_y(m.dim_y), dim_x(m.dim_x), dim_z(m.dim_z), len(m.len),
		format(m.format), ld(m.ld) {
	if(len > 0) {
		size_t width = (format == mat_row_major) ? dim_x*sizeof(float) : dim_y*sizeof(float);
		size_t height = (format == mat_row_major) ? dim_y : dim_x;
		int total_len = (format == mat_row_major) ? ld*dim_y*dim_z : ld*dim_x*dim_z;
		HANDLE_ERROR(cudaMalloc(&pbuf, total_len*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy2D(pbuf, ld*sizeof(float), m.pbuf, m.ld*sizeof(float), width, height * dim_z, cudaMemcpyDeviceToDevice));
	}
};

mat_device::mat_device(const mat_host &m, bool pitch): pbuf(NULL), dim_y(m.dim_y), dim_x(m.dim_x), dim_z(m.dim_z), len(m.len),
		format(m.format) {
	if(len > 0) {
		if(pitch) {
			size_t width = (format == mat_row_major) ? dim_x*sizeof(float) : dim_y*sizeof(float);
			size_t height = (format == mat_row_major) ? dim_y : dim_x;
			size_t p;

			HANDLE_ERROR(cudaMallocPitch(&pbuf, &p, width, height * dim_z));
			ld = p / sizeof(float);
			HANDLE_ERROR(cudaMemcpy2D(pbuf, ld*sizeof(float), m.data(), width, width, height * dim_z, cudaMemcpyHostToDevice));
		} else {
			HANDLE_ERROR(cudaMalloc(&pbuf, len*sizeof(float)));
			ld = (format == mat_row_major) ? dim_x : dim_y;
			HANDLE_ERROR(cudaMemcpy(pbuf, m.data(), len*sizeof(float), cudaMemcpyHostToDevice));
		}
	}
};

mat_device &mat_device::operator=(const mat_device &m) throw(illegal_mat_gpu_assignment) {
	int w_m = (m.format == mat_row_major) ? m.dim_x : m.dim_y;
	int w = (format == mat_row_major) ? dim_x : dim_y;
	if(w_m == m.ld && w == ld && len == m.len) {
		HANDLE_ERROR(cudaMemcpy(pbuf, m.data_dev_ptr(), len*sizeof(float), cudaMemcpyDeviceToDevice));
	} else if((m.dim_y == dim_y) && (m.dim_x == dim_x) && (m.dim_z == dim_z) && (m.format == format)) {
		size_t width = (format == mat_row_major) ? dim_x*sizeof(float) : dim_y*sizeof(float);
		size_t height = (format == mat_row_major) ? dim_y : dim_x;
		HANDLE_ERROR(cudaMemcpy2D(pbuf, ld*sizeof(float), m.pbuf, m.ld*sizeof(float), width, height * dim_z, cudaMemcpyDeviceToDevice));
	} else {
		throw illegal_mat_gpu_assignment("Illegale Zuweisung von mat_device-Objekten!");
	}
	return *this;
}


mat_device &mat_device::operator=(const mat_host &m) throw(illegal_mat_gpu_assignment) {
	if((m.dim_y == dim_y) && (m.dim_x == dim_x) && (m.dim_z == dim_z) && (m.format == format)) {
		size_t width = (format == mat_row_major) ? dim_x*sizeof(float) : dim_y*sizeof(float);
		size_t height = (format == mat_row_major) ? dim_y : dim_x;
		HANDLE_ERROR(cudaMemcpy2D(pbuf, ld*sizeof(float), m.data(), width, width, height * dim_z, cudaMemcpyHostToDevice));
	} else {
		throw illegal_mat_gpu_assignment("Illegale Zuweisung mat_device <- mat_host!");
	}
	return *this;
}

// function for conversion mat_device -> mat_host
void mat_gpu_to_host(mat_host &m_host, const mat_device &m) throw(illegal_mat_gpu_assignment) {
	if((m.dim_y == m_host.dim_y) && (m.dim_x == m_host.dim_x) && (m.dim_z == m_host.dim_z) && (m.format == m_host.format)) {
		size_t width = (m.format == mat_row_major) ? m.dim_x*sizeof(float) : m.dim_y*sizeof(float);
		size_t height = (m.format == mat_row_major) ? m.dim_y : m.dim_x;
		HANDLE_ERROR(cudaMemcpy2D(m_host.data(), width, m.data_dev_ptr(), m.leading_dim()*sizeof(float), width, height * m.dim_z, cudaMemcpyDeviceToHost));
	} else {
		throw illegal_mat_gpu_assignment("Illegale Zuweisung mat_host <- mat_device!");
	}
};

// output operator...
std::ostream &operator<<(std::ostream &stream, const mat_device &m) {
	mat_host m_host(m.dim_y, m.dim_x, m.dim_z, m.format, false, false);

	mat_gpu_to_host(m_host, m);

	stream << m_host;

	return stream;
}

// ------------------------------------------------------------------------
// member functions, etc. class sparse_mat_device
// ------------------------------------------------------------------------

//member functions...
sparse_mat_device::sparse_mat_device(int num_dim_y, int num_dim_x, int n_nonzero, T_sparse_mat_format storage_format, bool init): p_val(NULL), p_ptr(NULL), p_ind(NULL),
	dim_y(num_dim_y), dim_x(num_dim_x), nnz(n_nonzero), format(storage_format)
{
	if(nnz > 0) {
		int len_ptr = (format == sparse_mat_csc) ? dim_x + 1 : dim_y + 1;

		HANDLE_ERROR(cudaMalloc(&p_val, nnz*sizeof(float)));
		if(init)
			HANDLE_ERROR(cudaMemset(p_val, 0, nnz*sizeof(float)));
		HANDLE_ERROR(cudaMalloc(&p_ptr, len_ptr*sizeof(float)));
		if(init)
			HANDLE_ERROR(cudaMemset(p_ptr, 0, len_ptr*sizeof(float)));
		HANDLE_ERROR(cudaMalloc(&p_ind, nnz*sizeof(float)));
		if(init)
			HANDLE_ERROR(cudaMemset(p_ind, 0, nnz*sizeof(float)));
	}
};
sparse_mat_device::sparse_mat_device(const sparse_mat_device &m): p_val(NULL), p_ptr(NULL), p_ind(NULL), dim_y(m.dim_y),
	dim_x(m.dim_x), nnz(m.nnz), format(m.format)
{
	if(nnz > 0) {
		int len_ptr = (format == sparse_mat_csc) ? dim_x + 1 : dim_y + 1;

		HANDLE_ERROR(cudaMalloc(&p_val, nnz*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(p_val, m.p_val, nnz*sizeof(float), cudaMemcpyDeviceToDevice));
		HANDLE_ERROR(cudaMalloc(&p_ptr, len_ptr*sizeof(int)));
		HANDLE_ERROR(cudaMemcpy(p_ptr, m.p_ptr, len_ptr*sizeof(int), cudaMemcpyDeviceToDevice));
		HANDLE_ERROR(cudaMalloc(&p_ind, nnz*sizeof(int)));
		HANDLE_ERROR(cudaMemcpy(p_ind, m.p_ind, nnz*sizeof(int), cudaMemcpyDeviceToDevice));
	}
};

sparse_mat_device::sparse_mat_device(const sparse_mat_host &m): p_val(NULL), p_ptr(NULL), p_ind(NULL), dim_y(m.dim_y),
	dim_x(m.dim_x), nnz(m.nnz), format(m.format)
{
	if(nnz > 0) {
		int len_ptr = (format == sparse_mat_csc) ? dim_x + 1 : dim_y + 1;

		HANDLE_ERROR(cudaMalloc(&p_val, nnz*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(p_val, m.val(), nnz*sizeof(float), cudaMemcpyHostToDevice));
		HANDLE_ERROR(cudaMalloc(&p_ptr, len_ptr*sizeof(int)));
		HANDLE_ERROR(cudaMemcpy(p_ptr, m.ptr(), len_ptr*sizeof(int), cudaMemcpyHostToDevice));
		HANDLE_ERROR(cudaMalloc(&p_ind, nnz*sizeof(int)));
		HANDLE_ERROR(cudaMemcpy(p_ind, m.ind(), nnz*sizeof(int), cudaMemcpyHostToDevice));
	}
};

sparse_mat_device::~sparse_mat_device() {
	if(p_val) HANDLE_ERROR(cudaFree(p_val));
	if(p_ptr) HANDLE_ERROR(cudaFree(p_ptr));
	if(p_ind) HANDLE_ERROR(cudaFree(p_ind));
}

sparse_mat_device &sparse_mat_device::operator=(const sparse_mat_device &m) throw(illegal_mat_gpu_assignment) {
	if(m.nnz == nnz && m.format == format && m.dim_x == dim_x && m.dim_y == dim_y) {
		int len_ptr = (format == sparse_mat_csc) ? dim_x + 1 : dim_y + 1;

		HANDLE_ERROR(cudaMemcpy(p_val, m.p_val, nnz*sizeof(float), cudaMemcpyDeviceToDevice));
		HANDLE_ERROR(cudaMemcpy(p_ptr, m.p_ptr, len_ptr*sizeof(int), cudaMemcpyDeviceToDevice));
		HANDLE_ERROR(cudaMemcpy(p_ind, m.p_ind, nnz*sizeof(int), cudaMemcpyDeviceToDevice));
	} else {
		throw illegal_mat_gpu_assignment("Illegale Zuweisung von sparse_mat_device-Objekten!");
	}
	return *this;
}

sparse_mat_device &sparse_mat_device::operator=(const sparse_mat_host &m) throw(illegal_mat_gpu_assignment) {
	if(m.nnz == nnz && m.format == format && m.dim_x == dim_x && m.dim_y == dim_y) {
		int len_ptr = (format == sparse_mat_csc) ? dim_x + 1 : dim_y + 1;

		HANDLE_ERROR(cudaMemcpy(p_val, m.val(), nnz*sizeof(float), cudaMemcpyHostToDevice));
		HANDLE_ERROR(cudaMemcpy(p_ptr, m.ptr(), len_ptr*sizeof(int), cudaMemcpyHostToDevice));
		HANDLE_ERROR(cudaMemcpy(p_ind, m.ind(), nnz*sizeof(int), cudaMemcpyHostToDevice));
	} else {
		throw illegal_mat_gpu_assignment(("Illegale Zuweisung sparse_ mat_device <- sparse_mat_host!"));
	}
	return *this;
}

// function for conversion sparse_mat_device -> sparse_mat_host
void sparse_mat_gpu_to_host(sparse_mat_host &lhs, const sparse_mat_device &rhs) throw(illegal_mat_gpu_assignment) {
	if(lhs.nnz == rhs.nnz && lhs.format == rhs.format && lhs.dim_x == rhs.dim_x && lhs.dim_y == rhs.dim_y) {
		int len_ptr = (lhs.format == sparse_mat_csc) ? lhs.dim_x + 1 : lhs.dim_y + 1;

		HANDLE_ERROR(cudaMemcpy(lhs.val(), rhs.val(), lhs.nnz*sizeof(float), cudaMemcpyDeviceToHost));
		HANDLE_ERROR(cudaMemcpy(lhs.ptr(), rhs.ptr(), len_ptr*sizeof(int), cudaMemcpyDeviceToHost));
		HANDLE_ERROR(cudaMemcpy(lhs.ind(), rhs.ind(), lhs.nnz*sizeof(int), cudaMemcpyDeviceToHost));
	} else {
		throw illegal_mat_gpu_assignment(("Illegale Zuweisung sparse_ mat_host <- sparse_mat_device!"));
	}
}

// output operator...
std::ostream &operator<<(std::ostream &stream, const sparse_mat_device &m) {
	sparse_mat_host tmp(m.dim_y, m.dim_x, m.nnz, m.format);

	sparse_mat_gpu_to_host(tmp, m);

	stream << tmp;

	return stream;
}

// ------------------------------------------------------------------------
// member functions, etc. struct geometry
// ------------------------------------------------------------------------

geometry_device::geometry_device(const geometry_host &geom_host):  num_emitters(geom_host.num_emitters),
		num_receivers(geom_host.num_receivers), rv_x(geom_host.rv_x), rv_y(geom_host.rv_y), rv_z(geom_host.rv_z),
		scale_factor(geom_host.scale_factor), x_emitters(NULL), y_emitters(NULL), z_emitters(NULL), x_receivers(NULL),
		y_receivers(NULL), z_receivers(NULL)
{
	if(num_emitters > 0) {
		HANDLE_ERROR(cudaMalloc(&x_emitters, num_emitters*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(x_emitters, geom_host.x_emitters, num_emitters*sizeof(float), cudaMemcpyHostToDevice));
		HANDLE_ERROR(cudaMalloc(&y_emitters, num_emitters*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(y_emitters, geom_host.y_emitters, num_emitters*sizeof(float), cudaMemcpyHostToDevice));
		HANDLE_ERROR(cudaMalloc(&z_emitters, num_emitters*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(z_emitters, geom_host.z_emitters, num_emitters*sizeof(float), cudaMemcpyHostToDevice));
	}
	if(num_receivers > 0) {
		HANDLE_ERROR(cudaMalloc(&x_receivers, num_receivers*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(x_receivers, geom_host.x_receivers, num_receivers*sizeof(float), cudaMemcpyHostToDevice));
		HANDLE_ERROR(cudaMalloc(&y_receivers, num_receivers*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(y_receivers, geom_host.y_receivers, num_receivers*sizeof(float), cudaMemcpyHostToDevice));
		HANDLE_ERROR(cudaMalloc(&z_receivers, num_receivers*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(z_receivers, geom_host.z_receivers, num_receivers*sizeof(float), cudaMemcpyHostToDevice));
	}
}

geometry_device::geometry_device(const geometry_device &geom_dev):  num_emitters(geom_dev.num_emitters),
		num_receivers(geom_dev.num_receivers), rv_x(geom_dev.rv_x), rv_y(geom_dev.rv_y), rv_z(geom_dev.rv_z),
		scale_factor(geom_dev.scale_factor), x_emitters(NULL), y_emitters(NULL), z_emitters(NULL), x_receivers(NULL),
		y_receivers(NULL), z_receivers(NULL)
{
	if(num_emitters > 0) {
		HANDLE_ERROR(cudaMalloc(&x_emitters, num_emitters*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(x_emitters, geom_dev.x_emitters, num_emitters*sizeof(float), cudaMemcpyDeviceToDevice));
		HANDLE_ERROR(cudaMalloc(&y_emitters, num_emitters*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(y_emitters, geom_dev.y_emitters, num_emitters*sizeof(float), cudaMemcpyDeviceToDevice));
		HANDLE_ERROR(cudaMalloc(&z_emitters, num_emitters*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(z_emitters, geom_dev.z_emitters, num_emitters*sizeof(float), cudaMemcpyDeviceToDevice));
	}
	if(num_receivers > 0) {
		HANDLE_ERROR(cudaMalloc(&x_receivers, num_receivers*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(x_receivers, geom_dev.x_receivers, num_receivers*sizeof(float), cudaMemcpyDeviceToDevice));
		HANDLE_ERROR(cudaMalloc(&y_receivers, num_receivers*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(y_receivers, geom_dev.y_receivers, num_receivers*sizeof(float), cudaMemcpyDeviceToDevice));
		HANDLE_ERROR(cudaMalloc(&z_receivers, num_receivers*sizeof(float)));
		HANDLE_ERROR(cudaMemcpy(z_receivers, geom_dev.z_receivers, num_receivers*sizeof(float), cudaMemcpyDeviceToDevice));
	}
}

geometry_device::~geometry_device() {
	if(x_emitters) cudaFree(x_emitters);
	if(y_emitters) cudaFree(y_emitters);
	if(x_receivers) cudaFree(x_receivers);
	if(y_receivers) cudaFree(y_receivers);
}


#endif /* CONTAINER_DEVICE_H_ */