/* nag_tsa_cp_pelt_user (g13nbc) Example Program.
 *
 * Copyright 2024 Numerical Algorithms Group.
 *
 * Mark 30.2, 2024.
 */

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

/* Structure to hold extra information that the cost function requires */
typedef struct {
  Integer isinf;
  double shape;
  double *y;
} CostInfo;

/* Functions that are dependent on the cost function used */
#ifdef __cplusplus
extern "C" {
#endif
static void NAG_CALL costfn(Integer ts, Integer nr, const Integer r[],
                            double c[], Nag_Comm *comm, Integer *info);
#ifdef __cplusplus
}
#endif

static Integer get_data(Integer n, double *k, Nag_Comm *comm);
static void clean_data(Nag_Comm *comm);

int main(void) {
  /* Integer scalar and array declarations */
  Integer i, minss, n, ntau;
  Integer exit_status = 0;
  Integer *tau = 0;

  /* NAG structures and types */
  NagError fail;
  Nag_Comm comm;

  /* Double scalar and array declarations */
  double beta, k;

  /* Initialize the error structure */
  INIT_FAIL(fail);

  printf("nag_tsa_cp_pelt_user (g13nbc) Example Program Results\n\n");

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

  /* Read in the problem size, penalty and minimum segment size */
  scanf("%" NAG_IFMT "%lf%" NAG_IFMT "%*[^\n] ", &n, &beta, &minss);

  /* Read in other data, that (may be) dependent on the cost function */
  if (get_data(n, &k, &comm)) {
    printf("Set up failure\n");
    exit_status = 2;
    goto END;
  }

  /* Allocate output arrays */
  if (!(tau = NAG_ALLOC(n, Integer))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Call nag_tsa_cp_pelt_user (g13nbc) to detect change points */
  nag_tsa_cp_pelt_user(n, beta, minss, k, costfn, &ntau, tau, &comm, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_tsa_cp_pelt_user (g13nbc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Display the results */
  printf("  -- Change Points --\n");
  printf("  Number     Position\n");
  printf(" =====================\n");
  for (i = 0; i < ntau; i++) {
    printf(" %4" NAG_IFMT "       %6" NAG_IFMT "\n", i + 1, tau[i]);
  }

END:
  NAG_FREE(tau);
  clean_data(&comm);

  return (exit_status);
}

static void NAG_CALL costfn(Integer ts, Integer nr, const Integer r[],
                            double c[], Nag_Comm *comm, Integer *info) {
  double dn, shape, si;
  Integer i;
  CostInfo *ci;

  ci = (CostInfo *)comm->p;

  /* Get the shape parameter for the gamma distribution from comm */
  shape = ci->shape;

  /* Test which way around TS and R are (only needs to be done once) */
  if (ts < r[0]) {
    for (i = 0; i < nr; i++) {
      si = ci->y[r[i]] - ci->y[ts];

      if (si <= 0.0) {
        /* -Inf */
        ci->isinf = 1;
        c[i] = -X02ALC;
      } else {
        dn = (double)(r[i] - ts);
        c[i] = 2.0 * dn * shape * (log(si) - log(dn * shape));
      }
    }
  } else {
    for (i = 0; i < nr; i++) {
      si = ci->y[ts] - ci->y[r[i]];

      if (si <= 0.0) {
        /* -Inf */
        ci->isinf = 1;
        c[i] = -X02ALC;
      } else {
        dn = (double)(ts - r[i]);
        c[i] = 2.0 * dn * shape * (log(si) - log(dn * shape));
      }
    }
  }

  /* Set info nonzero to terminate execution for any reason */
  *info = 0;
}

static Integer get_data(Integer n, double *k, Nag_Comm *comm) {
  /* Read in data that is specific to the cost function */
  double shape;
  Integer i;
  CostInfo *ci;

  /* Allocate some memory for the additional information structure */
  /* This will be pointed to by comm->p */
  comm->p = 0;
  if (!(ci = NAG_ALLOC(1, CostInfo))) {
    printf("Allocation failure\n");
    return -1;
  }

  /* Read in the series of interest */
  /* NB: we are allocating n+1 elements to y as we manipulate the data
     in Y in a moment */
  if (!(ci->y = NAG_ALLOC(n + 1, double))) {
    printf("Allocation failure\n");
    return -1;
  }

  /* Referencing y from 1 here to aid manipulation later */
  for (i = 1; i <= n; i++)
    scanf("%lf", &(ci->y)[i]);
  scanf("%*[^\n] ");

  /* Read in the shape parameter for the Gamma distribution */
  scanf("%lf%*[^\n] ", &shape);

  /* Store the shape parameter in CostInfo structure */
  ci->shape = shape;

  /* Set the warning flag to 0 */
  ci->isinf = 0;

  /* The cost function is a function of the sum of y, so for efficiency we will
     calculate the cumulative sum. It should be noted that this may introduce
     some rounding issues with very extreme data */
  (ci->y)[0] = 0.0;
  for (i = 1; i <= n; i++)
    (ci->y)[i] += (ci->y)[i - 1];

  /* The value of k is defined by the cost function being used in this example
     a value of 0.0 is the required value */
  *k = 0.0;

  /* Store pointer to CostInfo structure in Nag_Comm */
  comm->p = (void *)ci;

  return 0;
}

static void clean_data(Nag_Comm *comm) {
  /* Free any memory allocated in get_data */
  CostInfo *ci;

  if (comm->p) {
    ci = (CostInfo *)comm->p;
    NAG_FREE(ci->y);
  }

  NAG_FREE(comm->p);
}