/* nag_tsa_cp_binary_user (g13nec) Example Program.
 *
 * Copyright 2023 Numerical Algorithms Group.
 *
 * Mark 29.2, 2023.
 */

#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 chgpfn(Nag_TS_SegSide side, Integer u, Integer w,
                            Integer minss, Integer *v, double cost[],
                            Nag_Comm *comm, Integer *info);
#ifdef __cplusplus
}
#endif

static double gamma_costfn(double si, Integer n, double shape, Integer *isinf);
static Integer get_data(Integer n, Nag_Comm *comm);
static void clean_data(Nag_Comm *comm);

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

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

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

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

  printf("nag_tsa_cp_binary_user (g13nec) Example Program Results\n\n");

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

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

  /* Read in other data, that (may be) dependent on the cost function */
  get_data(n, &comm);

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

  /* Call nag_tsa_cp_binary_user (g13nec) to detect change points */
  nag_tsa_cp_binary_user(n, beta, minss, mdepth, chgpfn, &ntau, tau, &comm,
                         &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_tsa_cp_binary_user (g13nec).\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 chgpfn(Nag_TS_SegSide side, Integer u, Integer w,
                            Integer minss, Integer *v, double cost[],
                            Nag_Comm *comm, Integer *info) {
  double shape, ys, this_cost, tcost[2];
  Integer i, nseg1, nseg2, isinf = 0;
  CostInfo *ci;

  ci = (CostInfo *)comm->p;

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

  /* In order to calculate the cost of having a change point at I, we need */
  /* to calculate sum y[j],j=u-1,...,i-1 and sum y[j],j=i,...w-1. We could */
  /* calculate the required values at each call to CHGPFN, but we reuse some */
  /* of these values, so will store the intermediate sums and only */
  /* recalculate the ones we required */

  /* If side = Nag_LeftSubSeg (i.e. we are working with a left hand */
  /* sub-segment), we already have sum y[j-1], j=u,..,i for this value of U, */
  /* so only need the backwards cumulative sum */
  if (side == Nag_FirstSegCall || side == Nag_LeftSubSeg) {
    /* ci->user[2*i] = sum y[j-1],j=i+1,...,w */
    ys = 0.0;
    for (i = w; i > u; i--) {
      ys += ci->y[i - 1];
      comm->user[2 * i - 3] = ys;
    }
  }

  /* Similarly, if side = Nag_RightSubSeg (i.e. we are working with a right */
  /* hand sub-segment), we already have SUM(Y(I+1:W)) for this value of W, */
  /* so only need the forwards cumulative sum */
  if (side == Nag_FirstSegCall || side == Nag_RightSubSeg) {
    /* comm->user[2*i-2] = sum y[j-1], j=u,...,i */
    ys = 0.0;
    for (i = u; i < w + 1; i++) {
      ys += ci->y[i - 1];
      comm->user[2 * i - 2] = ys;
    }
  }

  /* For SIDE = Nag_FirstSegCall we have calculated both sums, and because */
  /* the call with side = Nag_SecondSegCall directly follows we do not need */
  /* to recalculate anything */
  if (side == Nag_FirstSegCall) {
    /* Need to calculate the cost for the full segment */
    cost[0] = gamma_costfn(comm->user[2 * w - 2], w - u + 1, shape, &isinf);
    /* No need to populate the rest of COST or V */
  } else {
    /* Need to find a potential change point */
    *v = 0;
    cost[0] = 0.0;

    /* Calculate the widths of the sub-segment for the left most potential */
    /* change point, ensuring it has length at least minss */
    nseg1 = minss;
    nseg2 = w - u - minss + 1;

    /* Loop over all possible change point locations (conditional on the */
    /* length of both segments having length >= minss) */
    for (i = u + minss - 1; i < w - minss + 1; i++) {
      tcost[0] = gamma_costfn(comm->user[2 * i - 2], nseg1, shape, &isinf);
      tcost[1] = gamma_costfn(comm->user[2 * i - 1], nseg2, shape, &isinf);
      if (isinf != 0) {
        /* Total cost for change point is -Inf, so have found */
        /* the minimum */
        *v = i;
        cost[1] = tcost[0];
        cost[2] = tcost[1];
        cost[0] = (cost[1] < cost[2]) ? cost[1] : cost[2];
        break;
      } else {
        this_cost = tcost[0] + tcost[1];
      }

      if (this_cost < cost[0] || *v == 0) {
        /* Update the proposed change point location */
        *v = i;
        cost[1] = tcost[0];
        cost[2] = tcost[1];
        cost[0] = this_cost;
      }

      /* Update the size of the next segments */
      nseg1++;
      nseg2--;
    }
  }

  /* Store the ISINF flag */
  ci->isinf = isinf;

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

static double gamma_costfn(double si, Integer n, double shape, Integer *isinf) {
  /* Cost function for the gamma distribution */
  double tmp;

  if (si <= 0.0) {
    /* Cost is -Inf */
    *isinf = 1;
    return -X02ALC;
  } else {
    tmp = ((double)n) * shape;
    return 2.0 * tmp * (log(si) - log(tmp));
  }
}

static Integer get_data(Integer n, 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;
  comm->user = 0;
  if (!(ci = NAG_ALLOC(1, CostInfo))) {
    printf("Allocation failure\n");
    return -1;
  }

  /* Read in the series of interest */
  if (!(ci->y = NAG_ALLOC(n, double))) {
    printf("Allocation failure\n");
    return -1;
  }
  for (i = 0; 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;

  /* Allocate some workspace into comm that we will be using later */
  if (!(comm->user = NAG_ALLOC(2 * n, double))) {
    printf("Allocation failure\n");
    return -1;
  }

  /* 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);
  NAG_FREE(comm->user);
}