/* nag_correg_lmm_init_combine (g02jgc) Example Program.
 *
 * Copyright 2024 Numerical Algorithms Group.
 *
 * Mark 30.3, 2024.
 */
/* Pre-processor includes */
#include <math.h>
#include <nag.h>
#include <stdio.h>
#include <string.h>

char *read_line(char formula[], Integer nchar);

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

#define MAX_FORMULA_LEN 200
#define MAX_VNAME_LEN 200
#define MAX_PLAB_LEN 200
#define LCVALUE 2

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, ip,
      lvinfo, lisx, lenlab, sddat, lvnames, blicomm, blrcomm, totaln, blk, nblk;
  Integer *icomm = 0, *levels = 0, *isx = 0, *vinfo = 0, *bicomm = 0;

  /* Nag Types */
  NagError fail;
  Nag_IncludeIntercept intcpt;
  Nag_VariableType optype;

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

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

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

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

  printf("nag_correg_lmm_init_combine (g02jgc) Example Program Results\n\n");

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

  /* Read in the problem size */
  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 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);
  }

  /* Skip header */
  scanf("%*[^\n] ");

/* Read in the number of blocks of data */
  scanf("%" NAG_IFMT "%*[^\n] ", &nblk);

  /* Loop over each block of data */
  totaln = 0;
  for (blk = 0; blk < nblk; blk++) {
    /* Skip header */
    scanf("%*[^\n] ");

    /* Read in the number of observations in this block */
    scanf("%" NAG_IFMT "%*[^\n] ", &nobs);

    /* Keep a running total of the number of observations processed */
    totaln += nobs;

    /* Check whether we are supplying weights and allocate memory */
    if (weight == 'W') {
      if (!(wt = NAG_ALLOC(nobs, double))) {
        printf("Allocation failure\n");
        exit_status = -1;
        goto END;
      }
    }

    pddat = nobs;
    sddat = nvar;

    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] ");

    /*  Get the size of the communication arrays and allocate them */
    blicomm = 2;
    blrcomm = 0;
    if (!(bicomm = NAG_ALLOC(blicomm, Integer)) ||
        !(brcomm = NAG_ALLOC(blrcomm, 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, brcomm,
                        blrcomm, bicomm, blicomm, &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;
    }
    blicomm = bicomm[0];
    blrcomm = bicomm[1];
    NAG_FREE(bicomm);
    NAG_FREE(brcomm);
    if (!(bicomm = NAG_ALLOC(blicomm, Integer)) ||
        !(brcomm = NAG_ALLOC(blrcomm, 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, brcomm,
                        blrcomm, bicomm, blicomm, &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;
      }
    }

    if (blk == 0) {
      /* This is the first block of data, so move the communication arrays */
      /* into ICOMM and RCOMM */
      icomm = bicomm;
      bicomm = 0;
      licomm = blicomm;
      rcomm = brcomm;
      brcomm = 0;
      lrcomm = blrcomm;
    } else {
      /* Combine the current block with the previous blocks */
      /* using nag_correg_lmm_init_combine (g02jgc) */
      nag_correg_lmm_init_combine(hlmm, rcomm, lrcomm, icomm, licomm, brcomm,
                                  bicomm, &fail);
      if (fail.code == NW_ARRAY_SIZE) {
        /* icomm and / or rcomm are not sufficiently large to hold the */
        /* combined communication array */
        licomm = icomm[0];
        lrcomm = icomm[1];
        if (!(icomm = NAG_REALLOC(icomm, licomm, Integer)) ||
            !(rcomm = NAG_REALLOC(rcomm, lrcomm, double))) {
          printf("Allocation failure\n");
          exit_status = -1;
          goto END;
        }

        /* Re-run the routine with larger communication sizes */
        nag_correg_lmm_init_combine(hlmm, rcomm, lrcomm, icomm, licomm, brcomm,
                                    bicomm, &fail);
      }

      if (fail.code != NE_NOERROR) {
        printf("Error from nag_correg_lmm_init_combine (g02jgc).\n%s\n",
               fail.message);
        exit_status = 1;
        goto END;
      }
    }

    /* Free up the arrays for the current  block */
    NAG_FREE(wt);
    NAG_FREE(y);
    NAG_FREE(dat);
    NAG_FREE(bicomm);
    NAG_FREE(brcomm);
  }

  /* 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);
  }

  /* Get the number of fixed and random factors and the number of */
  /* variance components for the combined dataset */
  nag_correg_optget("NFF", &nff, &rvalue, cvalue, LCVALUE, &optype, icomm,
                    rcomm, &fail);
  if (fail.code == NE_NOERROR)
    nag_correg_optget("NRF", &nrf, &rvalue, cvalue, LCVALUE, &optype, icomm,
                      rcomm, &fail);
  if (fail.code == NE_NOERROR)
    nag_correg_optget("RNLSV", &rnlsv, &rvalue, cvalue, LCVALUE, &optype, icomm,
                      rcomm, &fail);
  if (fail.code == NE_NOERROR)
    nag_correg_optget("FNLSV", &fnlsv, &rvalue, cvalue, LCVALUE, &optype, icomm,
                      rcomm, &fail);
  if (fail.code == NE_NOERROR)
    nag_correg_optget("NVPR", &nvpr, &rvalue, cvalue, LCVALUE, &optype, icomm,
                      rcomm, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_correg_optget (g02zlc).\n%s\n", fail.message);
    exit_status = 1;
    goto END;
  }

  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)) || !(vplab = NAG_ALLOC(nvpr, char *))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }
  lenlab = MAX_PLAB_LEN;
  for (i = 0; i < nvpr; i++) {
    if (!(vplab[i] = NAG_ALLOC(lenlab, char))) {
      printf("Allocation failure\n");
      exit_status = -1;
      goto END;
    }
  }

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

  if (nvpr > 0) {
    /* Call nag_blgm_lm_submodel (g22ydc) to get the labels for the */
    /* variance components estimates */
    nag_blgm_lm_submodel(Nag_VarianceComponent, NULL, hlmm, &intcpt, &ip, lisx,
                         isx, nvpr, vplab, 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;
    }
    printf("\n");
    printf("Variance Components\n");
    printf("Component                        Estimate\n");
    for (i = 0; i < nvpr; i++) {
      printf("%-25s     %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 (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;
  }
}