/*
 * Classes for handling matrices and vectors:
 *
 * Class mat<Type>:         dense matrices/vectors
 * Class sparse_mat<Type>:  sparse matrices
 * Class geometry:         data for dynamic calculation of the measurement matrix
 *
 */

#ifndef CONTAINER_H_
#define CONTAINER_H_


#include <cstring>
#include <cstdlib>
#include <iostream>
#include <exception>

// enumerations...
enum T_alloc {allocated_in_constructor, preallocated_buffer};
enum T_mat_format {mat_col_major, mat_row_major};
enum T_sparse_mat_format {sparse_mat_csc, sparse_mat_csr};

// exceptions...
class illegal_mat_assignment: public std::exception
{
  virtual const char* what() const throw()
  {
    return "Illegale Zuweisung von mat-Objekten.";
  }
};

class illegal_sparse_mat_assignment: public std::exception
{
  virtual const char* what() const throw()
  {
    return "Illegale Zuweisung von sparse_mat-Objekten.";
  }
};

//---------------------------------------
// class mat...
//---------------------------------------

template <class Type>
class mat {
private:
	Type *pbuf;
	T_alloc alloc_info;
public:
	const int dim_y;
	const int dim_x;
	const int dim_z;
	const int len;
	const T_mat_format format;

	mat(int l_y, int l_x, int l_z, bool init=true, T_mat_format storage=mat_row_major);
	mat(int num_elements, bool init=true);
	mat(int l_y, int l_x, int l_z, Type *buffer, T_mat_format storage=mat_row_major);
	mat(int num_elements, Type *buffer);
	mat(const mat &m);
	~mat() { if(pbuf && (alloc_info == allocated_in_constructor)) delete[] pbuf; };

	Type *data() { return pbuf; }
	const Type *data() const { return pbuf; }

	Type &operator[](int i) { return pbuf[i]; }
	const Type &operator[](int i) const { return pbuf[i]; }

	mat &operator=(const mat<Type> &m) throw(illegal_mat_assignment);

	void randomFill(Type scl);
};


//---------------------------------------
// class sparse_mat...
//---------------------------------------

template<class Type>
class sparse_mat {
private:
	int *p_ptr;
	int *p_ind;
	Type *p_val;
	T_alloc alloc_info;
public:
	const int dim_y;
	const int dim_x;
	const int nnz;
	const T_sparse_mat_format format;

	sparse_mat(int l_y, int l_x, int n_nonzero, T_sparse_mat_format storage_format=sparse_mat_csc, bool init=true);
	sparse_mat(int l_y, int l_x, int n_nonzero, Type *buff_val, int *buff_ptr, int *buff_ind,
			T_sparse_mat_format storage_format=sparse_mat_csc);
	sparse_mat(const sparse_mat &m);
	~sparse_mat();

	Type *val() { return p_val; }
	const Type *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<Type> &operator=(const mat<Type> &m) throw(illegal_sparse_mat_assignment);
};

//---------------------------------------
// Struct geomety
//---------------------------------------

struct geometry {
	int *x_emitters;
	int *y_emitters;
	int *z_emitters;
	int *x_receivers;
	int *y_receivers;
	int *z_receivers;
	int num_emitters;
	int num_receivers;
	int rv_x;
	int rv_y;
	int rv_z;
	float scale_factor;
};

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

template<class Type>
mat<Type>::mat(int l_y, int l_x, int l_z, bool init, T_mat_format storage): pbuf(NULL), alloc_info(allocated_in_constructor),
	dim_y(l_y), dim_x(l_x), dim_z(l_z), len(l_y*l_x*l_z), format(storage) {
	if(len > 0) {
		pbuf = new Type[len];
		if(init)
			memset(pbuf, 0, len*sizeof(Type));
	}
};

template<class Type>
mat<Type>::mat(int num_elements, bool init): pbuf(NULL), alloc_info(allocated_in_constructor), dim_y(num_elements), dim_x(1),
	dim_z(1), len(num_elements), format(mat_row_major) {
	if(len > 0) {
		pbuf = new Type[len];
		if(init)
			memset(pbuf, 0, len*sizeof(Type));
	}
};

template<class Type>
mat<Type>::mat(int l_y, int l_x, int l_z, Type *buffer, T_mat_format storage): pbuf(buffer), dim_y(l_y), dim_x(l_x),
	dim_z(l_z), len(l_y*l_x*l_z), alloc_info(preallocated_buffer), format(storage) {};

template<class Type>
mat<Type>::mat(int num_elements, Type *buffer): pbuf(buffer), dim_y(num_elements), dim_x(1),
	dim_z(1), len(num_elements), alloc_info(preallocated_buffer), format(mat_row_major) {};

template<class Type>
mat<Type>::mat(const mat &m): pbuf(NULL), alloc_info(allocated_in_constructor), dim_y(m.dim_y), dim_x(m.dim_x),
	dim_z(m.dim_z), len(m.len), format(m.format) {
	if(len > 0) {
		pbuf = new Type[len];
		memcpy(pbuf, m.pbuf, len*sizeof(Type));
	}
};

template<class Type>
mat<Type> &mat<Type>::operator=(const mat<Type> &m) throw(illegal_mat_assignment) {
	if(m.len == len)
		memcpy(pbuf, m.pbuf, len*sizeof(Type));
	else
		throw illegal_mat_assignment();
	return *this;
}

template<class Type>
void mat<Type>::randomFill(Type scl) {
	for(int i=0; i<len; i++)
		pbuf[i] = ((Type)rand() / RAND_MAX - 0.5)/scl;
}

template<class Type>
std::ostream &operator<<(std::ostream &stream, const mat<Type> &m) {
	for(int z=0; z < m.dim_z; z++) {
		stream << "\nz=" << z+1 << ":\n\n";
		for(int i=0; i<m.dim_y; i++) {
			for(int j=0; j<m.dim_x; j++) {
				if(m.format == mat_row_major)
					stream << m[z*m.dim_x*m.dim_y + i*m.dim_x + j] << " ";
				else
					stream << m[z*m.dim_x*m.dim_y + i + j*m.dim_y] << " ";
			}
			stream << "\n";
		}
		return stream;
	}
}


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

template<class Type>
sparse_mat<Type>::sparse_mat(int l_y, int l_x, int n_nonzero, T_sparse_mat_format storage_format, bool init):
	p_val(NULL), p_ptr(NULL), p_ind(NULL), dim_y(l_y), dim_x(l_x), nnz(n_nonzero), format(storage_format),
	alloc_info(allocated_in_constructor)
{
	if(nnz > 0) {
		int len_ptr = (format == sparse_mat_csc) ? dim_x + 1 : dim_y + 1;

		p_val = new Type[nnz];
		p_ptr = new int[len_ptr];
		p_ind = new int[nnz];
		if(init) {
			memset(p_val, 0, nnz*sizeof(Type));
			memset(p_ptr, 0, (len_ptr)*sizeof(int));
			memset(p_ind, 0, nnz*sizeof(int));
		}
	}
};

template<class Type>
sparse_mat<Type>::sparse_mat(int l_y, int l_x, int n_nonzero, Type *buff_val, int *buff_ptr, int *buff_ind,
		T_sparse_mat_format storage_format): p_val(buff_val), p_ptr(buff_ptr), p_ind(buff_ind), dim_y(l_y), dim_x(l_x),
		nnz(n_nonzero), format(storage_format), alloc_info(preallocated_buffer) {}


template<class Type>
sparse_mat<Type>::sparse_mat(const sparse_mat &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), alloc_info(allocated_in_constructor)
{
	if(nnz > 0) {
		int len_ptr = (format == sparse_mat_csc) ? dim_x + 1 : dim_y + 1;

		p_val = new Type[nnz];
		memcpy(p_val, m.p_val, nnz*sizeof(Type));
		p_ptr = new int[len_ptr];
		memcpy(p_ptr, m.p_ptr, (len_ptr)*sizeof(int));
		p_ind = new int[nnz];
		memcpy(p_ind, m.p_ind, nnz*sizeof(int));
	}
};

template<class Type>
sparse_mat<Type>::~sparse_mat() {
	if(alloc_info == allocated_in_constructor) {
		if(p_val) delete[] p_val;
		if(p_ptr) delete[] p_ptr;
		if(p_ind) delete[] p_ind;
	}
}

template<class Type>
sparse_mat<Type> &sparse_mat<Type>::operator=(const mat<Type> &m) throw(illegal_sparse_mat_assignment){

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

		memcpy(p_val, m.p_val, nnz*sizeof(Type));
		memcpy(p_ptr, m.p_ptr, (len_ptr)*sizeof(int));
		memcpy(p_ind, m.p_ind, nnz*sizeof(int));
	} else {
		throw illegal_sparse_mat_assignment();
	}

	return *this;
}

// output operator ...
template<class Type>
std::ostream &operator<<(std::ostream &stream, const sparse_mat<Type> &m) {
	mat<Type> tmp(m.dim_y, m.dim_x);
	if(m.format == sparse_mat_csc) {
		for(int c=0; c < m.dim_x; c++)
			for(int i = m.ptr()[c]; i < m.ptr()[c+1]; i++)
				tmp[c*m.dim_y + m.ind()[i]] = m.val()[i];
	} else {
		for(int r=0; r < m.dim_y; r++)
			for(int i = m.ptr()[r]; i < m.ptr()[r+1]; i++)
				tmp[m.ind()[i]*m.dim_y + r] = m.val()[i];
	}
	stream << tmp;
	return stream;
}

#endif /* CONTAINER_H_ */