/****************************************************************************
**  SIONLIB     http://www.fz-juelich.de/jsc/sionlib                       **
*****************************************************************************
**  Copyright (c) 2008-2018                                                **
**  Forschungszentrum Juelich, Juelich Supercomputing Centre               **
**                                                                         **
**  See the file COPYRIGHT in the package base directory for details       **
****************************************************************************/

#define _XOPEN_SOURCE 700

#include <inttypes.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "sion_const.h"
#include "sion_debug.h"
#include "sion_error_handler.h"
#include "sion_filedesc.h"
#include "sion_generic.h"
#include "sion_omp_cb_gen.h"

#ifdef SION_OMP

static void *_lsion_omp_global_pointer;

int _lsion_omp_size_of_dtype(int dtype);

static void *_sion_omp_share_ptr(void *ptr);

int _lsion_register_callbacks_omp(void)
{
  int aid = 0;
  aid = lsion_generic_create_api("SIONlib_OMP_API");

  lsion_generic_register_create_local_commgroup_cb(aid, &_lsion_omp_create_lcg_cb);
  lsion_generic_register_free_local_commgroup_cb(aid, &_lsion_omp_free_lcg_cb);

  lsion_generic_register_barrier_cb(aid, &_lsion_omp_barrier_cb);
  lsion_generic_register_bcastr_cb(aid, &_lsion_omp_bcastr_cb);
  lsion_generic_register_gatherr_cb(aid, &_lsion_omp_gatherr_cb);
  lsion_generic_register_scatterr_cb(aid, &_lsion_omp_scatterr_cb);
  lsion_generic_register_gathervr_cb(aid, &_lsion_omp_gathervr_cb);
  lsion_generic_register_scattervr_cb(aid, &_lsion_omp_scattervr_cb);
  lsion_generic_register_gather_and_execute_cb(aid, &_lsion_omp_gather_process_cb);
  lsion_generic_register_execute_and_scatter_cb(aid, &_lsion_omp_process_scatter_cb);
  lsion_generic_register_get_capability_cb(aid, &_lsion_omp_get_capability_cb);

  return aid;
}

int _lsion_omp_create_lcg_cb(
  void **local_commgroup, void *global_commgroup, int grank, int gsize, int lrank, int lsize, int filenumber, int numfiles)
{
  int rc = SION_STD_SUCCESS;

  DPRINTFP((256, "_lsion_omp_create_lcg_cb", grank,
    " DUMMY for split global comm: grank=%d gsize=%d filenumber=%d, numfiles=%d, lrank=%d lsize=%d \n", grank, gsize, filenumber,
    numfiles, lrank, lsize));
  *local_commgroup = global_commgroup;

  return rc;
}

int _lsion_omp_free_lcg_cb(void *local_commgroup)
{
  int rc = SION_STD_SUCCESS;
  DPRINTFP((256, "_omp_free_lcg_cb", _SION_DEFAULT_RANK, " DUMMY for free local comm\n"));
  return rc;
}

int _lsion_omp_barrier_cb(void *commdata)
{
  int rc = SION_STD_SUCCESS;
  {
#pragma omp barrier
  }
  return rc;
}

int _lsion_omp_bcastr_cb(void *data, void *commdata, int dtype, int nelem, int root)
{
  int rc = SION_STD_SUCCESS;
  _omp_api_commdata *sapi = (_omp_api_commdata *)commdata;

  if (sapi->thread_num == root) {
    _lsion_omp_global_pointer = data;
  }

  /* threads sync global pointer */

  {
#pragma omp barrier
  }

  if ((sapi->thread_num != root) && (_lsion_omp_global_pointer != NULL)) {
    memcpy(data, _lsion_omp_global_pointer, nelem * _lsion_omp_size_of_dtype(dtype));
  }
  {
#pragma omp barrier
  }

  return rc;
}

/* indata: data that comes into fuinction, outdata: data that will be created by function */
int _lsion_omp_gatherr_cb(void *indata, void *outdata, void *commdata, int dtype, int nelem, int root)
{
  int rc = SION_STD_SUCCESS;
  _omp_api_commdata *sapi = (_omp_api_commdata *)commdata;

  /* set global pointer to outdata */
  if (sapi->thread_num == root) {
    _lsion_omp_global_pointer = outdata;
  }

  /* synchronize */
  {
#pragma omp barrier
  }

  /* copy indata -> outdata */
  memcpy((char *)_lsion_omp_global_pointer + sapi->thread_num * nelem * _lsion_omp_size_of_dtype(dtype), indata,
    nelem * _lsion_omp_size_of_dtype(dtype));

  /* synchronize again */
  {
#pragma omp barrier
  }
  return rc;
}

int _lsion_omp_scatterr_cb(void *indata, void *outdata, void *commdata, int dtype, int nelem, int root)
{
  int rc = SION_STD_SUCCESS;
  _omp_api_commdata *sapi = (_omp_api_commdata *)commdata;

  /* set global pointer to indata */
  if (sapi->thread_num == root) {
    _lsion_omp_global_pointer = indata;
  }

  /* synchronize */
  {
#pragma omp barrier
  }

  /* copy outdata <- indata */
  memcpy(outdata, (char *)_lsion_omp_global_pointer + sapi->thread_num * nelem * _lsion_omp_size_of_dtype(dtype),
    nelem * _lsion_omp_size_of_dtype(dtype));

  /* synchronize again */
  {
#pragma omp barrier
  }

  return rc;
}

int _lsion_omp_gathervr_cb(void *indata, void *outdata, void *commdata, int dtype, int *counts, int nelem, int root)
{
  int rc = SION_STD_SUCCESS;
  int *displs = NULL;
  int t, offset;
  _omp_api_commdata *sapi = (_omp_api_commdata *)commdata;

  /* compute offsets sizes for each thread and store these values in a temporary vector */
  if (sapi->thread_num == root) {
    displs = malloc(sapi->num_threads * sizeof(int));
    if (displs == NULL) {
      fprintf(stderr, "__sion_omp_gathervr_cb: cannot allocate temporary memory of size %zu (displs), aborting ...\n",
        (size_t)sapi->num_threads * sizeof(int));
      return -1;
    }
    offset = 0;
    for (t = 0; t < sapi->num_threads; t++) {
      displs[t] = offset;
      offset += counts[t];
      DPRINTFP(
        (256, "_lsion_omp_gathervr_cb", sapi->thread_num, " after compute %2d -> dpls=%2d count=%d\n", t, displs[t], counts[t]));
    }

    /* set global pointer to temporary vector */
    _lsion_omp_global_pointer = displs;
  }

  /* synchronize */
  {
#pragma omp barrier
  }

  /* get offset */
  offset = ((int *)_lsion_omp_global_pointer)[sapi->thread_num];

  /* synchronize again */
  {
#pragma omp barrier
  }

  /* free temporary vector and set global pointer to outdata */
  if (sapi->thread_num == root) {
    free(displs);
  }

  /* set global pointer to temporary vector */
  _lsion_omp_global_pointer = outdata;

  /* synchronize again */
  {
#pragma omp barrier
  }

  /* copy indata -> outdata */
  memcpy(
    (char *)_lsion_omp_global_pointer + offset * _lsion_omp_size_of_dtype(dtype), indata, nelem * _lsion_omp_size_of_dtype(dtype));

  /* synchronize again */
  {
#pragma omp barrier
  }
  return rc;
}

int _lsion_omp_scattervr_cb(void *indata, void *outdata, void *commdata, int dtype, int *counts, int nelem, int root)
{
  int rc = SION_STD_SUCCESS;
  int *displs = NULL;
  int t, offset;
  _omp_api_commdata *sapi = (_omp_api_commdata *)commdata;

  /* compute offsets for each thread and store these values in a temporary vector */
  if (sapi->thread_num == root) {
    displs = malloc(sapi->num_threads * sizeof(int));
    if (displs == NULL) {
      fprintf(stderr, "__sion_omp_gathervr_cb: cannot allocate temporary memory of size %zu (displs), aborting ...\n",
        (size_t)sapi->num_threads * sizeof(int));
      return -1;
    }
    offset = 0;
    for (t = 0; t < sapi->num_threads; t++) {
      displs[t] = offset;
      offset += counts[t];
      DPRINTFP(
        (256, "_lsion_omp_gathervr_cb", sapi->thread_num, " after compute %2d -> dpls=%2d count=%d\n", t, displs[t], counts[t]));
    }

    /* set global pointer to temporary vector */
    _lsion_omp_global_pointer = displs;
  }

  /* synchronize */
  {
#pragma omp barrier
  }

  /* get offset */
  offset = ((int *)_lsion_omp_global_pointer)[sapi->thread_num];

  /* synchronize again */
  {
#pragma omp barrier
  }

  /* free temporary vector and set global pointer to outdata */
  if (sapi->thread_num == root) {
    free(displs);
  }

  /* set global pointer to temporary vector */
  _lsion_omp_global_pointer = indata;

  /* synchronize again */
  {
#pragma omp barrier
  }

  /* copy indata -> outdata */
  memcpy(
    outdata, (char *)_lsion_omp_global_pointer + offset * _lsion_omp_size_of_dtype(dtype), nelem * _lsion_omp_size_of_dtype(dtype));

  /* synchronize again */
  {
#pragma omp barrier
  }
  return rc;
}

int64_t _lsion_omp_gather_process_cb(const void *data, int64_t fsblksize, void *commdata, int collector, int range_start,
  int range_end, int sid, int64_t process_cb(const void *, int64_t *, int), void *iterator,
  bool iterator_next_cb(void *, int64_t *))
{
  _omp_api_commdata *sapi = commdata;

  const void **datas = NULL;
  if (sapi->thread_num == collector) {
    datas = malloc(sapi->num_threads * sizeof(void *));
  }

  datas = _sion_omp_share_ptr(datas);

  if (datas == NULL) {
    return _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_RETURN,
      "_lsion_omp_gather_process_cb: cannot allocate temporary memory of size %zu (datas), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(void *));
  }

  datas[sapi->thread_num] = data;

  void **iterators = NULL;
  if (sapi->thread_num == collector) {
    iterators = malloc(sapi->num_threads * sizeof(void *));
  }

  iterators = _sion_omp_share_ptr(iterators);

  if (iterators == NULL) {
    free(datas);
    return _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_RETURN,
      "_lsion_omp_gather_process_cb: cannot allocate temporary memory of size %zu (iterators), aborting ...\n",
      (size_t)sapi->thread_num * sizeof(void *));
  }

  iterators[sapi->thread_num] = iterator;

  int64_t *bytes_writtens = NULL;
  if (sapi->thread_num == collector) {
    bytes_writtens = malloc(sapi->num_threads * sizeof(int64_t));
  }

  bytes_writtens = _sion_omp_share_ptr(bytes_writtens);

  if (bytes_writtens == NULL) {
    free(datas);
    free(iterators);
    return _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_RETURN,
      "_lsion_omp_gather_process_cb: cannot allocate temporary memory of size %zu (bytes_writtens), aborting...\n",
      (size_t)sapi->thread_num * sizeof(int64_t));
  }

// make sure all data buffer pointers and iterators have been gathered
#pragma omp barrier

  /* scan all other tasks */
  if (sapi->thread_num == collector) {
    for (int t = range_start; t <= range_end; ++t) {
      const char *p = datas[t];
      int64_t spec[] = {-1, -1};
      int64_t total_bytes_written = 0;

      while (iterator_next_cb(iterators[t], spec)) {
        DPRINTFP((256, "_lsion_omp_gather_process_cb", -1, "on master t=%d spec=%" PRId64 ",%" PRId64 "\n", t, spec[0], spec[1]));

        int64_t bytes_written = process_cb(p, spec, sid);

        p += bytes_written;
        total_bytes_written += bytes_written;
        if (bytes_written != spec[1]) {
          break;
        }
      }

      bytes_writtens[t] = total_bytes_written;
    }
  }

// make sure collector thread has updated write counts
#pragma omp barrier

  int64_t rc = (sapi->thread_num == collector) ? 0 : bytes_writtens[sapi->thread_num];

// make sure everyone has read their write count
#pragma omp barrier

  if (sapi->thread_num == collector) {
    free(datas);
    free(iterators);
    free(bytes_writtens);
  }

  return rc;
}

int64_t _lsion_omp_process_scatter_cb(void *data, int64_t fsblksize, void *commdata, int collector, int range_start, int range_end,
  int sid, int64_t process_cb(void *, int64_t *, int), void *iterator, bool iterator_next_cb(void *, int64_t *))
{
  _omp_api_commdata *sapi = commdata;

  void **datas = NULL;
  if (sapi->thread_num == collector) {
    datas = malloc(sapi->num_threads * sizeof(void *));
  }

  datas = _sion_omp_share_ptr(datas);

  if (datas == NULL) {
    return _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_RETURN,
      "_lsion_omp_process_scatter_cb: cannot allocate temporary memory of size %zu (datas), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(void *));
  }

  datas[sapi->thread_num] = data;

  void **iterators = NULL;
  if (sapi->thread_num == collector) {
    iterators = malloc(sapi->num_threads * sizeof(void *));
  }

  iterators = _sion_omp_share_ptr(iterators);

  if (iterators == NULL) {
    return _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_RETURN,
      "_lsion_omp_process_scatter_cb: cannot allocate temporary memory of size %zu (iterators), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(void *));
  }

  iterators[sapi->thread_num] = iterator;

  int64_t *bytes_reads = NULL;
  if (sapi->thread_num == collector) {
    bytes_reads = malloc(sapi->num_threads * sizeof(int64_t));
  }

  bytes_reads = _sion_omp_share_ptr(bytes_reads);

  if (bytes_reads == NULL) {
    return _lsion_errorprint(SION_SIZE_NOT_VALID, _SION_ERROR_RETURN,
      "_lsion_omp_process_scatter_cb: cannot allocate temporary memory of size %zu (bytes_reads), aborting ...\n",
      (size_t)sapi->num_threads * sizeof(int64_t));
  }

// make sure all data buffer pointers and iterators have been gathered
#pragma omp barrier

  /* scan all other tasks */
  if (sapi->thread_num == collector) {
    for (int t = range_start; t <= range_end; ++t) {
      char *p = datas[t];
      int64_t spec[] = {-1, -1};
      int64_t total_bytes_read = 0;

      while (iterator_next_cb(iterators[t], spec)) {
        DPRINTFP((256, "_lsion_omp_process_scatter_cb", -1, "on master t=%d spec=%" PRId64 ",%" PRId64 "\n", t, spec[0], spec[1]));

        int64_t bytes_read = process_cb(p, spec, sid);

        p += bytes_read;
        total_bytes_read += bytes_read;
        if (bytes_read != spec[1]) {
          break;
        }
      }

      bytes_reads[t] = total_bytes_read;
    }
  }

// make sure collector thread has updated read counts
#pragma omp barrier

  int64_t rc = (sapi->thread_num == collector) ? 0 : bytes_reads[sapi->thread_num];

// make sure everyone has read their read count
#pragma omp barrier

  if (sapi->thread_num == collector) {
    free(datas);
    free(iterators);
    free(bytes_reads);
  }

  return rc;
}

int _lsion_omp_get_capability_cb(void *commdata)
{
  int rc = SION_CAPABILITY_NONE;
  _omp_api_commdata *sapi = (_omp_api_commdata *)commdata;

  if (sapi->thread_num == 0) {
    rc = SION_CAPABILITY_FULL;
    DPRINTFP((256, "_lsion_omp_get_capability_cb", sapi->thread_num, "FULL capability\n"));
  } else {
    rc = SION_CAPABILITY_ONLY_SENDER;
    DPRINTFP((256, "_lsion_omp_get_capability_cb", sapi->thread_num, "ONLY SENDER capability\n"));
  }
  return rc;
}

static void *_sion_omp_share_ptr(void *ptr)
{
// make sure everyone is done with the global pointer
#pragma omp barrier

#pragma omp master
  _lsion_omp_global_pointer = ptr;

// make sure the master thread has updated the global pointer
#pragma omp barrier

  void *ret = _lsion_omp_global_pointer;

// make sure everyone has read the global pointer
#pragma omp barrier
  return ret;
}

int _lsion_omp_size_of_dtype(int dtype)
{
  switch (dtype) {
  case _SION_INT32:
    return sizeof(int32_t);
    break;
  case _SION_INT64:
    return sizeof(int64_t);
    break;
  case _SION_CHAR:
    return 1;
    break;
  default:
    return sizeof(int64_t);
  }
}

/* end of ifdef OPENMP */
#endif
