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

void print_results(Nag_OrderType order, Integer n, Integer nff, Integer nlsv,
                   Integer nrf, Integer fixed[], Integer nrndm, Integer rndm[],
                   Integer lrndm, Integer nvpr, Integer vpr[], double gamma[],
                   Integer effn, Integer rnkx, Integer ncov, double lnlike,
                   Integer id[], Integer pdid, double b[], double se[]);

#define RNDM(I, J)                                                             \
  rndm[(order == Nag_ColMajor) ? ((J - 1) * lrndm + I - 1)                     \
                               : ((I - 1) * nrndm + J - 1)]
#define DAT(I, J)                                                              \
  dat[(order == Nag_ColMajor) ? ((J - 1) * pddat + I - 1)                      \
                              : ((I - 1) * pddat + J - 1)]
#define ID(I, J) id[((J - 1) * pdid + I - 1)]

int main(void) {
  /* IO file pointers */

  /* Integer scalar and array declarations */
  Integer exit_status = 0;
  Integer pdid, licomm, lrcomm, tdczz, lb, pdcxx, pdcxz, pdczz, pddat, effn, i,
      j, lvpr, n, ncol, ncov, lfixed, nff, nl, nlsv, nrndm, nrf, nv, nvpr, rnkx,
      lwt, size_dat, lrndm;
  Integer *fixed = 0, *icomm = 0, *id = 0, *levels = 0, *rndm = 0;
  Integer *vpr = 0;
  Integer ticomm[2];

  /* NAG structures */
  NagError fail;
  Nag_OrderType order = Nag_RowMajor;

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

  /* Character scalars */
  char weight;

  /* Use the default options */
  Integer *iopt = 0;
  Integer liopt = 0;
  double *ropt = 0;
  Integer lropt = 0;

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

  printf("nag_correg_mixeff_hier_ml (g02jec) 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 "%" NAG_IFMT "%*[^\n] ",
        &weight, &n, &ncol, &nrndm, &nvpr);

  /* Maximum size for fixed and rndm */
  lfixed = ncol + 2;
  lrndm = 2 * ncol + 3;

  if (order == Nag_ColMajor) {
    pddat = n;
    size_dat = pddat * ncol;
  } else {
    pddat = ncol;
    size_dat = pddat * n;
  }

  /* Allocate some memory */
  if (!(y = NAG_ALLOC(n, double)) || !(vpr = NAG_ALLOC(nvpr, Integer)) ||
      !(levels = NAG_ALLOC(ncol, Integer)) ||
      !(gamma = NAG_ALLOC(nvpr + 1, double)) ||
      !(fixed = NAG_ALLOC(lfixed, Integer)) ||
      !(rndm = NAG_ALLOC(lrndm * nrndm, Integer)) ||
      !(dat = NAG_ALLOC(size_dat, double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

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

  /* Read in the number of levels associated with each of the
     independent variables */
  for (i = 0; i < ncol; i++)
    scanf("%" NAG_IFMT "", &levels[i]);
  scanf("%*[^\n] ");

  /* Read in the fixed part of the model */
  /* Skip the heading */
  scanf("%*[^\n] ");
  /* Number of variables */
  scanf("%" NAG_IFMT "%*[^\n] ", &fixed[0]);
  nv = fixed[0];
  if (nv + 2 > lfixed) {
    printf(" ** Problem size too large, increase array sizes\n");
    printf("LFIXED,NV+2 = %" NAG_IFMT ", %" NAG_IFMT "\n", lfixed, nv + 2);
    exit_status = -1;
    goto END;
  }
  /* Intercept */
  scanf("%" NAG_IFMT "%*[^\n] ", &fixed[1]);
  /* Variable IDs */
  if (nv > 0) {
    for (i = 2; i < nv + 2; i++)
      scanf("%" NAG_IFMT "", &fixed[i]);
    scanf("%*[^\n] ");
  }

  /* Read in the random part of the model */
  lvpr = 0;
  pdid = 0;
  for (j = 1; j <= nrndm; j++) {
    scanf("%*[^\n] ");
    /* Number of variables */
    scanf("%" NAG_IFMT "%*[^\n] ", &RNDM(1, j));
    nv = RNDM(1, j);
    if ((nv + 3) > lrndm) {
      printf(" ** Problem size too large, increase array sizes\n");
      printf("LRNDM,NV+2 = %" NAG_IFMT ", %" NAG_IFMT "\n", lrndm, nv + 2);
      exit_status = -1;
      goto END;
    }
    /* Intercept */
    scanf("%" NAG_IFMT "%*[^\n] ", &RNDM(2, j));
    /* Variable IDs */
    if (nv > 0) {
      for (i = 3; i <= nv + 2; i++)
        scanf("%" NAG_IFMT "", &RNDM(i, j));
      scanf("%*[^\n] ");
    }
    /* Number of subject variables */
    scanf("%" NAG_IFMT "%*[^\n] ", &RNDM(nv + 3, j));
    nl = RNDM(nv + 3, j);
    if (nv + nl + 2 > lrndm) {
      printf(" ** Problem size too large, increase array sizes\n");
      printf("LRNDM,NV+NL++2 = %" NAG_IFMT ", %" NAG_IFMT "\n", lrndm,
             nv + nl + 2);
      exit_status = -1;
      goto END;
    }
    /* Subject variable IDs */
    if (nl > 0) {
      for (i = nv + 4; i <= nv + nl + 3; i++)
        scanf("%" NAG_IFMT "", &RNDM(i, j));
      scanf("%*[^\n] ");
    }
    pdid = MAX(pdid, nl);
    lvpr += RNDM(2, j) + nv;
  }
  pdid += 3;

  /* Read in the dependent and independent data */
  for (i = 1; i <= n; i++) {
    scanf("%lf", &y[i - 1]);
    for (j = 1; j <= ncol; j++)
      scanf("%lf", &DAT(i, j));
    if (lwt > 0)
      scanf("%lf", &wt[i - 1]);
    scanf("%*[^\n] ");
  }

  /* Read in VPR */
  for (i = 0; i < lvpr; i++)
    scanf("%" NAG_IFMT "", &vpr[i]);
  scanf("%*[^\n] ");

  /* Read in GAMMA */
  for (i = 0; i < nvpr; i++)
    scanf("%lf", &gamma[i]);
  scanf("%*[^\n] ");

  /* Get the size of the communication arrays */
  licomm = 2;
  lrcomm = 1;
  nag_correg_mixeff_hier_init(order, n, ncol, dat, pddat, levels, y, wt, fixed,
                              lfixed, nrndm, rndm, lrndm, &nff, &nlsv, &nrf,
                              trcomm, lrcomm, ticomm, licomm, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_correg_mixeff_hier_init (g02jcc).\n%s\n",
           fail.message);
    exit_status = 1;
    goto END;
  }
  licomm = ticomm[0];
  lrcomm = ticomm[1];

  /* Allocate the communication arrays */
  if (!(icomm = NAG_ALLOC(licomm, Integer)) ||
      !(rcomm = NAG_ALLOC(lrcomm, double))) {
    printf("Allocation failure 4\n");
    exit_status = -1;
    goto END;
  }

  /* Pre-process the data */
  nag_correg_mixeff_hier_init(order, n, ncol, dat, pddat, levels, y, wt, fixed,
                              lfixed, nrndm, rndm, lrndm, &nff, &nlsv, &nrf,
                              rcomm, lrcomm, icomm, licomm, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_correg_mixeff_hier_init (g02jcc).\n%s\n",
           fail.message);
    exit_status = 1;
    goto END;
  }

  /* Allocate the output arrays */
  lb = nff + nrf * nlsv;
  tdczz = nrf * nlsv;
  pdcxx = nff;
  pdcxz = nff;
  pdczz = nrf;
  if (!(b = NAG_ALLOC(lb, double)) || !(cxx = NAG_ALLOC(pdcxx * nff, double)) ||
      !(cxz = NAG_ALLOC(pdcxz * tdczz, double)) ||
      !(czz = NAG_ALLOC(pdczz * tdczz, double)) ||
      !(se = NAG_ALLOC(lb, double)) || !(id = NAG_ALLOC(pdid * lb, Integer))) {
    printf("Allocation failure 5\n");
    exit_status = -1;
    goto END;
  }

  /*  Perform the analysis */
  nag_correg_mixeff_hier_ml(lvpr, vpr, nvpr, gamma, &effn, &rnkx, &ncov,
                            &lnlike, lb, id, pdid, b, se, czz, pdczz, cxx,
                            pdcxx, cxz, pdcxz, rcomm, icomm, iopt, liopt, ropt,
                            lropt, &fail);
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_correg_mixeff_hier_ml (g02jec).\n%s\n",
           fail.message);
    exit_status = 1;
    if (fail.code != NW_NOT_CONVERGED && fail.code != NW_TOO_MANY_ITER &&
        fail.code != NW_KT_CONDITIONS && fail.code != NE_NEG_ELEMENT)
      goto END;
  }

  /* Display the output */
  print_results(order, n, nff, nlsv, nrf, fixed, nrndm, rndm, lrndm, nvpr, vpr,
                gamma, effn, rnkx, ncov, lnlike, id, pdid, b, se);

END:

  NAG_FREE(wt);
  NAG_FREE(y);
  NAG_FREE(vpr);
  NAG_FREE(levels);
  NAG_FREE(gamma);
  NAG_FREE(fixed);
  NAG_FREE(rndm);
  NAG_FREE(dat);
  NAG_FREE(icomm);
  NAG_FREE(rcomm);
  NAG_FREE(b);
  NAG_FREE(cxx);
  NAG_FREE(cxz);
  NAG_FREE(czz);
  NAG_FREE(se);
  NAG_FREE(id);

  return exit_status;
}

void print_results(Nag_OrderType order, Integer n, Integer nff, Integer nlsv,
                   Integer nrf, Integer fixed[], Integer nrndm, Integer rndm[],
                   Integer lrndm, Integer nvpr, Integer vpr[], double gamma[],
                   Integer effn, Integer rnkx, Integer ncov, double lnlike,
                   Integer id[], Integer pdid, double b[], double se[]) {
  Integer aid, i, k, l, ns, nv, p, pb, tb, tdid, vid, same, pk;

  /* Display the output */
  printf(" Number of observations (N)                    =  %" NAG_IFMT "\n",
         n);
  printf(" Number of random factors (NRF)                =  %" NAG_IFMT "\n",
         nrf);
  printf(" Number of fixed factors (NFF)                 =  %" NAG_IFMT "\n",
         nff);
  printf(" Number of subject levels (NLSV)               =  %" NAG_IFMT "\n",
         nlsv);
  printf(" Rank of X (RNKX)                              =  %" NAG_IFMT "\n",
         rnkx);
  printf(" Effective N (EFFN)                            =  %" NAG_IFMT "\n",
         effn);
  printf(" Number of nonzero variance components (NCOV) =  %" NAG_IFMT "\n",
         ncov);

  printf(" Parameter Estimates\n");
  tdid = nff + nrf * nlsv;

  if (nrf > 0) {
    printf("\n");
    printf(" Random Effects\n");
  }

  pb = -999;
  pk = -9;
  for (k = 1; k <= nrf * nlsv; k++) {
    tb = ID(1, k);
    if (tb != -999) {
      vid = ID(2, k);
      nv = RNDM(1, tb);
      ns = RNDM(3 + nv, tb);

      if (pb != tb) {
        same = 0;
      } else {
        same = 1;
        for (l = 1; l <= ns; l++) {
          if (ID(3 + l, k) != ID(3 + l, pk)) {
            same = 0;
            break;
          }
        }
      }

      if (!same) {
        if (k != 1)
          printf("\n");
        printf("  Subject: ");
        for (l = 1; l <= ns; l++)
          printf(" Variable %2" NAG_IFMT " (Level %1" NAG_IFMT ") ",
                 RNDM(3 + nv + l, tb), ID(3 + l, k));
        printf("\n");
      }
      pb = tb;
      pk = k;

      if (vid == 0) {
        /* Intercept */
        printf("    Intercept                 %10.4f %10.4f\n", b[k], se[k]);
      } else {
        /* VID'th variable specified in RNDM */
        aid = RNDM(2 + vid, tb);
        if (ID(3, k) == 0) {
          printf("     Variable %2" NAG_IFMT "", aid);
          printf("               %10.4f %10.4f\n", b[k - 1], se[k - 1]);
        } else {
          printf("     Variable %2" NAG_IFMT "", aid);
          printf(" (Level %1" NAG_IFMT ")     %10.4f %10.4f\n", ID(3, k),
                 b[k - 1], se[k - 1]);
        }
      }
    }
  }

  if (nff > 0) {
    printf("\n");
    printf(" Fixed Effects\n");
  }
  for (k = nrf * nlsv + 1; k <= tdid; k++) {
    vid = ID(2, k);
    if (vid != -999) {
      if (vid == 0) {
        /* Intercept */
        printf("   Intercept                   %10.4f %10.4f\n", b[k - 1],
               se[k - 1]);
      } else {
        /* VID'th variable specified in FIXED */
        aid = fixed[2 + vid - 1];
        if (ID(3, k) == 0) {
          printf("   Variable %2" NAG_IFMT "", aid);
          printf("                 %10.4f %10.4f\n", b[k - 1], se[k - 1]);
        } else {
          printf("   Variable %2" NAG_IFMT "", aid);
          printf(" (Level %1" NAG_IFMT ")       %10.4f %10.4f\n", ID(3, k),
                 b[k - 1], se[k - 1]);
        }
      }
    }
  }

  printf("\n");
  printf(" Variance Components\n");
  printf("  Estimate      Parameter        Subject\n");
  for (k = 1; k <= nvpr; k++) {
    printf("%10.5f     ", gamma[k - 1]);
    p = 0;
    for (tb = 1; tb <= nrndm; tb++) {
      nv = RNDM(1, tb);
      ns = RNDM(3 + nv, tb);
      for (i = 1; i <= nv + RNDM(2, tb); i++) {
        p++;
        if (vpr[p - 1] == k) {
          printf("Variable %2" NAG_IFMT "     Variables ", RNDM(2 + i, tb));
          for (l = 1; l <= ns; l++)
            printf("%2" NAG_IFMT " ", RNDM(3 + nv + l, tb));
        }
      }
    }
    printf("\n");
  }
  printf("\n");
  printf("SIGMA**2         = %15.5f\n", gamma[nvpr]);
  printf("-2LOG LIKELIHOOD = %15.5f\n", lnlike);
}