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

#pragma once

#include "util.h"
#include "File.h"

template <typename T>
class matrix
{
public:
  typedef enum {None, HostOnly, DeviceOnly, HostAndDevice} SpawnType;

  //Returns whether spawn type includes host
  static bool onHost(SpawnType st) { return !!(st & HostOnly); }
  //Returns whether spawn type includes device
  static bool onDevice(SpawnType st) { return !!(st & DeviceOnly); }

  //Returns the total size of all matrix data residing in host memory
  static bytes totalBytesOnHost() { return s_iTotalHostBytes; }
  //Returns the total size of all matrix data residing in device memory
  static bytes totalBytesOnDevice() { return s_iTotalDeviceBytes; }

  //Creates resizable matrix
  matrix();
  //Creates matrix
  matrix(nat rows, nat columns, SpawnType st = HostOnly);
  //Creates matrix and sets all elements to given value
  matrix(nat rows, nat columns, T value, SpawnType st = HostOnly);
  //Creates matrix from file
  matrix(char* path, SpawnType st = HostOnly);
  //Destroys matrix
  ~matrix();

  //Sets the dimensions of a matrix
  void resize(nat rows, nat columns);
  //Creates matrix array in memory
  void spawn(SpawnType st = HostAndDevice);
  //Destroys matrix array in memory
  void despawn(SpawnType st = HostAndDevice);

  //Returns address on host
  T* hp() const { return m_pHost; }
  //Returns address on device
  T* dp() const { return m_pDevice; }

  //Returns row address on host
  T* hr(nat row) const;
  //Returns row address on device
  T* dr(nat row) const;

  //Host memory accessor
  T& operator()(nat row, nat column, bool trans = false) const;

  //Returns the number of rows
  nat rows() const { return m_iRows; }
  //Returns the number of columns
  nat columns() const { return m_iColumns; }

  //Returns the number of elements
  nat cells() const { return m_iCells; }
  //Returns the size of an element in bytes
  bytes cellSize() const { return m_iCellSize; }

  //Returns the leading dimension of the transpose in device memory (pitch / cell size)
  nat ldt() const { return m_iLeadingDimensionTr; }   //device only
  //Returns the pitch in device memory in bytes
  bytes pitch() const { return m_iPitch; }            //device only
  //Returns the width in device memory in bytes
  bytes width() const { return m_iWidth; }            //device only

  //Displays matrix info
  void dispInfo() const;
  //Displays matrix data
  void dispData() const;

  //Reads matrix data from file into host data
  void read(char* path, nat rows, nat columns);
  //Writes matrix data to file
  void write(char* path) const;

  //Sets all elements to given value in host and device memory
  void init(T value);
  //Sets all elements to given value in host memory
  void initHost(T value);
  //Sets all elements to given value in device memory
  void initDevice(T value);

  //Copies device data to host
  void syncHostWithDevice();
  //Copies host data to device
  void syncDeviceWithHost();

  //Returns whether host and device data are the same
  bool synced() const;
  //Returns whether host and device data are consistent with each other (approx.) (real only)
  bool consistent() const;

protected:
  static bytes s_iTotalHostBytes;
  static bytes s_iTotalDeviceBytes;

  T* m_pHost;
  T* m_pDevice;

  nat m_iRows;
  nat m_iColumns;

  nat m_iCells;
  bytes m_iCellSize;

  nat m_iLeadingDimensionTr;
  bytes m_iPitch;
  bytes m_iWidth;

private:
  void prepare(nat rows, nat columns, SpawnType st = None);
};

typedef matrix<real> Matrix;
typedef matrix<nat> Natrix;

//Computes transA(A(0 .. heightA - 1, 0 .. widthA - 1)) x transB(B(0 .. heightB - 1, 0 .. widthB - 1)) and stores the result in C(0 .. height(transA(A)), 0 .. width(transB(B)) on host
void matrix_mult_h(const Matrix& A, nat heightA, nat widthA, const Matrix& B, nat heightB, nat widthB, Matrix& C, bool transA = false, bool transB = false);

//Computes transA(A(0 .. heightA - 1, 0 .. widthA - 1)) x transB(B(0 .. heightB - 1, 0 .. widthB - 1)) and stores the result in C(0 .. height(transA(A)), 0 .. width(transB(B)) on device
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 = false, bool transB = false);

//Computes transA(A) x transB(B) and stores the result in C on host
void matrix_mult_h(const Matrix& A, const Matrix& B, Matrix& C, bool transA = false, bool transB = false);

//Computes transA(A) x transB(B) and stores the result in C on device
void matrix_mult_d(cublasHandle_t& handle, const Matrix& A, const Matrix& B, Matrix& C, bool transA = false, bool transB = false);
