/*
  This file is part of the SC Library.
  The SC Library provides support for parallel scientific applications.

  Copyright (C) 2010 The University of Texas System
  Additional copyright (C) 2011 individual authors

  The SC Library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.

  The SC Library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public
  License along with the SC Library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
  02110-1301, USA.
*/

#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif
#include <sc_containers.h>
#include <sc_sort.h>

typedef struct sc_psort_peer
{
  int                 received;
  int                 sent;
  int                 prank;
  size_t              length;
  char               *buffer;
  char               *my_start;
}
sc_psort_peer_t;

typedef struct sc_psort
{
  sc_MPI_Comm         mpicomm;
  int                 num_procs, rank;
  size_t              size;
  size_t              my_lo, my_hi, my_count;
  size_t             *gmemb;
  char               *my_base;
  int                 (*compar) (const void *, const void *);
}
sc_psort_t;

#ifndef SC_HAVE_QSORT_R

/* qsort is not reentrant, so we do the inverse static */
static int          (*sc_compare) (const void *, const void *);
static int
sc_icompare (const void *v1, const void *v2)
{
  return sc_compare (v2, v1);
}

#else

#ifndef SC_HAVE_BSD_QSORT_R

/* qsort_r comparator functions conforming to GNU standard */

static int
sc_compare_r (const void *v1, const void *v2, void *arg)
{
  sc_psort_t         *pst = (sc_psort_t *) arg;
  return pst->compar (v1, v2);
}

static int
sc_icompare_r (const void *v1, const void *v2, void *arg)
{
  sc_psort_t         *pst = (sc_psort_t *) arg;
  return pst->compar (v2, v1);
}

#else /* SC_HAVE_BSD_QSORT_R */

/* qsort_r comparator functions conforming to BSD standard */

static int
sc_compare_r (void *arg, const void *v1, const void *v2)
{
  sc_psort_t         *pst = (sc_psort_t *) arg;
  return pst->compar (v1, v2);
}

static int
sc_icompare_r (void *arg, const void *v1, const void *v2)
{
  sc_psort_t         *pst = (sc_psort_t *) arg;
  return pst->compar (v2, v1);
}

#endif /* SC_HAVE_BSD_QSORT_R */
#endif /* SC_HAVE_QSORT_R */

static              size_t
sc_bsearch_cumulative (const size_t * cumulative, size_t nmemb,
                       size_t pos, size_t guess)
{
  size_t              proc_low, proc_high;

  proc_low = 0;
  proc_high = nmemb - 1;

  for (;;) {
    SC_ASSERT (proc_low <= proc_high);
    SC_ASSERT (proc_low < nmemb && proc_high < nmemb);
    SC_ASSERT (proc_low <= guess && guess <= proc_high);

    /* check if pos is on a lower owner than guess */
    if (pos < cumulative[guess]) {
      proc_high = guess - 1;
      guess = (proc_low + proc_high + 1) / 2;
      continue;
    }

    /* check if pos is on a higher owner than guess */
    if (cumulative[guess + 1] <= pos) {
      proc_low = guess + 1;
      guess = (proc_low + proc_high) / 2;
      continue;
    }

    /* otherwise guess is the correct owner */
    break;
  }

  /* make sure we found a valid owner with nonzero count */
  SC_ASSERT (guess < nmemb);
  SC_ASSERT (cumulative[guess] <= pos && pos < cumulative[guess + 1]);
  return guess;
}

static void
sc_merge_bitonic (sc_psort_t * pst, size_t lo, size_t hi, int dir)
{
  const size_t        n = hi - lo;

  if (n > 1 && pst->my_hi > lo && pst->my_lo < hi) {
    const int           rank = pst->rank;
    const size_t        size = pst->size;
    size_t              k, n2;
    size_t              lo_end, hi_beg;
    size_t              lo_length, hi_length;
    size_t              offset, max_length;
    int                 lo_owner, hi_owner;
    int                 num_peers, remaining, remaining2, outcount, outcount2;
    int                 mpiret;
    sc_array_t          a, *pa = &a;
    sc_array_t          r, *pr = &r;
    sc_array_t          s, *ps = &s;
    int                *wait_indices, *wait_indices2;
    sc_MPI_Status      *recv_statuses, *recv_statuses2;
    sc_psort_peer_t    *peer;

    for (k = 1; k < n;) {
      k = k << 1;
    }
    n2 = k >> 1;
    SC_ASSERT (n2 >= n / 2 && n2 < n);

    lo_end = lo + n - n2;
    hi_beg = lo + n2;
    SC_ASSERT (lo_end <= hi_beg && lo_end - lo == hi - hi_beg);

    sc_array_init (pa, sizeof (sc_psort_peer_t));
    sc_array_init (pr, sizeof (sc_MPI_Request));
    sc_array_init (ps, sizeof (sc_MPI_Request));

    /* loop 1: initiate communication */
    lo_owner = hi_owner = rank;
    offset = max_length = 0;
    for (offset = 0; offset < lo_end - lo; offset += max_length) {
      lo_owner =
        (int) sc_bsearch_cumulative (pst->gmemb, (size_t) pst->num_procs,
                                     lo + offset, (size_t) lo_owner);
      lo_length = pst->gmemb[lo_owner + 1] - (lo + offset);
      hi_owner =
        (int) sc_bsearch_cumulative (pst->gmemb, (size_t) pst->num_procs,
                                     hi_beg + offset, (size_t) hi_owner);
      hi_length = pst->gmemb[hi_owner + 1] - (hi_beg + offset);
      max_length = lo_end - (lo + offset);
      max_length = SC_MIN (max_length, SC_MIN (lo_length, hi_length));
      SC_ASSERT (max_length > 0);

      if (lo_owner == rank && hi_owner != rank) {
        char               *lo_data;
        sc_MPI_Request     *rreq, *sreq;
        const int           bytes = (int) (max_length * size);

        /* receive high part, send low part */
        peer = (sc_psort_peer_t *) sc_array_push (pa);
        rreq = (sc_MPI_Request *) sc_array_push (pr);
        sreq = (sc_MPI_Request *) sc_array_push (ps);
        lo_data = pst->my_base + (lo + offset - pst->my_lo) * size;

        peer->received = 0;
        peer->sent = 0;
        peer->prank = hi_owner;
        peer->length = max_length;
        peer->buffer = SC_ALLOC (char, bytes);
        peer->my_start = lo_data;
        mpiret = sc_MPI_Irecv (peer->buffer, bytes, sc_MPI_BYTE,
                               peer->prank, SC_TAG_PSORT_HI, pst->mpicomm,
                               rreq);
        SC_CHECK_MPI (mpiret);

        SC_ASSERT (lo_data >= pst->my_base);
        SC_ASSERT (lo_data + bytes <= pst->my_base + pst->my_count * size);
        mpiret = sc_MPI_Isend (lo_data, bytes, sc_MPI_BYTE,
                               peer->prank, SC_TAG_PSORT_LO, pst->mpicomm,
                               sreq);
        SC_CHECK_MPI (mpiret);
      }
      else if (lo_owner != rank && hi_owner == rank) {
        char               *hi_data;
        sc_MPI_Request     *rreq, *sreq;
        const int           bytes = (int) (max_length * size);

        /* receive low part, send high part */
        peer = (sc_psort_peer_t *) sc_array_push (pa);
        rreq = (sc_MPI_Request *) sc_array_push (pr);
        sreq = (sc_MPI_Request *) sc_array_push (ps);
        hi_data = pst->my_base + (hi_beg + offset - pst->my_lo) * size;

        peer->received = 0;
        peer->sent = 0;
        peer->prank = lo_owner;
        peer->length = max_length;
        peer->buffer = SC_ALLOC (char, bytes);
        peer->my_start = hi_data;

        mpiret = sc_MPI_Irecv (peer->buffer, bytes, sc_MPI_BYTE,
                               peer->prank, SC_TAG_PSORT_LO, pst->mpicomm,
                               rreq);
        SC_CHECK_MPI (mpiret);

        SC_ASSERT (hi_data >= pst->my_base);
        SC_ASSERT (hi_data + bytes <= pst->my_base + pst->my_count * size);
        mpiret = sc_MPI_Isend (hi_data, bytes, sc_MPI_BYTE,
                               peer->prank, SC_TAG_PSORT_HI, pst->mpicomm,
                               sreq);
        SC_CHECK_MPI (mpiret);
      }
    }

    /* loop 2: local computation */
    lo_owner = hi_owner = rank;
    offset = max_length = 0;
    for (offset = 0; offset < lo_end - lo; offset += max_length) {
      lo_owner =
        (int) sc_bsearch_cumulative (pst->gmemb, (size_t) pst->num_procs,
                                     lo + offset, (size_t) lo_owner);
      lo_length = pst->gmemb[lo_owner + 1] - (lo + offset);
      hi_owner =
        (int) sc_bsearch_cumulative (pst->gmemb, (size_t) pst->num_procs,
                                     hi_beg + offset, (size_t) hi_owner);
      hi_length = pst->gmemb[hi_owner + 1] - (hi_beg + offset);
      max_length = lo_end - (lo + offset);
      max_length = SC_MIN (max_length, SC_MIN (lo_length, hi_length));
      SC_ASSERT (max_length > 0);

      if (lo_owner == rank && hi_owner == rank) {
        size_t              zz;
        char               *lo_data, *hi_data;
        char               *temp = SC_ALLOC (char, size);

        /* local comparisons only */
        lo_data = pst->my_base + (lo + offset - pst->my_lo) * size;
        hi_data = pst->my_base + (hi_beg + offset - pst->my_lo) * size;
        for (zz = 0; zz < max_length; ++zz) {
          if (dir == (pst->compar (lo_data, hi_data) > 0)) {
            memcpy (temp, lo_data, size);
            memcpy (lo_data, hi_data, size);
            memcpy (hi_data, temp, size);
          }
          lo_data += size;
          hi_data += size;
        }
        SC_FREE (temp);
      }
    }

    /* loop 3: receive and compute with received data */
    outcount = 0;
    outcount2 = 0;
    num_peers = (int) pa->elem_count;
    wait_indices = SC_ALLOC (int, num_peers);
    wait_indices2 = SC_ALLOC (int, num_peers);
    recv_statuses = SC_ALLOC (sc_MPI_Status, num_peers);
    recv_statuses2 = SC_ALLOC (sc_MPI_Status, num_peers);
    for (remaining = num_peers, remaining2 = num_peers;
         remaining > 0 || remaining2 > 0;
         remaining -= outcount, remaining2 -= outcount2) {
      int                 i;
      outcount = 0;
      outcount2 = 0;
      if (remaining > 0) {
        mpiret = sc_MPI_Waitsome (num_peers, (sc_MPI_Request *) pr->array,
                                  &outcount, wait_indices, recv_statuses);
        SC_CHECK_MPI (mpiret);
        SC_ASSERT (outcount != sc_MPI_UNDEFINED);
        SC_ASSERT (outcount > 0);
        for (i = 0; i < outcount; ++i) {
          size_t              zz;
          char               *lo_data, *hi_data;
#ifdef SC_ENABLE_DEBUG
          sc_MPI_Status      *jstatus;

          jstatus = &recv_statuses[i];
#endif

          /* retrieve peer information */
          peer = (sc_psort_peer_t *) sc_array_index_int (pa, wait_indices[i]);
          if (peer->sent) {
            /* only overwrite data,
             * if made sure it has already been sent to the other process */
            SC_ASSERT (!peer->received);
            SC_ASSERT (peer->prank != rank);
            SC_ASSERT (peer->prank == jstatus->MPI_SOURCE);

            /* comparisons with remote peer */
            if (rank < peer->prank) {
              lo_data = peer->my_start;
              hi_data = peer->buffer;
              for (zz = 0; zz < peer->length; ++zz) {
                if (dir == (pst->compar (lo_data, hi_data) > 0)) {
                  memcpy (lo_data, hi_data, size);
                }
                lo_data += size;
                hi_data += size;
              }
            }
            else {
              lo_data = peer->buffer;
              hi_data = peer->my_start;
              for (zz = 0; zz < peer->length; ++zz) {
                if (dir == (pst->compar (lo_data, hi_data) > 0)) {
                  memcpy (hi_data, lo_data, size);
                }
                lo_data += size;
                hi_data += size;
              }
            }

            /* close down this peer */
            SC_FREE (peer->buffer);
          }
          peer->received = 1;
        }
      }

      if (remaining2 > 0) {
        mpiret =
          sc_MPI_Waitsome (num_peers, (sc_MPI_Request *) ps->array,
                           &outcount2, wait_indices2, recv_statuses2);
        SC_CHECK_MPI (mpiret);
        SC_ASSERT (outcount2 != sc_MPI_UNDEFINED);
        SC_ASSERT (outcount2 > 0);
        for (i = 0; i < outcount2; ++i) {
          size_t              zz;
          char               *lo_data, *hi_data;

          /* retrieve peer information */
          peer =
            (sc_psort_peer_t *) sc_array_index_int (pa, wait_indices2[i]);

          if (peer->received) {
            SC_ASSERT (!peer->sent);

            /* comparisons with remote peer */
            if (rank < peer->prank) {
              lo_data = peer->my_start;
              hi_data = peer->buffer;
              for (zz = 0; zz < peer->length; ++zz) {
                if (dir == (pst->compar (lo_data, hi_data) > 0)) {
                  memcpy (lo_data, hi_data, size);
                }
                lo_data += size;
                hi_data += size;
              }
            }
            else {
              lo_data = peer->buffer;
              hi_data = peer->my_start;
              for (zz = 0; zz < peer->length; ++zz) {
                if (dir == (pst->compar (lo_data, hi_data) > 0)) {
                  memcpy (hi_data, lo_data, size);
                }
                lo_data += size;
                hi_data += size;
              }
            }

            /* close down this peer */
            SC_FREE (peer->buffer);
          }
          peer->sent = 1;
        }
      }
    }
    SC_ASSERT (remaining == 0);
    SC_ASSERT (remaining2 == 0);
    SC_FREE (recv_statuses);
    SC_FREE (wait_indices);
    SC_FREE (recv_statuses2);
    SC_FREE (wait_indices2);

    /* clean up */
    if (num_peers > 0) {
      mpiret = sc_MPI_Waitall (num_peers, (sc_MPI_Request *) ps->array,
                               sc_MPI_STATUSES_IGNORE);
      SC_CHECK_MPI (mpiret);
    }
    sc_array_reset (pa);
    sc_array_reset (pr);
    sc_array_reset (ps);

    /* recursive merge */
    sc_merge_bitonic (pst, lo, lo + n2, dir);
    sc_merge_bitonic (pst, lo + n2, hi, dir);
  }
}

static void
sc_psort_bitonic (sc_psort_t * pst, size_t lo, size_t hi, int dir)
{
  const size_t        n = hi - lo;

  if (n > 1 && pst->my_hi > lo && pst->my_lo < hi) {
    if (lo >= pst->my_lo && hi <= pst->my_hi) {
#ifndef SC_HAVE_QSORT_R
      qsort (pst->my_base + (lo - pst->my_lo) * pst->size,
             n, pst->size, dir ? sc_compare : sc_icompare);
#else
#ifndef SC_HAVE_BSD_QSORT_R
      qsort_r (pst->my_base + (lo - pst->my_lo) * pst->size,
               n, pst->size, dir ? sc_compare_r : sc_icompare_r, pst);

#else
      qsort_r (pst->my_base + (lo - pst->my_lo) * pst->size,
               n, pst->size, pst, dir ? sc_compare_r : sc_icompare_r);
#endif
#endif
    }
    else {
      const size_t        n2 = n / 2;

      sc_psort_bitonic (pst, lo, lo + n2, !dir);
      sc_psort_bitonic (pst, lo + n2, hi, dir);
      sc_merge_bitonic (pst, lo, hi, dir);
    }
  }
}

void
sc_psort (sc_MPI_Comm mpicomm, void *base, size_t *nmemb, size_t size,
          int (*compar) (const void *, const void *))
{
  int                 mpiret;
  int                 num_procs, rank;
  int                 i;
  size_t              total;
  size_t             *gmemb;
  sc_psort_t          pst;

#ifndef SC_HAVE_QSORT_R
  SC_ASSERT (sc_compare == NULL);
#endif

  /* get basic MPI information */
  mpiret = sc_MPI_Comm_size (mpicomm, &num_procs);
  SC_CHECK_MPI (mpiret);
  mpiret = sc_MPI_Comm_rank (mpicomm, &rank);
  SC_CHECK_MPI (mpiret);

  /* alloc global offset array */
  gmemb = SC_ALLOC (size_t, num_procs + 1);
  gmemb[0] = 0;
  for (i = 0; i < num_procs; ++i) {
    gmemb[i + 1] = gmemb[i] + nmemb[i];
  }

  /* set up internal state and call recursion */
  pst.mpicomm = mpicomm;
  pst.num_procs = num_procs;
  pst.rank = rank;
  pst.size = size;
  pst.my_lo = gmemb[rank];
  pst.my_hi = gmemb[rank + 1];
  pst.my_count = nmemb[rank];
  SC_ASSERT (pst.my_lo + pst.my_count == pst.my_hi);
  pst.gmemb = gmemb;
  pst.my_base = (char *) base;
  pst.compar = compar;
#ifndef SC_HAVE_QSORT_R
  sc_compare = compar;
#endif
  total = gmemb[num_procs];
  SC_GLOBAL_LDEBUGF ("Total values to sort %lld\n", (long long) total);
  sc_psort_bitonic (&pst, 0, total, 1);

  /* clean up and free memory */
#ifndef SC_HAVE_QSORT_R
  sc_compare = NULL;
#endif
  SC_FREE (gmemb);
}
