/**
 * @file xt_mpi_stripe_parse_func.h
 *
 * @copyright Copyright  (C)  2022 Jörg Behrens <behrens@dkrz.de>
 *                                 Moritz Hanke <hanke@dkrz.de>
 *                                 Thomas Jahns <jahns@dkrz.de>
 *
 * @author Jörg Behrens <behrens@dkrz.de>
 *         Moritz Hanke <hanke@dkrz.de>
 *         Thomas Jahns <jahns@dkrz.de>
 */
/*
 * Keywords:
 * Maintainer: Jörg Behrens <behrens@dkrz.de>
 *             Moritz Hanke <hanke@dkrz.de>
 *             Thomas Jahns <jahns@dkrz.de>
 * URL: https://dkrz-sw.gitlab-pages.dkrz.de/yaxt/
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are  permitted provided that the following conditions are
 * met:
 *
 * Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 *
 * Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 *
 * Neither the name of the DKRZ GmbH nor the names of its contributors
 * may be used to endorse or promote products derived from this software
 * without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif


#define XT_TOKEN_PASTE2_(a,b) a##b
#define XT_TOKEN_PASTE2(a,b) XT_TOKEN_PASTE2_(a,b)
#define XT_TOKEN_PASTE3_(a,b,c) a##b##c
#define XT_TOKEN_PASTE3(a,b,c) XT_TOKEN_PASTE3_(a,b,c)

#define XT_MPI_OFFSET_EXT \
  struct XT_TOKEN_PASTE3(Xt_,XT_MPI_STRP_PRS_PREFIX,offset_ext)
#define XT_MPI_STRP_PRS_MATCH_BLOCK_VEC \
  XT_TOKEN_PASTE2(XT_MPI_STRP_PRS_PREFIX,match_block_vec)
#define XT_MPI_STRP_PRS_MATCH_INDEXED \
  XT_TOKEN_PASTE2(XT_MPI_STRP_PRS_PREFIX,match_indexed)
#define XT_MPI_STRP_PRS_MATCH_SIMPLE_VEC \
  XT_TOKEN_PASTE2(XT_MPI_STRP_PRS_PREFIX,match_simple_vec)
#define XT_MPI_STRP_PRS_MATCH_CONTIGUOUS \
  XT_TOKEN_PASTE2(XT_MPI_STRP_PRS_PREFIX,match_contiguous)
#define XT_MPI_STRP_PRS_GEN_FALLBACK_TYPE                \
  XT_TOKEN_PASTE2(XT_MPI_STRP_PRS_PREFIX,gen_fallback_type)
#define XT_MPI_STRP_PRS_ENTRY \
  XT_TOKEN_PASTE3(parse_,XT_MPI_STRP_PRS_PREFIX,stripe)
#define XT_MPI_STRP_PRS_ENTRY_LIBINTERNAL \
  XT_TOKEN_PASTE3(xt_mpi_parse_,XT_MPI_STRP_PRS_PREFIX,stripe)
#define XT_MPI_STRP_PRS_DRIVER \
  XT_TOKEN_PASTE3(xt_mpi_generate_datatype_,XT_MPI_STRP_PRS_PREFIX,stripe)

/**
 * @return true if matched, false if not matched
 */
static bool
XT_MPI_STRP_PRS_MATCH_BLOCK_VEC(
  size_t *pstart_,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_disp,
  const int *v_size,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_stride,
  size_t vlen,
  MPI_Aint *disp, MPI_Datatype *dt,
  struct Xt_mpi_strp_prs_params *params)
{
  // using at least 3 vectors
  size_t p = *pstart_, pstart = p;
  if (p+2 >= vlen || v_stride[p] != XT_MPI_STRP_PRS_UNITSTRIDE
      || v_stride[p+1] != XT_MPI_STRP_PRS_UNITSTRIDE ) return false;
  int bl = v_size[p];
  assert(bl > 0);
  if (v_size[p+1] != bl) return false;

  XT_MPI_STRP_PRS_AOFS_TYPE vstride = v_disp[p+1] - v_disp[p];

  p += 2;
  while( p < vlen && v_stride[p] == XT_MPI_STRP_PRS_UNITSTRIDE
         && v_size[p] == bl && v_disp[p] - v_disp[p-1] == vstride ) {
    p++;
  }
  size_t n = p - pstart;
  if (n<3) return false;
  *pstart_ = p;

  XT_MPI_STRP_PRS_AOFS_TYPE disp_ = n == vlen ? 0 : v_disp[pstart];
  *disp = XT_MPI_STRP_PRS_DISP_ADJUST(disp_);

  *dt = XT_MPI_STRP_PRS_BLOCK_VEC_CREATE(
    &params->ddt_list, (int)n, bl, vstride, params->old_type, params->comm);

  XT_MPI_STRP_PRS_AOFS_TYPE start = v_disp[pstart] - disp_;

  if (start) {
    MPI_Datatype dt1 = *dt;
    // (start != 0) => add offset:
    xt_mpi_call(MPI_Type_create_hindexed(
                  1, &(int){1},
                  &(MPI_Aint){XT_MPI_STRP_PRS_DISP_ADJUST(start)}, dt1, dt),
                params->comm);
    Xt_mpi_ddt_cache_entry_release(&params->ddt_list, &dt1, params->comm);
  }
  return n != 0;
}

static bool
XT_MPI_STRP_PRS_MATCH_INDEXED(
  size_t *pstart_,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_disp,
  const int *v_size,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_stride,
  size_t vlen,
  MPI_Aint *disp, MPI_Datatype *dt,
  struct Xt_mpi_strp_prs_params *params) {
  // we only accept non-trivial matches
  size_t p = *pstart_, pstart = p;
  if (p >= vlen || v_stride[p] != XT_MPI_STRP_PRS_UNITSTRIDE || v_size[p] < 2)
    return false;

  do
    ++p;
  while (p < vlen && v_stride[p] == XT_MPI_STRP_PRS_UNITSTRIDE);

  size_t n = p - pstart;

  if (n < 2) return false;
  *pstart_ = p;

  XT_MPI_STRP_PRS_AOFS_TYPE start = n == vlen ? 0 : v_disp[pstart];
  *disp = XT_MPI_STRP_PRS_DISP_ADJUST(start);
  XT_MPI_STRP_PRS_AOFS_TYPE *restrict d
    = xmalloc(n * sizeof (int) + n * sizeof (*d));
  int *restrict bl = (int * restrict)(d + n);
  bool hom_bl = true;
  d[0] = v_disp[pstart] - start;
  int bl0 = bl[0] = v_size[pstart];
  for (size_t i = 1; i < n; i++) {
    size_t iv = pstart + i;
    d[i] = v_disp[iv] - start;
    bl[i] = v_size[iv];
    hom_bl &= (bl[i] == bl0);
  }

  if (hom_bl) {
    *dt = XT_MPI_STRP_PRS_INDEXED_BLOCK_CREATE(
      &params->ddt_list, (int)n, bl0, d, params->old_type, params->comm);
  } else {
    *dt = XT_MPI_STRP_PRS_INDEXED_CREATE(
      &params->ddt_list, (int)n, bl, d, params->old_type, params->comm);
  }

  free(d);
  return n != 0;
}

static bool
XT_MPI_STRP_PRS_MATCH_SIMPLE_VEC(
  size_t *pstart_,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_disp,
  const int *v_size,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_stride,
  size_t vlen,
  MPI_Aint *disp, MPI_Datatype *dt,
  struct Xt_mpi_strp_prs_params *params) {
  // we only accept non-trivial matches (nsteps>2) with stride /= 1
  // using only one vector from v
  size_t p = *pstart_;
  if (p >= vlen) return false;
  int nstrides = v_size[p];
  XT_MPI_STRP_PRS_AOFS_TYPE stride = v_stride[p];
  if (nstrides < 2 || stride == XT_MPI_STRP_PRS_UNITSTRIDE ) return false;

  *pstart_ = p + 1;

  XT_MPI_STRP_PRS_AOFS_TYPE disp_ = vlen > 1 ? v_disp[p] : 0;
  *disp = XT_MPI_STRP_PRS_DISP_ADJUST(disp_);

  *dt = XT_MPI_STRP_PRS_BLOCK_VEC_CREATE(
    &params->ddt_list, nstrides, 1, stride, params->old_type, params->comm);

  XT_MPI_STRP_PRS_AOFS_TYPE start = v_disp[p] - disp_;
  if (start) {
    MPI_Datatype dt1 = *dt;

    // (start != 0) => add offset:
    xt_mpi_call(MPI_Type_create_hindexed(
                  1, &(int){1},
                  &(MPI_Aint){XT_MPI_STRP_PRS_DISP_ADJUST(start)}, dt1, dt),
                params->comm);
    Xt_mpi_ddt_cache_entry_release(&params->ddt_list, &dt1, params->comm);
  }
  return nstrides != 0;
}

static bool
XT_MPI_STRP_PRS_MATCH_CONTIGUOUS(
  size_t *pstart_,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_disp,
  const int *v_size,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_stride,
  size_t vlen,
  MPI_Aint *restrict disp, MPI_Datatype *dt,
  struct Xt_mpi_strp_prs_params *params) {
  size_t p = *pstart_;
  if (p >= vlen || v_stride[p] != XT_MPI_STRP_PRS_UNITSTRIDE || v_size[p] < 2)
    return false;

  XT_MPI_STRP_PRS_AOFS_TYPE disp_ = vlen > 1 ? v_disp[p] : 0;
  *disp = XT_MPI_STRP_PRS_DISP_ADJUST(disp_);
  XT_MPI_STRP_PRS_AOFS_TYPE d = v_disp[p] - disp_;

  if (!d)
    *dt = Xt_mpi_ddt_cache_acquire_contiguous(
      &params->ddt_list, v_size[p], params->old_type, params->comm);
  else
    *dt = XT_MPI_STRP_PRS_INDEXED_BLOCK_CREATE(
      &params->ddt_list, 1, v_size[p], &d, params->old_type, params->comm);

  *pstart_ = p+1;
  return true;
}

static void
XT_MPI_STRP_PRS_GEN_FALLBACK_TYPE(
  size_t set_start, size_t set_end,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_disp,
  const int *v_size,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_stride,
  size_t vlen,
  MPI_Aint *disp,
  MPI_Datatype *dt, struct Xt_mpi_strp_prs_params *params) {
  size_t ia = set_start;
  size_t ib = set_end;
  if (ib <= ia || ib > vlen) return;

  int n = 0;
  for (size_t i=ia; i < ib; i++)
    n += v_size[i];

  /* todo: given the guarantees for v that fceb584 introduced,
   * this check should never fire */
  assert(n>0);

  // generate absolute datatype if ia == 0 && ib == vlen,
  // else generate relative datatype that gets embedded by the caller
  XT_MPI_STRP_PRS_AOFS_TYPE start = (ia == 0 && ib == vlen) ? 0 : v_disp[ia];

  *disp = XT_MPI_STRP_PRS_DISP_ADJUST(start);

  XT_MPI_STRP_PRS_AOFS_TYPE *restrict d = xmalloc(sizeof (*d) * (size_t)n);
  size_t p=0;
#ifndef NDEBUG
  /* did any element of v have non-positive size? */
  bool found_np = false;
#endif

  for (size_t i=ia; i < ib; i++) {
#ifndef NDEBUG
    found_np |= v_size[i] <= 0;
#endif
    size_t v_i_size = (size_t)(v_size[i] > 0 ? v_size[i] : 0);
    for (size_t k=0; k < v_i_size; k++) {
      d[p] = v_disp[i] + (XT_MPI_STRP_PRS_AOFS_TYPE)k * v_stride[i] - start;
      p++;
    }
  }
  assert(!found_np);

  if (n==1 && d[0] == 0) {
    *dt = params->old_type;
  } else {
    *dt = XT_MPI_STRP_PRS_INDEXED_BLOCK_CREATE(
      &params->ddt_list, n, 1, d, params->old_type, params->comm);
  }
  free(d);
}

static struct Xt_mpi_dt_ofs
XT_MPI_STRP_PRS_ENTRY(
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_disp,
  const int *v_size,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_stride,
  size_t vlen,
  struct Xt_mpi_strp_prs_params *params);

struct Xt_mpi_dt_ofs
XT_MPI_STRP_PRS_ENTRY_LIBINTERNAL(
  size_t vlen,
  const XT_MPI_STRP_PRS_AOFS_TYPE v_disp[vlen],
  const int v_size[vlen],
  const XT_MPI_STRP_PRS_AOFS_TYPE v_stride[vlen],
  struct Xt_mpi_strp_prs_params *params)
{
  return XT_MPI_STRP_PRS_ENTRY(v_disp, v_size, v_stride, vlen, params);
}

static struct Xt_mpi_dt_ofs
XT_MPI_STRP_PRS_ENTRY(
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_disp,
  const int *v_size,
  const XT_MPI_STRP_PRS_AOFS_TYPE *v_stride,
  size_t vlen,
  struct Xt_mpi_strp_prs_params *params)
{
  /* [set_start,set_end) describes the prefix of non-matching
   * elements in v that then need to be handled with gen_fallback_type */
  size_t set_start = 0, set_end = 0;
  MPI_Aint *restrict wdisp
    = xmalloc(sizeof(MPI_Datatype) * vlen + sizeof (MPI_Aint) * vlen);
  MPI_Datatype *restrict wdt = (MPI_Datatype *)(wdisp + vlen);
  /* [p,vlen) is the part of v that still needs matching performed */
  /* m is the index of the next datatype and displacements to write
   * to wdt and wdisp respectively */
  size_t p = 0, m = 0;
  while (p<vlen) {
    /* depending on whether there is a non-empty prefix, the datatype
     * and displacement corresponding to a match need to be written
     * to wdt[m+1] and wdisp[m+1] or wdt[m] and wdisp[m] respectively */
    size_t mm = m + (set_start < set_end);
    if ((XT_MPI_STRP_PRS_MATCH_BLOCK_VEC(
           &p, v_disp, v_size, v_stride, vlen,
           wdisp+mm, wdt+mm, params))
        || (XT_MPI_STRP_PRS_MATCH_INDEXED(
              &p, v_disp, v_size, v_stride, vlen,
              wdisp+mm, wdt+mm, params))
        || (XT_MPI_STRP_PRS_MATCH_SIMPLE_VEC(
              &p, v_disp, v_size, v_stride, vlen,
              wdisp+mm, wdt+mm, params))
        || (XT_MPI_STRP_PRS_MATCH_CONTIGUOUS(
              &p, v_disp, v_size, v_stride, vlen,
              wdisp+mm, wdt+mm, params)) ) {
      /* in case a match is found, generate fallback datatype for
       * non-matching, preceding extents */
      if (set_start < set_end) {
        XT_MPI_STRP_PRS_GEN_FALLBACK_TYPE(
          set_start, set_end, v_disp, v_size, v_stride, vlen,
          wdisp+m, wdt+m, params);
        m++;
      }
      m++;
      set_start = p;
    } else {
      /* assign ext investigated last to prefix */
      set_end = ++p;
    }
  }
  if (set_start <  set_end) {
    XT_MPI_STRP_PRS_GEN_FALLBACK_TYPE(
      set_start, set_end, v_disp, v_size, v_stride, vlen,
      wdisp+m, wdt+m, params);
    m++;
  }
  size_t wlen = m;
  MPI_Datatype result_dt, old_type = params->old_type;
  MPI_Aint disp0;
  if (wlen == 1 ) {
    disp0 = wdisp[0];
    if (wdt[0] == old_type && !xt_mpi_ddt_is_named(old_type))
      xt_mpi_call(MPI_Type_dup(old_type, wdt), params->comm);
    result_dt = wdt[0];
  } else {
    int *restrict wblocklength
      = wlen * sizeof (int) <= (vlen - wlen) * sizeof (*wdt)
      ? (void *)(wdt + wlen) : xmalloc(wlen * sizeof (*wblocklength));
    disp0 = wdisp[0];
    for(size_t i=0; i<wlen; i++) {
      wdisp[i] -= disp0;
      wblocklength[i] = 1;
    }
    result_dt = Xt_mpi_ddt_cache_acquire_struct(
      &params->ddt_list, (int)wlen, wblocklength, wdisp, wdt, params->comm);
    if (wlen * sizeof (int) > (vlen - wlen) * sizeof (*wdt))
      free(wblocklength);
    for (size_t i = 0; i < wlen; i++)
      if (wdt[i] != old_type)
        Xt_mpi_ddt_cache_entry_release(&params->ddt_list, wdt+i, params->comm);
  }
  xt_mpi_call(MPI_Type_commit(&result_dt), params->comm);
  free(wdisp);
  return (struct Xt_mpi_dt_ofs){ .dt = result_dt, .displacement = disp0 };
}

MPI_Datatype
XT_MPI_STRP_PRS_DRIVER(const XT_MPI_OFFSET_EXT *v,
                       int count, MPI_Datatype old_type,
                       MPI_Comm comm)
{
  size_t count_ = (size_t)0;
  for (int i=0; i<count; ++i)
    count_ += (size_t)(v[i].size > 0);
  if (count_ < 1) return MPI_DATATYPE_NULL;
  XT_MPI_STRP_PRS_AOFS_TYPE *restrict v_disp
    = xmalloc(2 * count_ * sizeof (*v_disp) + count_ * sizeof (int)),
    *restrict v_stride = v_disp+count_;
  int *restrict v_size = (int *)(v_stride + count_);
  if (count_) {
    for (size_t i=0, j=0; i<(size_t)count; ++i) {
      v_disp[j] = v[i].start;
      v_stride[j] = v[i].stride;
      v_size[j] = v[i].size;
      j += v[i].size > 0;
    }
  } else {
    v_disp = v_stride = NULL;
    v_size = NULL;
  }
  struct Xt_mpi_strp_prs_params params;
  xt_init_mpi_strp_prs_params(&params, old_type, &xt_default_config, comm);
  struct Xt_mpi_dt_ofs dt_ofs
    = XT_MPI_STRP_PRS_ENTRY(v_disp, v_size, v_stride, count_, &params);
  if (dt_ofs.displacement != 0) {
    MPI_Datatype datatype = dt_ofs.dt;
    Xt_Type_create_hindexed(1, (int [1]){ 1 },
                            &dt_ofs.displacement, datatype, &dt_ofs.dt, comm);
    xt_mpi_call(MPI_Type_free(&datatype), comm);
    xt_mpi_call(MPI_Type_commit(&dt_ofs.dt), comm);
  }
  if (count_)
    free(v_disp);
  return dt_ofs.dt;
}


#undef XT_MPI_STRP_PRS_ENTRY
#undef XT_MPI_STRP_PRS_GEN_FALLBACK_TYPE
#undef XT_MPI_STRP_PRS_MATCH_CONTIGUOUS
#undef XT_MPI_STRP_PRS_MATCH_SIMPLE_VEC
#undef XT_MPI_STRP_PRS_MATCH_INDEXED
#undef XT_MPI_STRP_PRS_MATCH_BLOCK_VEC
#undef XT_MPI_OFFSET_EXT

#undef XT_TOKEN_PASTE2
#undef XT_TOKEN_PASTE2_
#undef XT_TOKEN_PASTE3
#undef XT_TOKEN_PASTE3_

/*
 * Local Variables:
 * c-basic-offset: 2
 * coding: utf-8
 * indent-tabs-mode: nil
 * show-trailing-whitespace: t
 * require-trailing-newline: t
 * End:
 */
