#pragma once

#include "Matrix.h"

#define RANGE(i, from, to)                                                     \
	int i = from;                                                              \
	i < to;                                                                    \
	i += 1

namespace linalg {

namespace kernel {

//
// step 10
// CUDA kernel add
//
template <typename T>
__device__ void add(const Matrix<T>* a, const Matrix<T>* b, Matrix<T>* res) {
	auto x = (blockIdx.x * blockDim.x) + threadIdx.x;
	auto y = (blockIdx.y * blockDim.y) + threadIdx.y;
	if (x >= res->cols())
		return;
	if (y >= res->rows())
		return;
	auto a_ref = (const Matrix<T>&)(a);
	auto b_ref = (const Matrix<T>&)(b);
	auto res_ref = (Matrix<T>&)(res);
	auto res_ptr = &res_ref(x, y);
	*res_ptr = *(&a_ref(x, y)) + *b_ref(x, y);
}

//
// step 12
// CUDA kernel apply
//

} // namespace kernel

template <typename T>
__host__ Matrix<T>::Matrix(int rows, int cols)
	: m_data_ptr(nullptr), m_rows(rows), m_cols(cols), m_pitch(0) {
	auto line_width = cols * sizeof(T);
	// step 07
	cudaMallocPitch(&this->m_data_ptr, &this->m_pitch, //
					line_width, rows				   //
	);
}

template <typename T> __host__ void Matrix<T>::free() {
	// step 07
	cudaFree(this->m_data_ptr);
}

template <typename T>
__host__ void Matrix<T>::to_cuda(const std::vector<T>& values) {
	// step 08
	auto vec_line_width = this->m_cols * sizeof(T);
	auto vec_arr_ptr = &values.front();
	cudaMemcpy2D(this->m_data_ptr, this->m_pitch, vec_arr_ptr, vec_line_width,
				 vec_line_width, this->m_rows, cudaMemcpyHostToDevice);
}

template <typename T>
__host__ void Matrix<T>::to_cpu(std::vector<T>& values) const {
	// step 08
	auto vec_line_width = this->m_cols * sizeof(T);
	auto vec_arr_ptr = &values.front();
	cudaMemcpy2D(vec_arr_ptr, vec_line_width, this->m_data_ptr, this->m_pitch,
				 vec_line_width, this->m_rows, cudaMemcpyDeviceToHost);
}

template <typename T>
__device__ const T& Matrix<T>::operator()(int i, int j) const {
	// step 09
	if (i >= this->m_cols)
		return NULL;
	if (j >= this->m_rows)
		return NULL;
	auto offset = (j * this->m_pitch) + (i * sizeof(T));
	auto base_ptr = (u_int8_t*)(this->m_data_ptr);
	auto result = base_ptr + offset;
	return (const T&)(result);
}

template <typename T> __device__ T& Matrix<T>::operator()(int i, int j) {
	// step 09
	// if (i >= this->m_cols)
	//     return nullptr;
	// if (j >= this->m_rows)
	//     return nullptr;
	auto offset = (j * this->m_pitch) + (i * sizeof(T));
	auto base_ptr = (u_int8_t*)(this->m_data_ptr);
	auto result = base_ptr + offset;
	return (T&)(result);
}

template <typename T>
__host__ Matrix<T> Matrix<T>::operator+(const Matrix<T>& other) const {
	// step 11
	auto width = min(this->m_cols, other.m_cols);
	auto height = min(this->m_rows, other.m_rows);
	auto res = Matrix<T>(width, height);
	auto threads_per_block = dim3(32, 32, 1);
	auto blocks = dim3(width / 32 + 1, height / 32 + 1, 1);
	kernel::add<T><<<blocks, threads_per_block>>>(this, &other, &res);
	return res;
}

template <typename T>
__host__ Matrix<T> Matrix<T>::operator-(const Matrix<T>& other) const {
	// step 12
}

template <typename T>
__host__ Matrix<T> Matrix<T>::operator*(const Matrix<T>& other) const {
	// step 12
}

template <typename T>
__host__ Matrix<T> Matrix<T>::operator/(const Matrix<T>& other) const {
	// step 12
}

} // namespace linalg
