/* ========================================================================== *
 *  Matrix class for host-device operation.
 *
 *  Author: Firas Safadi (fsafadi@ulg.ac.be)
 * ========================================================================== */

#include "Matrix.h"

#define MAX_CELLS_PER_THREAD 4
#define MAX_THREADS_PER_BLOCK 256

#define REAL_OUTPUT_WIDTH 16
#define REAL_OUTPUT_PRECISION 4

#ifndef DOUBLE_PRECISION
#define PRECISION_BITS 15
#define CUBLAS_GEMM cublasSgemm
#else
#define PRECISION_BITS 40
#define CUBLAS_GEMM cublasDgemm
#endif

template <typename T>
void matrix<T>::prepare(nat rows, nat columns, SpawnType st)
{
  m_pHost = NULL;
  m_pDevice = NULL;

  m_iRows = rows;
  m_iColumns = columns;

  m_iCells = m_iRows*m_iColumns;
  m_iCellSize = sizeof(T);

  m_iWidth = m_iColumns*m_iCellSize;
  m_iPitch = m_iWidth;
  m_iLeadingDimensionTr = m_iColumns;

  spawn(st);
}

template <typename T>
matrix<T>::matrix()
{
  prepare(0, 0);
}

template <typename T>
matrix<T>::matrix(nat rows, nat columns, SpawnType st)
{
  prepare(rows, columns, st);
}

template <typename T>
matrix<T>::matrix(nat rows, nat columns, T value, SpawnType st)
{
  prepare(rows, columns, st);

  if (m_pHost)
    initHost(value);

  if (m_pDevice)
    initDevice(value);
}

template <typename T>
matrix<T>::matrix(char* path, SpawnType st)
{
  File matrix_file(path);

  nat rows, columns;

  matrix_file >> rows;
  matrix_file >> columns;

  prepare(rows, columns, st);

  if (m_pHost) {
    for (nat i = 0; i < m_iCells; i++)
      matrix_file >> m_pHost[i];

    if (m_pDevice)
      syncDeviceWithHost();
  }

  else if (m_pDevice) {
    T* arr;

    cuda_call(cudaMallocHost(&arr, m_iCells*m_iCellSize));

    for (nat i = 0; i < m_iCells; i++)
      matrix_file >> arr[i];

    cuda_call(cudaMemcpy2D(m_pDevice, m_iPitch, arr, m_iWidth, m_iWidth, m_iRows, cudaMemcpyHostToDevice));

    cuda_call(cudaFreeHost(arr));
  }

  matrix_file.close();
}

template <typename T>
matrix<T>::~matrix()
{
  despawn(HostAndDevice);
}

template <typename T>
void matrix<T>::resize(nat rows, nat columns)
{
  if (m_pHost || m_pDevice)
    error("Cannot resize spawned matrix");

  prepare(rows, columns);
}

template <typename T>
void matrix<T>::spawn(SpawnType st)
{
  if (st != None && m_iCells == 0)
    error("Cannot spawn zero dim matrix");

  if (onHost(st)) {
    if (m_pHost)
      error("Matrix already spawned on host");

    cuda_call(cudaMallocHost(&m_pHost, m_iCells*m_iCellSize));
    s_iTotalHostBytes += m_iCells*m_iCellSize;
  }

  if (onDevice(st)) {
    if (m_pDevice)
      error("Matrix already spawned on device");

    cuda_call(cudaMallocPitch(&m_pDevice, &m_iPitch, m_iWidth, m_iRows));
    m_iLeadingDimensionTr = m_iPitch / m_iCellSize;
    s_iTotalDeviceBytes += m_iRows*m_iPitch;

    if (m_iLeadingDimensionTr*m_iCellSize != m_iPitch)
      error("Cannot create matrix: pitch not a multiple of cell size");
  }
}

template <typename T>
void matrix<T>::despawn(SpawnType st)
{
  if (onHost(st) && m_pHost) {
    s_iTotalHostBytes -= m_iCells*m_iCellSize;
    cuda_call(cudaFreeHost(m_pHost));
    m_pHost = NULL;
  }

  if (onDevice(st) && m_pDevice) {
    s_iTotalDeviceBytes -= m_iRows*m_iPitch;
    cuda_call(cudaFree(m_pDevice));
    m_pDevice = NULL;
  }
}

template <typename T>
T* matrix<T>::hr(nat row) const
{
  if (!m_pHost)
    error("No matrix on host to access");

  if (row >= m_iRows)
    error("Matrix index out of range");
  
  return m_pHost + row*m_iColumns;
}

template <typename T>
T* matrix<T>::dr(nat row) const
{
  if (!m_pDevice)
    error("No matrix on device to access");

  if (row >= m_iRows)
    error("Matrix index out of range");

  return m_pDevice + row*m_iLeadingDimensionTr;
}

template <typename T>
T& matrix<T>::operator()(nat row, nat column, bool trans) const
{
  if (!m_pHost)
    error("No matrix on host to access");

  if (trans)
    std::swap(row, column);

  if (row >= m_iRows || column >= m_iColumns)
    error("Matrix index out of range");

  return m_pHost[row*m_iColumns + column];
}

template <typename T>
void matrix<T>::dispInfo() const
{
  char buf[BUF_SIZE];

  std::cout << std::endl;
  std::cout << "Address on host: 0x" << m_pHost << std::endl;
  std::cout << "Address on device: 0x" << m_pDevice << std::endl;

  std::cout << std::endl;
  std::cout << "Dimensions: " << m_iRows << "x" << m_iColumns << std::endl;

  std::cout << std::endl;
  std::cout << "Cell size: " << m_iCellSize << " bytes" << std::endl;

  std::cout << std::endl;
  std::cout << "Size on host: " << size_to_string(buf, m_pHost ? m_iCells*m_iCellSize : 0) << std::endl;
  std::cout << "Size on device: " << size_to_string(buf, m_pDevice ? m_iRows*m_iPitch : 0) << std::endl;

  std::cout << std::endl;
  std::cout << "Width on device: " << size_to_string(buf, m_pDevice ? m_iWidth : 0) << std::endl;
  std::cout << "Pitch on device: " << size_to_string(buf, m_pDevice ? m_iPitch : 0) << std::endl;
}

template <typename T>
void matrix<T>::dispData() const
{
  nat pre;

  if (typeid(T) == typeid(real)) {
    pre = std::cout.precision();
    std::cout.setf(std::ios::scientific);
    std::cout << std::setprecision(REAL_OUTPUT_PRECISION);
  }

  std::cout << std::endl;
  std::cout << "Host data:" << std::endl;

  if (m_pHost) {
    for (nat i = 0; i < m_iRows; i++) {
      for (nat j = 0; j < m_iColumns; j++)
        std::cout << std::setw(REAL_OUTPUT_WIDTH) << m_pHost[i*m_iColumns + j];
      std::cout << std::endl;
    }
  }

  else {
    std::cout << "None" << std::endl;
  }

  std::cout << std::endl;
  std::cout << "Device data:" << std::endl;

  if (m_pDevice) {
    T* arr;

    cuda_call(cudaMallocHost(&arr, m_iCells*m_iCellSize));

    cuda_call(cudaMemcpy2D(arr, m_iWidth, m_pDevice, m_iPitch, m_iWidth, m_iRows, cudaMemcpyDeviceToHost));

    for (nat i = 0; i < m_iRows; i++) {
      for (nat j = 0; j < m_iColumns; j++)
        std::cout << std::setw(REAL_OUTPUT_WIDTH) << arr[i*m_iColumns + j];
      std::cout << std::endl;
    }

    cuda_call(cudaFreeHost(arr));
  }

  else {
    std::cout << "None" << std::endl;
  }

  if (typeid(T) == typeid(real)) {
    std::cout << std::setprecision(pre);
    std::cout.unsetf(std::ios::scientific);
  }
}

template <typename T>
void matrix<T>::read(char* path, nat rows, nat columns)
{
  File f(path);

  if (!m_pHost)
    error("No matrix on host to read into");

  if (rows == 0 || columns == 0)
    error("Cannot read zero dim matrix");

  if (rows > m_iRows || columns > m_iColumns)
    error("File data does not fit into matrix");

  for (nat i = 0; i < rows; i++)
    for (nat j = 0; j < columns; j++)
      f >> m_pHost[i*m_iColumns + j];

  f.close();
}

template <typename T>
void matrix<T>::write(char* path) const
{
  File f(path, std::ios_base::out);

  if (typeid(T) == typeid(real))
    f << std::scientific << std::setprecision(REAL_OUTPUT_PRECISION);

  f << File::endl;
  f << "Host data:" << File::endl;

  if (m_pHost) {
    for (nat i = 0; i < m_iRows; i++) {
      for (nat j = 0; j < m_iColumns; j++)
        f << std::setw(REAL_OUTPUT_WIDTH) << m_pHost[i*m_iColumns + j];
      f << File::endl;
    }
  }

  else {
    f << "None" << File::endl;
  }

  f << File::endl;
  f << "Device data:" << File::endl;

  if (m_pDevice) {
    T* arr;

    cuda_call(cudaMallocHost(&arr, m_iCells*m_iCellSize));

    cuda_call(cudaMemcpy2D(arr, m_iWidth, m_pDevice, m_iPitch, m_iWidth, m_iRows, cudaMemcpyDeviceToHost));

    for (nat i = 0; i < m_iRows; i++) {
      for (nat j = 0; j < m_iColumns; j++)
        f << std::setw(REAL_OUTPUT_WIDTH) << arr[i*m_iColumns + j];
      f << File::endl;
    }

    cuda_call(cudaFreeHost(arr));
  }

  else {
    f << "None" << File::endl;
  }

  f.close();
}

template <typename T>
void matrix<T>::init(T value)
{
  initHost(value);
  initDevice(value);
}

template <typename T>
void matrix<T>::initHost(T value)
{
  if (!m_pHost)
    error("No matrix on host to initialize");

  std::fill(m_pHost, m_pHost + m_iCells, value);
}

template <typename T>
__global__ void init_matrix_K(T* M, nat rows, nat columns, nat ldt, T value)
{
  nat x = threadIdx.x + blockIdx.x*blockDim.x;

  if (x < columns) {
    nat y = blockIdx.y*MAX_CELLS_PER_THREAD;

    T* end = M + MIN(rows, y + MAX_CELLS_PER_THREAD)*ldt + x;

    for (M += y*ldt + x; M < end; M += ldt)
      *M = value;
  }
}

template <typename T>
void matrix<T>::initDevice(T value)
{
  if (!m_pDevice)
    error("No matrix on device to initialize");

  dim3 grid(RDIV(m_iColumns, MAX_THREADS_PER_BLOCK), RDIV(m_iRows, MAX_CELLS_PER_THREAD));

  AsyncKernelLaunch(init_matrix_K, grid, MAX_THREADS_PER_BLOCK, m_pDevice, m_iRows, m_iColumns, m_iLeadingDimensionTr, value);
}

template <typename T>
void matrix<T>::syncHostWithDevice()
{
  if (!m_pHost)
    error("No matrix on host to copy to");

  if (!m_pDevice)
    error("No matrix on device to copy from");

  cuda_call(cudaMemcpy2D(m_pHost, m_iWidth, m_pDevice, m_iPitch, m_iWidth, m_iRows, cudaMemcpyDeviceToHost));
}

template <typename T>
void matrix<T>::syncDeviceWithHost()
{
  if (!m_pHost)
    error("No matrix on host to copy from");

  if (!m_pDevice)
    error("No matrix on device to copy to");

  cuda_call(cudaMemcpy2D(m_pDevice, m_iPitch, m_pHost, m_iWidth, m_iWidth, m_iRows, cudaMemcpyHostToDevice));
}

template <typename T>
bool matrix<T>::synced() const
{
  if (!m_pHost || !m_pDevice)
    return false;

  T* arr;

  cuda_call(cudaMallocHost(&arr, m_iCells*m_iCellSize));

  cuda_call(cudaMemcpy2D(arr, m_iWidth, m_pDevice, m_iPitch, m_iWidth, m_iRows, cudaMemcpyDeviceToHost));

  bool res = !std::memcmp(m_pHost, arr, m_iCells*m_iCellSize);

  cuda_call(cudaFreeHost(arr));

  return res;
}

template <typename T>
bool matrix<T>::consistent() const
{
  if (typeid(T) != typeid(real))
    error("Consistency test requires real types");

  if (!m_pHost || !m_pDevice)
    return false;

  real* arr;

  cuda_call(cudaMallocHost(&arr, m_iCells*m_iCellSize));

  cuda_call(cudaMemcpy2D(arr, m_iWidth, m_pDevice, m_iPitch, m_iWidth, m_iRows, cudaMemcpyDeviceToHost));

  nat i;

  for (i = 0; i < m_iCells; i++)
    if (!approx_equal(m_pHost[i], arr[i], PRECISION_BITS))
      break;

  cuda_call(cudaFreeHost(arr));

  return i == m_iCells;
}

bytes matrix<real>::s_iTotalHostBytes = 0;
bytes matrix<real>::s_iTotalDeviceBytes = 0;

bytes matrix<nat>::s_iTotalHostBytes = 0;
bytes matrix<nat>::s_iTotalDeviceBytes = 0;

template class matrix<real>;
template class matrix<nat>;

static void validate_matrix_mult(const Matrix& A, nat heightA, nat widthA, const Matrix& B, nat heightB, nat widthB, Matrix& C, bool transA, bool transB)
{
  if (heightA > A.rows() || widthA > A.columns())
    error("Matrix multiplication: specified dimensions do not define a submatrix of A");

  if (heightB > B.rows() || widthB > B.columns())
    error("Matrix multiplication: specified dimensions do not define a submatrix of B");

  if (transA)
    std::swap(heightA, widthA);

  if (transB)
    std::swap(heightB, widthB);

  if (widthA != heightB)
    error("Matrix multiplication: inner dimensions must match");

  if (C.rows() < heightA || C.columns() < widthB)
    error("Matrix multiplication: result does not fit in C");
}

void matrix_mult_h(const Matrix& A, nat heightA, nat widthA, const Matrix& B, nat heightB, nat widthB, Matrix& C, bool transA, bool transB)
{
  if (!A.hp() || !B.hp() || !C.hp())
    error("Matrix multiplication: valid host pointers required for A, B and C");

  validate_matrix_mult(A, heightA, widthA, B, heightB, widthB, C, transA, transB);

  if (transA)
    std::swap(heightA, widthA);

  if (transB)
    std::swap(heightB, widthB);

  for (nat i = 0; i < heightA; i++) {
    for (nat j = 0; j < widthB; j++) {
      real sum = 0.0;
      for (nat k = 0; k < widthA; k++)
        sum += A(i, k, transA)*B(k, j, transB);
      C(i, j) = sum;
    }
  }
}

void matrix_mult_d(cublasHandle_t& handle, const Matrix& A, nat heightA, nat widthA, const Matrix& B, nat heightB, nat widthB, Matrix& C, bool transA, bool transB)
{
  if (!A.dp() || !B.dp() || !C.dp())
    error("Matrix multiplication: valid device pointers required for A, B and C");

  validate_matrix_mult(A, heightA, widthA, B, heightB, widthB, C, transA, transB);

  cublasOperation_t opA = CUBLAS_OP_N;
  cublasOperation_t opB = CUBLAS_OP_N;

  if (transA) {
    std::swap(heightA, widthA);
    opA = CUBLAS_OP_T;
  }

  if (transB) {
    std::swap(heightB, widthB);
    opB = CUBLAS_OP_T;
  }

  real alpha = 1.0, beta = 0.0;

  //Because BLAS works in column-major order, the following call actually computes C^T = B^T x A^T (note the inverted dimensions)
  cublas_call(CUBLAS_GEMM(handle, opB, opA, widthB, heightA, widthA, &alpha, B.dp(), B.ldt(), A.dp(), A.ldt(), &beta, C.dp(), C.ldt()));
}

void matrix_mult_h(const Matrix& A, const Matrix& B, Matrix& C, bool transA, bool transB)
{
  matrix_mult_h(A, A.rows(), A.columns(), B, B.rows(), B.columns(), C, transA, transB);
}

void matrix_mult_d(cublasHandle_t& handle, const Matrix& A, const Matrix& B, Matrix& C, bool transA, bool transB)
{
  matrix_mult_d(handle, A, A.rows(), A.columns(), B, B.rows(), B.columns(), C, transA, transB);
}
