/* nag_sort_rank_sort (m01dsc) Example Program.
 *
 * Copyright 2024 Numerical Algorithms Group.
 *
 * Mark 30.2, 2024.
 *
 */

#include <nag.h>
#include <stdio.h>

#ifdef __cplusplus
extern "C" {
#endif
static Integer NAG_CALL compare(const Nag_Pointer a, const Nag_Pointer b);
static Integer weighted_quantiles(Integer n, const double vec[],
                                  const double weights[], Integer nq,
                                  double quant[], double quant_vec[],
                                  NagError *fail);
static int ex1(void);
static int ex2(void);

#ifdef __cplusplus
}
#endif

int main(void) {
  Integer exit_status_ex1;
  Integer exit_status_ex2;

  printf("nag_sort_rank_sort (m01dsc) Example Program Results\n\n");

  /* Skip heading in data file */
  scanf("%*[^\n] ");

  exit_status_ex1 = ex1();
  exit_status_ex2 = ex2();

  return (exit_status_ex1 == 0 && exit_status_ex2 == 0) ? 0 : 1;
}

static int ex1(void) {
  Integer exit_status = 0;
  NagError fail;
  double *vec = 0;
  ptrdiff_t step;
  size_t i, n, *rank = 0, step_u;

  INIT_FAIL(fail);

  printf("Example 1\n");

  /* Skip heading in data file */
  scanf(" %*[^\n] ");
  scanf("%" NAG_UFMT "%" NAG_UFMT "", &n, &step_u);
  step = (ptrdiff_t)step_u;
  if (n >= 1) {
    if (!(vec = NAG_ALLOC(n, double)) || !(rank = NAG_ALLOC(n, size_t))) {
      printf("Allocation failure\n");
      exit_status = -1;
      goto END;
    }
  } else {
    printf("Invalid n or step.\n");
    exit_status = 1;
    return exit_status;
  }
  for (i = 0; i < n; ++i)
    scanf("%lf", &vec[i]);
  /* nag_sort_rank_sort (m01dsc).
   * Rank sort of set of values of arbitrary data type
   */
  nag_sort_rank_sort((Pointer)vec, n, step * (ptrdiff_t)(sizeof(double)),
                     compare, Nag_Ascending, rank, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_sort_rank_sort (m01dsc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }
  printf("    Data      Rank\n");
  for (i = 0; i < n; ++i)
    printf("   %7.4f   %4" NAG_UFMT "\n", vec[i], rank[i]);
END:
  NAG_FREE(vec);
  NAG_FREE(rank);

  return exit_status;
}

int ex2(void) {
  /* Scalars */
  Integer exit_status = 0, i, n, nq;
  /* Arrays */
  double *vec = 0, *quant = 0, *quant_vec = 0, *weights = 0;
  /* Nag Types */
  NagError fail;

  INIT_FAIL(fail);

  printf("\n\nExample 2\n");

  /* Skip heading in data file */
  scanf(" %*[^\n] ");
  scanf("%" NAG_IFMT "%" NAG_IFMT "", &n, &nq);
  if (n >= 1 && nq >= 1) {
    if (!(vec = NAG_ALLOC(n, double)) || !(quant = NAG_ALLOC(nq, double)) ||
        !(quant_vec = NAG_ALLOC(nq, double)) ||
        !(weights = NAG_ALLOC(n, double))) {
      printf("Allocation failure\n");
      exit_status = -1;
      goto END;
    }
  } else {
    if (n < 1) {
      printf("Invalid n.\n");
    } else {
      printf("Invalid nq.\n");
    }
    exit_status = 1;
    goto END;
  }
  for (i = 0; i < n; ++i)
    scanf("%lf", &vec[i]);
  for (i = 0; i < n; ++i)
    scanf("%lf", &weights[i]);
  for (i = 0; i < nq; ++i)
    scanf("%lf", &quant[i]);

  /* Find quantiles of set of values of data type double */
  exit_status =
      weighted_quantiles(n, vec, weights, nq, quant, quant_vec, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from weighted_quantiles.\n%s\n", fail.message);
    goto END;
  }

  printf("    Quantile   Result\n\n");
  for (i = 0; i < nq; ++i) {
    printf("   %7.4f     %7.4f\n", quant[i], quant_vec[i]);
  }

END:
  NAG_FREE(vec);
  NAG_FREE(quant);
  NAG_FREE(quant_vec);
  NAG_FREE(weights);

  return exit_status;
}

Integer weighted_quantiles(Integer n, const double vec[],
                           const double weights[], Integer nq, double quant[],
                           double quant_vec[], NagError *fail) {
  Integer i, index, exit_status = 0;
  double item, *cum_weights = 0;
  ptrdiff_t step;
  size_t *rank = 0;

  /* Skip heading in data file */
  scanf("%*[^\n]");

  if (!(cum_weights = NAG_ALLOC(n, double)) || !(rank = NAG_ALLOC(n, size_t))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  step = (ptrdiff_t)1;

  /* nag_sort_rank_sort (m01dsc).
   * Rank sort of set of values of arbitrary data type
   */
  nag_sort_rank_sort((Pointer)vec, (size_t)n,
                     step * (ptrdiff_t)(sizeof(double)), compare, Nag_Ascending,
                     rank, fail);
  if (fail->code != NE_NOERROR) {
    printf("Error from nag_sort_rank_sort (m01dsc).\n%s\n", fail->message);
    exit_status = 1;
    goto END;
  }

  /* nag_sort_permute_invert (m01zac).
   * Inverts a permutation converting a rank vector to an
   * index vector or vice versa
   */
  nag_sort_permute_invert(rank, n, fail);
  if (fail->code != NE_NOERROR) {
    printf("Error from nag_sort_permute_invert (m01zac).\n%s\n", fail->message);
    exit_status = 1;
    goto END;
  }

  /* obtain the cumulative weights */
  cum_weights[0] = weights[rank[0]];
  for (i = 1; i < n; ++i)
    cum_weights[i] = cum_weights[i - 1] + weights[rank[i]];

  /* find the quantiles */
  for (i = 0; i < nq; ++i) {
    /* scale the quantiles w.r.t the total cumulative sum */
    item = quant[i] * cum_weights[n - 1];

    /* nag_sort_realvec_search (m01nac)
     * Binary search in set of double precision numbers
     */
    index =
        nag_sort_realvec_search(Nag_FALSE, cum_weights, 0, n - 1, item, fail);
    if (fail->code != NE_NOERROR) {
      printf("Error from nag_sort_realvec_search (m01nac).\n%s\n",
             fail->message);
      exit_status = 1;
      goto END;
    }

    if (index != n - 1)
      index++;

    /* convert the index returned by M01NAF into a value
       from the (sorted) original data */
    quant_vec[i] = vec[rank[index]];
  }

END:
  NAG_FREE(cum_weights);
  NAG_FREE(rank);

  return exit_status;
}

static Integer NAG_CALL compare(const Nag_Pointer a, const Nag_Pointer b) {
  double x = *((const double *)a);
  double y = *((const double *)b);
  return (x < y ? -1 : (x == y ? 0 : 1));
}