#ifndef TP_MEMORY_H
#define TP_MEMORY_H

#include <stdio.h>
#include <stdlib.h>

#include "tp_utils.hpp"
#include "tp_memory_interface.hpp"

#include <pmi.h>

class Memory: public iMemory {
public:
  /**********************/
  /*  |-- Allocation    */
  /**********************/
  void memAlloc ( int64_t buffSize, mem_t mem, int masterRank, char* fileName, MPI_Comm comm ) {
    int rank;
    
    this->mem_        = mem;
    this->masterRank_ = masterRank;
    this->comm_       = comm;
    this->buffSize_   = buffSize;
    strcpy ( this->fileName_, fileName);
    
    MPI_Comm_rank ( this->comm_, &rank );

    switch ( this->mem_ )
      {
      case DDR:
	if ( rank == this->masterRank_ ) {
	  this->buffer_ = malloc ( this->buffSize_ );
	  MPI_Win_create ( this->buffer_, this->buffSize_, 1, MPI_INFO_NULL, this->comm_, &this->RMAWin_ );
	}
	else
	  MPI_Win_create ( NULL, 0, 1, MPI_INFO_NULL, this->comm_, &this->RMAWin_ );
	
	MPI_Win_fence (0, this->RMAWin_ );
	break;
      default:
	printMsg ( ERROR, "Unable to allocate memory (mem = %s)\n", this->memName () );
	MPI_Abort ( MPI_COMM_WORLD, -1 );
      }
  }

  void memFree  ( ) {
    int rank;
    MPI_Comm_rank ( MPI_COMM_WORLD, &rank );
    
    switch ( this->mem_ )
      {
      case DDR:
	MPI_Win_free ( &this->RMAWin_ );
	if ( rank == this->masterRank_ )
	  free ( this->buffer_ );
	break;
      default:
	printMsg ( ERROR, "Unable to free memory (mem = %s)\n", this->memName () );
	MPI_Abort ( MPI_COMM_WORLD, -1 );
      }
  }

  /**********************/
  /*  |-- I/O           */
  /**********************/
  int memWrite ( void* srcBuffer, int64_t srcSize, int64_t offset ) {
    int err;
    MPI_Status status;

    switch ( this->mem_ )
      {
      case DDR:
	err = MPI_Put ( srcBuffer, srcSize, MPI_BYTE, 0, offset, srcSize, MPI_BYTE, this->RMAWin_ );
	break;
      default:
	printMsg ( ERROR, "Error while writing data (mem = %s)\n", this->memName () );
	MPI_Abort ( MPI_COMM_WORLD, -1 );
      }

    return err;
  }

  int memRead  ( void* srcBuffer, int64_t srcSize, int64_t offset ) {
    int err;
    MPI_Status status;

    switch ( this->mem_ )
      {
      case DDR:
	err = MPI_Get ( srcBuffer, srcSize, MPI_BYTE, 0, offset, srcSize, MPI_BYTE, this->RMAWin_ );
	break;
      default:
	printMsg ( ERROR, "Error while reading data (mem = %s)\n", this->memName () );
	MPI_Abort ( MPI_COMM_WORLD, -1 );
      }

    return err;
  }
 
  int memFlush ( ) {
    int err;
    MPI_Status status;

    switch ( this->mem_ )
      {
      case DDR:
	MPI_Win_fence ( 0, this->RMAWin_ );
        break;
      default:
        printMsg ( ERROR, "Error while flushing data (mem = %s)\n", this->memName () );
        MPI_Abort ( MPI_COMM_WORLD, -1 );
      }

    return err;

  }

  /**********************/
  /*  |-- Utils         */
  /**********************/
  char* memName  ( ) {
    switch ( this->mem_ )
      {
      case DDR:
	return "DDR";
	break;
      case HBM:
	return "HBM";
	break;
      case SSD:
	return "SSD";
	break;
      case HDD:
	return "HDD";
	break;
      default:
	printMsg ( ERROR, "Wrong memory type!\n" );
	MPI_Abort ( MPI_COMM_WORLD, -1 );
      }
  }

  mem_t memType  ( char* name ) {
    if ( ! strcmp ( "DDR", name ) ) return DDR;
    if ( ! strcmp ( "HBM", name ) ) return HBM;
    if ( ! strcmp ( "SSD", name ) ) return SSD;
    if ( ! strcmp ( "HDD", name ) ) return HDD;

    printMsg ( ERROR, "Wrong memory name!\n" );
    MPI_Abort ( MPI_COMM_WORLD, -1 );
  }

  /************************/
  /*  |-- Characteristics */
  /************************/
  int64_t memBandwidth ( ) {
    return 0;
  }

  int64_t memLatency   ( ) {
    return 0;
  }

  int64_t memCapacity  ( ) {
    return 0;
  }
};

#endif // TP_MEMORY_H