#include "tp_memory.hpp"

Memory::Memory ()
{
  this->request_ = NULL;
  this->mmapAllocatorRank_ = 0;
}


Memory::~Memory ()
{
}


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

  switch ( this->mem_ )
    {
    case DDR:
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Allocate memory on DDR (%s:%d)\n", __FILE__, __LINE__ );
	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;
    case HBM:
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Allocate memory on HBM (%s:%d)\n", __FILE__, __LINE__ );
	this->buffer_ = hbw_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;
    case NLS:
    case PFS:
      char *name;
      int resultlen;
      name = (char *) malloc ( MPI_MAX_PROCESSOR_NAME * sizeof ( char ) );
      MPI_Get_processor_name( name, &resultlen );
      
      strcpy ( this->fileName_, fileName );
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Open file %s on %s from node %s (%s:%d)\n", this->fileName_, this->memName (), name, __FILE__, __LINE__ );
      }
      err = MPI_File_open( this->comm_, this->fileName_, MPI_MODE_RDWR | MPI_MODE_CREATE,
			   MPI_INFO_NULL, &this->fileHandle_ );
      
      if ( err != MPI_SUCCESS )
	printMsg ( ERROR, "Error while opening the file %s on %s from node %s (%s:%d)\n", this->fileName_, this->memName (), name, __FILE__, __LINE__ );
      break;
    case NVR:
      sprintf ( this->fileName_, "%s/%s", this->memPath (), fileName );
      
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Map file %s in DRAM (%s:%d)\n", this->fileName_, __FILE__, __LINE__ );
	// Add test for memory capacity
	/*
	 *  O_RDWR | O_CREAT : read, write, create
	 *  S_IRWXU : read, write, execute/search by owner 
	 */
	this->fd_ = open(this->fileName_, O_RDWR | O_CREAT, S_IRWXU);
	
	if ( this->fd_ == -1 ) {
	  printMsg ( ERROR, "Error while opening the file %s (%s:%d)\n", this->fileName_, __FILE__, __LINE__ );
	  MPI_Abort ( MPI_COMM_WORLD, -1 );	  
	}

	if ( ftruncate ( this->fd_, this->buffSize_ ) == -1 ) {
	  printMsg ( ERROR, "Error while truncating the file %s (%s:%d)\n", this->fileName_, __FILE__, __LINE__ );
	  MPI_Abort ( MPI_COMM_WORLD, -1 );	  
	}
	
	/*
	 *  PROT_READ | PROT_WRITE : read, write
	 *  MAP_SHARED : Updates to the mapping are visible to other
	 *               processes mapping the same region, and (in
	 *               the case of file-backed mappings) are carried
	 *               through to the underlying file. 
	 */
	this->buffer_ = mmap(0, this->buffSize_, PROT_READ | PROT_WRITE, MAP_SHARED, this->fd_, 0);
	if (this->buffer_ == MAP_FAILED) {
	  printMsg ( ERROR, "Mmap of file %s has failed (%s:%d)\n", this->fileName_, __FILE__, __LINE__ );
          MPI_Abort ( MPI_COMM_WORLD, -1 );
	}
	
	err = MPI_Win_create ( this->buffer_, this->buffSize_, 1, MPI_INFO_NULL, this->comm_, &this->RMAWin_ );
      }
      else
	err = MPI_Win_create ( NULL, 0, 1, MPI_INFO_NULL, this->comm_, &this->RMAWin_ );

      if ( err != MPI_SUCCESS)
	printMsg ( ERROR, "Unable to create RMA window on mmapped file %s (%s:%d)\n", this->fileName_, __FILE__, __LINE__);
      
      this->memFlush ();
      break;
    default:
      printMsg ( ERROR, "Unable to allocate memory (mem = %s)\n", this->memName () );
      MPI_Abort ( MPI_COMM_WORLD, -1 );
    }
}


void Memory::memFree  ( ) {
  int rank, err;
  MPI_Comm_rank ( MPI_COMM_WORLD, &rank );
    
  switch ( this->mem_ )
    {
    case DDR:
      MPI_Win_free ( &this->RMAWin_ );
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Free memory on DDR (%s:%d)\n", __FILE__, __LINE__ );
	free ( this->buffer_ );
      }
      break;
    case HBM:
      MPI_Win_free ( &this->RMAWin_ );
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Free memory on HBM (%s:%d)\n", __FILE__, __LINE__ );
	hbw_free ( this->buffer_ );
      }
      break;
    case NLS:
    case PFS:
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Close file %s on %s (%s:%d)\n", this->fileName_, this->memName (), __FILE__, __LINE__ );
      }
      MPI_File_close ( &this->fileHandle_ );
      break;
    case NVR:
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Free memory on NVRAM and unmap file %s (%s:%d)\n", this->fileName_, __FILE__, __LINE__ );
	close ( this->fd_ );
	err = munmap ( this->buffer_, this->buffSize_ );
	if ( err == -1 ) {
	  printMsg ( ERROR, "Error while unmaping the file %s (%s:%d)\n", this->fileName_, __FILE__, __LINE__ );
          MPI_Abort ( MPI_COMM_WORLD, -1 );
	}
      }
      
      err = MPI_Win_free ( &this->RMAWin_ );
      if ( err != MPI_SUCCESS)
        printMsg ( ERROR, "Unable to free RMA window on mmapped file %s (%s:%d)\n", this->fileName_, __FILE__, __LINE__);
      
      break;
    default:
      printMsg ( ERROR, "Unable to free memory (mem = %s)\n", this->memName () );
      MPI_Abort ( MPI_COMM_WORLD, -1 );
    }
}


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

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

  return err;
}


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

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

  return err;
}

 
int Memory::memFlush ( ) {
  int err, rank;
  MPI_Status status;

  MPI_Comm_rank ( MPI_COMM_WORLD, &rank );

  switch ( this->mem_ )
    {
    case DDR:
    case HBM:
      MPI_Win_fence ( 0, this->RMAWin_ );
      break;
    case NLS:
    case PFS:
      if ( this->request_ != NULL )
	MPI_Wait ( &this->request_, &status );
      break;
    case NVR:
      err = MPI_Win_fence ( 0, this->RMAWin_ );
      if ( err != MPI_SUCCESS)
       	printMsg ( ERROR, "Unable to perform a synchronization on the RMA window (mem = %s)\n", this->memName () );
      
      if ( this->masterRank_ ) {
	printMsg ( DEBUG, "Sync memory and file %s (%s:%d)\n", this->fileName_, __FILE__, __LINE__ );
	err = msync( this->buffer_, this->buffSize_, MS_ASYNC );
	if ( err == -1 ) {
	  printMsg ( ERROR, "Error while syncing memory and file %s (%s:%d)\n", this->fileName_, __FILE__, __LINE__ );
          MPI_Abort ( MPI_COMM_WORLD, -1 );
	}
      }
      
      err = 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;
}


int Memory::memUncache  ( ) {
  return 0;
}


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


char* Memory::memName  ( mem_t mem ) {
  switch ( mem )
    {
    case DDR:
      return "DDR";
      break;
    case HBM:
      return "HBM";
      break;
    case PFS:
      return "PFS";
      break;
    case NLS:
      return "NLS";
      break;
    case NVR:
      return "NVR";
      break;
    default:
      printMsg ( ERROR, "Wrong memory type!\n" );
      MPI_Abort ( MPI_COMM_WORLD, -1 );
    }
}


mem_t Memory::memTypeByName  ( char* name ) {
  if ( ! strcmp ( "DDR", name ) ) return DDR;
  if ( ! strcmp ( "HBM", name ) ) return HBM;
  if ( ! strcmp ( "PFS", name ) ) return PFS;
  if ( ! strcmp ( "NLS", name ) ) return NLS;
  if ( ! strcmp ( "NVR", name ) ) return NVR;

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


mem_t Memory::memTypeByPath  ( char* path ) {
  if ( ! strncmp ( path, "/lus/theta-fs0", 14 ) ) return PFS;
  if ( ! strncmp ( path, "/local/scratch", 14 ) ) return NLS;
  
  printMsg ( ERROR, "No memory tier corresponding to %s!\n", path );
  MPI_Abort ( MPI_COMM_WORLD, -1 );
}


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


int64_t Memory::memLatency   ( ) {
  return 0;
}


int64_t Memory::memCapacity  ( ) {
  return 0;
}


bool Memory::memPersistency  ( ) {
  return false;
}


char* Memory::memPath ( ) {
  if ( this->mem_ == NVR )
    return "/local/scratch";
  else
    return "";      
}


/*
 *  kBps
 */
int64_t Memory::memBandwidth ( mem_t mem ) {
    switch ( mem )
    {
    case DDR:
      return 90000000;
      break;
    case HBM:
      return 350000000;
      break;
    case PFS:
      return 1800000;
      break;
    case NLS:
      return 1800000;
      break;
    case NVR:
      return 400000;
      break;
    default:
      printMsg ( ERROR, "Wrong memory type!\n" );
      MPI_Abort ( MPI_COMM_WORLD, -1 );
    }
}

/*
 *  ms
 */
int64_t Memory::memLatency  ( mem_t mem ) {
    switch ( mem )
    {
    case DDR:
      return 2;
      break;
    case HBM:
      return 1;
      break;
    case PFS:
      return 30;
      break;
    case NLS:
      return 5;
      break;
    case NVR:
      return 5;
      break;
    default:
      printMsg ( ERROR, "Wrong memory type!\n" );
      MPI_Abort ( MPI_COMM_WORLD, -1 );
    }
}


int64_t Memory::memCapacity  ( mem_t mem ) {
  return 0;
}


bool Memory::memPersistency  ( mem_t mem ) {
  return false;
}