/* nag_correg_lars_xtx (g02mbc) Example Program.
 *
 * Copyright 2022 Numerical Algorithms Group.
 *
 * Mark 28.5, 2022.
 */
/* Pre-processor includes */
#include <nag.h>
#include <stdio.h>

int main(void) {
  /* Integer scalar and array declarations */
  Integer i, j, k, ip, ldb, lddtd, m, mnstep, n, nstep, lropt, pm, pm2, pddy;
  Integer *isx = 0;
  Integer exit_status = 0;

  /* NAG structures and types */
  NagError fail;
  Nag_LARSModelType mtype;
  Nag_LARSPreProcess intcpt, pred;
  Nag_SumSquare mean;

  /* Double scalar and array declarations */
  double sw;
  double *b = 0, *dtd = 0, *fitsum = 0, *dy = 0, *ropt = 0, *wmean = 0;

  /* Character scalar and array declarations */
  char cmtype[40], cpred[40], cintcpt[40];

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

  printf("nag_correg_lars_xtx (g02mbc) Example Program Results\n\n");

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

  /* Read in the problem size */
  scanf("%" NAG_IFMT "%" NAG_IFMT "%*[^\n] ", &n, &m);

  /* Read in the model specification */
  scanf("%39s%39s%39s%" NAG_IFMT "%*[^\n] ", cmtype, cpred, cintcpt, &mnstep);
  mtype = (Nag_LARSModelType)nag_enum_name_to_value(cmtype);
  pred = (Nag_LARSPreProcess)nag_enum_name_to_value(cpred);
  intcpt = (Nag_LARSPreProcess)nag_enum_name_to_value(cintcpt);

  /* Using all variables */
  isx = 0;
  ip = m;

  /* Optional arguments (using defaults) */
  lropt = 0;
  ropt = 0;

  /* Allocate memory for the augmented matrix [D y] and its cross-product */
  pddy = n;
  if (!(dy = NAG_ALLOC(pddy * (m + 1), double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Read in the augmented matrix [D y] and calculate cross-product matrices
     (NB: Datasets with a large number of observations can be split into
     blocks with the resulting cross-product matrices being combined
     using g02bzc) */
  for (i = 0; i < n; i++) {
    for (j = 0; j < m + 1; j++) {
      scanf("%lf", &dy[j * pddy + i]);
    }
  }
  scanf("%*[^\n] ");

  pm = m * (m + 1) / 2;
  pm2 = (m + 1) * (m + 2) / 2 - 1;

  /* We are calculating the cross-product matrix using nag_correg_ssqmat
     (g02buc) which returns it in packed storage */
  lddtd = 1;
  if (!(wmean = NAG_ALLOC(m + 1, double)) ||
      !(dtd = NAG_ALLOC(pm2 + 1, double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  mean = (intcpt == Nag_LARS_Centered) ? Nag_AboutMean : Nag_AboutZero;

  /* Calculate the cross-product matrices using g02buc */
  nag_correg_ssqmat(Nag_ColMajor, mean, n, m + 1, dy, pddy, NULL, &sw, wmean,
                    dtd, NAGERR_DEFAULT);
  /* The first PM+1 elements of dtd contain the cross-products of D
     elements PM to PM2-1 contains cross-product of D with y and element PM2
     contains cross-product of y with itself */

  /* Allocate output arrays */
  ldb = ip;
  if (!(b = NAG_ALLOC(ldb * (mnstep + 1), double)) ||
      !(fitsum = NAG_ALLOC(6 * (mnstep + 1), double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Call g02mbc to the model */
  nag_correg_lars_xtx(mtype, pred, intcpt, n, m, dtd, lddtd, isx, &dtd[pm],
                      dtd[pm2], mnstep, &ip, &nstep, b, ldb, fitsum, ropt,
                      lropt, &fail);
  if (fail.code != NE_NOERROR) {
    if (fail.code != NW_OVERFLOW_WARN && fail.code != NW_POTENTIAL_PROBLEM &&
        fail.code != NW_LIMIT_REACHED) {
      printf("Error from nag_correg_lars_xtx (g02mbc).\n%s\n", fail.message);
      exit_status = 1;
      goto END;
    } else {
      printf("Warning from nag_correg_lars_xtx (g02mbc).\n%s\n", fail.message);
      exit_status = 2;
    }
  }

  /* Display the parameter estimates */
  printf("  Step ");
  for (i = 0; i < MAX(ip - 2, 0) * 5; i++)
    printf(" ");
  printf(" Parameter Estimate\n ");
  for (i = 0; i < 5 + ip * 10; i++)
    printf("-");
  printf("\n");
  for (k = 0; k < nstep; k++) {
    printf("  %3" NAG_IFMT "", k + 1);
    for (j = 0; j < ip; j++) {
      printf(" %9.3f", b[k * ldb + j]);
    }
    printf("\n");
  }
  printf("\n");
  printf(" alpha: %9.3f\n", wmean[m]);
  printf("\n");
  printf("  Step     Sum      RSS       df       Cp       Ck     Step Size\n ");
  for (i = 0; i < 64; i++)
    printf("-");
  printf("\n");
  for (k = 0; k < nstep; k++) {
    printf("  %3" NAG_IFMT " %9.3f %9.3f %6.0f  %9.3f %9.3f %9.3f\n", k + 1,
           fitsum[k * 6], fitsum[k * 6 + 1], fitsum[k * 6 + 2],
           fitsum[k * 6 + 3], fitsum[k * 6 + 4], fitsum[k * 6 + 5]);
  }
  printf("\n");
  printf(" sigma^2: %9.3f\n", fitsum[nstep * 6 + 4]);

END:
  NAG_FREE(dy);
  NAG_FREE(wmean);
  NAG_FREE(dtd);
  NAG_FREE(b);
  NAG_FREE(fitsum);
  NAG_FREE(ropt);

  return (exit_status);
}