/* nag_correg_lmm_fit (g02jhc) Example Program.
 *
 * Copyright 2023 Numerical Algorithms Group.
 *
 * Mark 29.3, 2023.
 */
/* Pre-processor includes */
#include <math.h>
#include <nag.h>
#include <stdio.h>
#include <string.h>

char *read_line(char formula[], Integer nchar);
int custom_labels(Nag_DesignMatrixReturn what, void *hlmm, char *const vnames[],
                  char **plab[], Integer lenlab);

#define DAT(I, J) dat[(J)*pddat + (I)]

#define MAX_FORMULA_LEN 200
#define MAX_VNAME_LEN 200
#define MAX_PLAB_LEN 200

int main(void) {
  /* Integer scalar and array declarations */
  Integer exit_status = 0;
  Integer licomm, lrcomm, lb, pdcxx, pdcxz, pdczz, pddat, effn, i, j, nobs,
      nvar, ncov, nff, rnlsv, fnlsv, nrndm, nrf, nvpr, rnkx, nzz, nxx, lvinfo,
      lisx, lenlab, sddat, lvnames;
  Integer *icomm = 0, *levels = 0, *isx = 0, *vinfo = 0;

  /* Nag Types */
  NagError fail;

  /* Double scalar and array declarations */
  double lnlike;
  double *b = 0, *cxx = 0, *cxz = 0, *czz = 0, *dat = 0, *gamma = 0, *rcomm = 0,
         *se = 0, *wt = 0, *y = 0;

  /* Character scalar and array declarations */
  char weight;
  char formula[MAX_FORMULA_LEN];
  char **vnames = 0, **xplab = 0, **zplab = 0, **vplab = 0;

  /* Void pointers */
  void *hddesc = 0, *hfixed = 0, *hlmm = 0;
  void **hrndm = 0;

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

  printf("nag_correg_lmm_fit (g02jhc) Example Program Results\n\n");

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

  /* Read in the initial arguments */
  scanf("%c%" NAG_IFMT "%" NAG_IFMT "%" NAG_IFMT "%*[^\n] ", &weight, &nobs,
        &nvar, &lvnames);

  /* Read in number of levels and names for the variables */
  if (!(levels = NAG_ALLOC(nvar, Integer)) ||
      !(vnames = NAG_ALLOC(lvnames, char *))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }
  for (i = 0; i < lvnames; i++) {
    if (!(vnames[i] = NAG_ALLOC(MAX_VNAME_LEN, char))) {
      printf("Allocation failure\n");
      exit_status = -1;
      goto END;
    }
  }
  for (i = 0; i < nvar; i++)
    scanf("%" NAG_IFMT "", &levels[i]);
  scanf("%*[^\n] ");

  if (lvnames > 0) {
    for (i = 0; i < lvnames; i++)
      scanf("%50s", vnames[i]);
    scanf("%*[^\n] ");
  }

  /* Call nag_blgm_lm_describe_data (g22ybc) to get a description of */
  /* the data matrix */
  nag_blgm_lm_describe_data(&hddesc, nobs, nvar, levels, lvnames, vnames,
                            &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_blgm_lm_describe_data (g22ybc).\n%s\n",
           fail.message);
    exit_status = 1;
    goto END;
  }

  /* Read in the data matrix and response variable */
  pddat = nobs;
  sddat = nvar;
  if (weight == 'W' || weight == 'w') {
    if (!(wt = NAG_ALLOC(nobs, double))) {
      printf("Allocation failure\n");
      exit_status = -1;
      goto END;
    }
  }
  if (!(dat = NAG_ALLOC(pddat * sddat, double)) ||
      !(y = NAG_ALLOC(nobs, double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }
  for (i = 0; i < nobs; i++) {
    for (j = 0; j < nvar; j++) {
      scanf("%lf", &DAT(i, j));
    }
    scanf("%lf", &y[i]);
    if (wt) {
      scanf("%lf", &wt[i]);
    }
  }
  scanf("%*[^\n] ");

  /* Read in the formula for the fixed part of the model, remove comments */
  /* and call nag_blgm_lm_formula (g22yac) to parse it */
  read_line(formula, MAX_FORMULA_LEN);
  nag_blgm_lm_formula(&hfixed, formula, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_blgm_lm_formula (g22yac).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Read in the number of random statements */
  scanf("%" NAG_IFMT "%*[^\n] ", &nrndm);
  if (!(hrndm = NAG_ALLOC(nrndm, void *))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }
  for (i = 0; i < nrndm; i++)
    hrndm[i] = NULL;

  /* Read in the formula for the random parts of the model, remove comments */
  /* and call nag_blgm_lm_formula (g22yac) to parse them */
  for (i = 0; i < nrndm; i++) {
    read_line(formula, MAX_FORMULA_LEN);
    nag_blgm_lm_formula(&hrndm[i], formula, &fail);
  }

  /*  Get the size of the communication arrays and allocate them */
  licomm = 2;
  lrcomm = 0;
  if (!(icomm = NAG_ALLOC(licomm, Integer)) ||
      !(rcomm = NAG_ALLOC(lrcomm, double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }
  nag_correg_lmm_init(&hlmm, hddesc, hfixed, nrndm, hrndm, nobs, y, wt, dat,
                      pddat, sddat, &fnlsv, &nff, &rnlsv, &nrf, &nvpr, rcomm,
                      lrcomm, icomm, licomm, &fail);
  if (fail.code != NW_ARRAY_SIZE) {
    printf("Error from nag_correg_lmm_init (g02jfc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }
  licomm = icomm[0];
  lrcomm = icomm[1];
  NAG_FREE(icomm);
  NAG_FREE(rcomm);
  if (!(icomm = NAG_ALLOC(licomm, Integer)) ||
      !(rcomm = NAG_ALLOC(lrcomm, double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Pre-process the data using nag_correg_lmm_init (g02jfc) */
  nag_correg_lmm_init(&hlmm, hddesc, hfixed, nrndm, hrndm, nobs, y, wt, dat,
                      pddat, sddat, &fnlsv, &nff, &rnlsv, &nrf, &nvpr, rcomm,
                      lrcomm, icomm, licomm, &fail);
  if (fail.code != NE_NOERROR) {
    if (fail.code == NW_POTENTIAL_PROBLEM) {
      printf("Warning from nag_correg_lmm_init (g02jfc).\n%s\n", fail.message);
    } else {
      printf("Error from nag_correg_lmm_init (g02jfc).\n%s\n", fail.message);
      exit_status = 1;
      goto END;
    }
  }

  /* Use nag_blgm_handle_free (g22zac) to clean-up the handles that */
  /* are no longer required */
  nag_blgm_handle_free(&hddesc, &fail);
  nag_blgm_handle_free(&hfixed, &fail);
  if (hrndm) {
    for (i = 0; i < nrndm; i++)
      nag_blgm_handle_free(&hrndm[i], &fail);
    NAG_FREE(hrndm);
  }

  nzz = nrf * rnlsv;
  nxx = nff * fnlsv;
  lb = nxx + nzz;

  /* We don't want the output from CXX, CZZ and CXZ */
  pdczz = 0;
  pdcxx = 0;
  pdcxz = 0;

  /* Allocate the rest of the output arrays */
  if (!(gamma = NAG_ALLOC(nvpr + 1, double)) || !(b = NAG_ALLOC(lb, double)) ||
      !(se = NAG_ALLOC(lb, double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Use MIVQUE for the initial values for gamma */
  for (i = 0; i <= nvpr; i++)
    gamma[i] = -1.0;

  /* Fit the model using nag_correg_lmm_fit (g02jhc) */
  nag_correg_lmm_fit(hlmm, nvpr, gamma, &effn, &rnkx, &ncov, &lnlike, lb, b, se,
                     czz, pdczz, cxx, pdcxx, cxz, pdcxz, rcomm, icomm, &fail);
  if (fail.code != NE_NOERROR) {
    if (fail.code == NW_NOT_CONVERGED || fail.code == NW_KT_CONDITIONS ||
        fail.code == NE_NEG_ELEMENT) {
      printf("Warning from nag_correg_lmm_fit (g02jhc).\n%s\n", fail.message);
    } else {
      printf("Error from nag_correg_lmm_fit (g02jhc).\n%s\n", fail.message);
      exit_status = 1;
      goto END;
    }
  }

  lvinfo = 0;
  lisx = 0;
  if (!(vinfo = NAG_ALLOC(lvinfo, Integer)) ||
      !(isx = NAG_ALLOC(lisx, Integer))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }
  lenlab = MAX_PLAB_LEN;

  /* Print the results */
  printf("Number of observations           = %" NAG_IFMT "\n", nobs);
  printf("Total number of random factors   = %" NAG_IFMT "\n", nzz);
  printf("Total number of fixed factors    = %" NAG_IFMT "\n", nxx);
  printf("Rank of X                        = %" NAG_IFMT "\n", rnkx);
  printf("Effective N                      = %" NAG_IFMT "\n", effn);
  printf("-2Log Likelihood                 = %-.4f\n", lnlike);
  printf("Sigma**2                         = %-.4f\n", gamma[nvpr]);

  printf("Parameter Estimates\n");
  if (nzz > 0) {
    /* Get the labels for the random parameter estimates */
    if (custom_labels(Nag_Z, hlmm, vnames, &zplab, lenlab))
      goto END;
    printf("\n");
    printf("Random Effects\n");
    printf("Parameter");
    for (i = 0; i < 42; i++)
      printf(" ");
    printf("Estimate     Standard Error\n");
    for (i = 0; i < nzz; i++) {
      if (zplab[i][1] != '\0')
        printf("%-47s %10.4f     %10.4f\n", zplab[i], b[i], se[i]);
    }
  }

  if (nxx > 0) {
    /* Get the labels for the fixed parameter estimates */
    if (custom_labels(Nag_X, hlmm, vnames, &xplab, lenlab))
      goto END;
    printf("\n");
    printf("Fixed Effects\n");
    printf("Parameter");
    for (i = 0; i < 42; i++)
      printf(" ");
    printf("Estimate     Standard Error\n");
    for (i = 0; i < nxx; i++) {
      if (xplab[i][1] != '\0')
        printf("%-47s %10.4f     %10.4f\n", xplab[i], b[nzz + i], se[nzz + i]);
    }
  }

  if (nvpr > 0) {
    /* Get the labels for the variance components estimates */
    if (custom_labels(Nag_VarianceComponent, hlmm, vnames, &vplab, lenlab))
      goto END;
    printf("\n");
    printf("Variance Components\n");
    printf("Component");
    for (i = 0; i < 42; i++)
      printf(" ");
    printf("Estimate\n");
    for (i = 0; i < nvpr; i++) {
      printf("%-47s %10.4f\n", vplab[i], gamma[i]);
    }
  }

END:
  /* Use nag_blgm_handle_free (g22zac) to clean-up the remaining G22 handles */
  nag_blgm_handle_free(&hlmm, &fail);

  NAG_FREE(levels);
  if (vnames) {
    for (i = 0; i < lvnames; i++)
      NAG_FREE(vnames[i]);
    NAG_FREE(vnames);
  }
  NAG_FREE(dat);
  NAG_FREE(y);
  NAG_FREE(wt);
  NAG_FREE(icomm);
  NAG_FREE(rcomm);
  NAG_FREE(gamma);
  NAG_FREE(b);
  NAG_FREE(se);
  NAG_FREE(vinfo);
  NAG_FREE(isx);
  if (xplab) {
    for (i = 0; i < nxx; i++)
      NAG_FREE(xplab[i]);
    NAG_FREE(xplab);
  }
  if (zplab) {
    for (i = 0; i < nzz; i++)
      NAG_FREE(zplab[i]);
    NAG_FREE(zplab);
  }
  if (vplab) {
    for (i = 0; i < nvpr; i++)
      NAG_FREE(vplab[i]);
    NAG_FREE(vplab);
  }
  nag_blgm_handle_free(&hddesc, &fail);
  nag_blgm_handle_free(&hfixed, &fail);
  if (hrndm) {
    for (i = 0; i < nrndm; i++)
      nag_blgm_handle_free(&hrndm[i], &fail);
    NAG_FREE(hrndm);
  }

  return exit_status;
}

char *read_line(char formula[], Integer nchar) {
  /* Read in a line from stdin and remove any comments */
  char *pch;

  /* Read in the model formula */
  if (fgets(formula, nchar, stdin)) {
    /* Strip comments from formula */
    pch = strstr(formula, "::");
    if (pch)
      *pch = '\0';
    return formula;
  } else {
    return 0;
  }
}

int custom_labels(Nag_DesignMatrixReturn what, void *hlmm, char *const vnames[],
                  char **plab[], Integer lenlab) {
  /* Get custom labels for the fixed or random parameter estimates or */
  /* estimates of the variance components as obtained from fitting a linear */
  /* mixed effects regression model using nag_correg_lmm_fit (g02jhc) */

  /* Integer scalar and array declarations */
  Integer exit_status = 0;
  Integer i, ip, j, k, li, lisx, lplab, lvinfo, nv, vi;
  Integer *isx = 0, *vinfo = 0;

  /* Character scalar and array declarations */
  char term[MAX_PLAB_LEN];
  char *this_lab = 0;
  Nag_IncludeIntercept intcpt;

  /* NAG types */
  NagError fail;

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

  /* Get the require array sizes */
  lisx = 0;
  lvinfo = 3;
  lplab = 0;
  if (!(vinfo = NAG_ALLOC(lvinfo, Integer))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }
  nag_blgm_lm_submodel(what, NULL, hlmm, &intcpt, &ip, lisx, isx, lplab, *plab,
                       lenlab, lvinfo, vinfo, &fail);
  if (fail.code != NE_NOERROR && fail.code != NW_ARRAY_SIZE) {
    printf("Error from nag_blgm_lm_submodel (g22ydc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  lplab = vinfo[1];
  lvinfo = vinfo[2];

  /* Reallocate the output arrays */
  NAG_FREE(vinfo);
  if (!(vinfo = NAG_ALLOC(lvinfo, Integer)) ||
      !((*plab) = NAG_ALLOC(lplab, char *))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }
  for (i = 0; i < lplab; i++) {
    if (!((*plab)[i] = NAG_ALLOC(lenlab, char))) {
      printf("Allocation failure\n");
      exit_status = -1;
      goto END;
    }
  }

  /* Generate the labelling information using nag_blgm_lm_submodel (g22ydc) */
  nag_blgm_lm_submodel(what, NULL, hlmm, &intcpt, &ip, lisx, isx, lplab, *plab,
                       lenlab, lvinfo, vinfo, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_blgm_lm_submodel (g22ydc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  /* Because we call nag_blgm_lm_submodel with lplab > 0, plab will contain */
  /* the default labels created by nag_blgm_lm_submodel, so we could return */
  /* at this point, however we will use the information in vinfo to create */
  /* our own custom labels */
  k = 0;
  for (j = 0; j < ip; j++) {
    nv = vinfo[k];
    k += 1;
    this_lab = (*plab)[j];
    if (nv == 0) {
      sprintf(this_lab, "Intercept");
    } else {
      for (i = 0; i < nv; i++) {
        vi = vinfo[k] - 1;
        li = vinfo[k + 1];

        if (li != 0)
          sprintf(term, "%s (Lvl %" NAG_IFMT ")%c", vnames[vi], li, '\0');
        else
          strncpy(term, vnames[vi], MAX_PLAB_LEN);

        if (i == 0) {
          strncpy(this_lab, term, MAX_PLAB_LEN);
          this_lab += strlen(this_lab);
        } else {
          this_lab += sprintf(this_lab, ".%s", term);
        }

        /* We are ignoring the contrast identifier when */
        /* constructing these labels */
        k += 3;
      }
    }
  }

END:

  NAG_FREE(vinfo);

  return (exit_status);
}