/* nag_reml_mixed_regsn (g02jac) Example Program.
 *
 * Copyright 2017 Numerical Algorithms Group.
 *
 * Mark 26.1, 2017.
 */

#include <stdio.h>
#include <nag.h>
#include <nag_stdlib.h>
#include <nagg02.h>

int main(void)
{

  /* Scalars */
  double like, tol;
  Integer cwid, df, exit_status, fint, i, j, k, l, lb, maxit, n, ncol, nff;
  Integer nfv, nrf, nrv, nvpr, tddat, rint, svid, warnp, yvid, fnlevel;
  Integer rnlevel, lgamma, fl;
  /* Nag types */
  NagError fail;

  /* Arrays */
  double *b = 0, *dat = 0, *gamma = 0, *se = 0;
  Integer *fvid = 0, *levels = 0, *rvid = 0, *vpr = 0;

#define DAT(I, J) dat[(I-1)*tddat + J - 1]

  exit_status = 0;
  INIT_FAIL(fail);

  printf("nag_reml_mixed_regsn (g02jac) Example Program Results\n\n");

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

  /* Read in the problem size information */
  scanf("%" NAG_IFMT "%" NAG_IFMT "%" NAG_IFMT "%" NAG_IFMT "%" NAG_IFMT
        "%*[^\n] ", &n, &ncol, &nfv, &nrv, &nvpr);

  /* Check problem size */
  if (n < 0 || ncol < 0 || nfv < 0 || nrv < 0 || nvpr < 0) {
    printf("Invalid problem size, at least one of n, ncol, nfv, "
           "nrv or nvpr is < 0\n");
    exit_status = 1;
    goto END;
  }

  /* Allocate memory first lot of memory */
  if (!(levels = NAG_ALLOC(ncol, Integer)) ||
      !(fvid = NAG_ALLOC(nfv, Integer)) ||
      !(rvid = NAG_ALLOC(nrv, Integer)) || !(vpr = NAG_ALLOC(nrv, Integer)))
  {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Read in number of levels for each variable */
  for (i = 1; i <= ncol; ++i) {
    scanf("%" NAG_IFMT "", &levels[i - 1]);
  }
  scanf("%*[^\n] ");

  /* Read in model information */
  scanf("%" NAG_IFMT "", &yvid);
  for (i = 1; i <= nfv; ++i) {
    scanf("%" NAG_IFMT "", &fvid[i - 1]);
  }
  for (i = 1; i <= nrv; i++) {
    scanf("%" NAG_IFMT "", &rvid[i - 1]);
  }
  scanf("%" NAG_IFMT "%" NAG_IFMT "%" NAG_IFMT "%" NAG_IFMT "%*[^\n] ", &svid,
        &cwid, &fint, &rint);
  scanf("%*[^\n] ");

  /* Read in the variance component flag */
  for (i = 1; i <= nrv; ++i) {
    scanf("%" NAG_IFMT "", &vpr[i - 1]);
  }
  scanf("%*[^\n] ");

  /* If no subject specified, then ignore rint */
  if (svid == 0) {
    rint = 0;
  }

  /* Count the number of levels in the fixed parameters */
  for (i = 1, fnlevel = 0; i <= nfv; ++i) {
    fl = levels[fvid[i - 1] - 1] - 1;
    fnlevel += (fl < 1) ? 1 : fl;
  }
  if (fint == 1) {
    fnlevel++;
  }

  /* Count the number of levels in the random parameters */
  for (i = 1, rnlevel = 0; i <= nrv; ++i) {
    rnlevel += levels[rvid[i - 1] - 1];
  }
  if (rint) {
    rnlevel++;
  }

  /* Calculate the sizes of the output arrays */
  if (rint == 1) {
    lgamma = nvpr + 2;
  }
  else {
    lgamma = nvpr + 1;
  }
  if (svid) {
    lb = fnlevel + levels[svid - 1] * rnlevel;
  }
  else {
    lb = fnlevel + rnlevel;
  }

  tddat = ncol;

  /* Allocate remaining memory */
  if (!(dat = NAG_ALLOC(n * ncol, double)) ||
      !(gamma = NAG_ALLOC(lgamma, double)) ||
      !(b = NAG_ALLOC(lb, double)) || !(se = NAG_ALLOC(lb, double)))
  {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Read in the Data matrix */
  for (i = 1; i <= n; ++i) {
    for (j = 1; j <= ncol; ++j) {
      scanf("%lf", &DAT(i, j));
    }
  }

  /* Read in the initial values for GAMMA */
  for (i = 1; i < lgamma; ++i) {
    scanf("%lf", &gamma[i - 1]);
  }

  /* Read in the maximum number of iterations */
  scanf("%" NAG_IFMT "%*[^\n] ", &maxit);

  /* Run the analysis */
  tol = 0.;
  warnp = 0;
  /* nag_reml_mixed_regsn (g02jac).
   * Linear mixed effects regression using Restricted Maximum
   * Likelihood (REML)
   */
  nag_reml_mixed_regsn(n, ncol, dat, tddat, levels, yvid, cwid, nfv, fvid,
                       fint, nrv, rvid, nvpr, vpr, rint, svid, gamma, &nff,
                       &nrf, &df, &like, lb, b, se, maxit, tol, &warnp,
                       &fail);

  /* Report the results */
  if (fail.code == NE_NOERROR) {
    /* Output results */
    if (warnp != 0) {
      printf("Warning: At least one variance component was ");
      printf("estimated to be negative and then reset to zero");
      printf("\n");
    }
    printf("Fixed effects (Estimate and Standard Deviation)\n\n");
    k = 1;
    if (fint == 1) {
      printf("Intercept             %10.4f%10.4f\n", b[k - 1], se[k - 1]);
      ++k;
    }
    for (i = 1; i <= nfv; ++i) {
      for (j = 1; j <= levels[fvid[i - 1] - 1]; ++j) {
        if (levels[fvid[i - 1] - 1] != 1 && j == 1)
          continue;
        printf("Variable%4" NAG_IFMT " Level%4" NAG_IFMT "%10.4f%10.4f\n",
               i, j, b[k - 1], se[k - 1]);
        ++k;
      }
    }

    printf("\n");
    printf("Random Effects (Estimate and Standard Deviation)\n");
    if (svid == 0) {
      for (i = 1; i <= nrv; ++i) {
        for (j = 1; j <= levels[rvid[i - 1] - 1]; ++j) {
          printf("%s%4" NAG_IFMT "%s%4" NAG_IFMT "%10.4f%10.4f\n",
                 "Variable", i, " Level", j, b[k - 1], se[k - 1]);
          ++k;
        }
      }
    }
    else {
      for (l = 1; l <= levels[svid - 1]; ++l) {
        if (rint == 1) {
          printf("%s%4" NAG_IFMT "%s%10.4f%10.4f\n",
                 "Intercept for Subject Level", l, "         ",
                 b[k - 1], se[k - 1]);
          ++k;
        }
        for (i = 1; i <= nrv; ++i) {
          for (j = 1; j <= levels[rvid[i - 1] - 1]; ++j) {
            printf("%s%4" NAG_IFMT "%s%4" NAG_IFMT "%s%4" NAG_IFMT ""
                   "%10.4f%10.4f\n", "Subject Level", l,
                   " Variable", i, " Level", j, b[k - 1], se[k - 1]);
            ++k;
          }
        }
      }
    }

    printf("\n");
    printf("%s\n", " Variance Components");
    for (i = 1; i <= nvpr + rint; ++i) {
      printf("%4" NAG_IFMT "%10.4f\n", i, gamma[i - 1]);
    }
    printf("%s%10.4f\n\n", "SIGMA^2     = ", gamma[nvpr + rint]);

    printf("%s%10.4f\n\n", "-2LOG LIKE  = ", like);
    printf("%s%" NAG_IFMT "\n", "DF          = ", df);
  }
  else {
    printf("Routine nag_reml_mixed_regsn (g02jac) failed, with error "
           "message:\n%s\n", fail.message);
  }

END:
  NAG_FREE(b);
  NAG_FREE(dat);
  NAG_FREE(gamma);
  NAG_FREE(se);
  NAG_FREE(fvid);
  NAG_FREE(levels);
  NAG_FREE(rvid);
  NAG_FREE(vpr);
  return exit_status;
}