/* nag_reml_mixed_regsn (g02jac) Example Program.
 *
 * Copyright 2004 Numerical Algorithms Group.
 *
 * Mark 8, 2004.
 */

#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");
  lb = 25;
  /* Skip heading in data file */
  scanf("%*[^\n] ");

  /* Read in the problem size information */
  scanf("%ld%ld%ld%ld%ld%*[^\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("%ld", &levels[i - 1]);
    }
  scanf("%*[^\n] ");

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

  /* Read in the variance component flag */
  for (i = 1; i <= nrv; ++i)
    {
      scanf("%ld", &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("%ld%*[^\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%4ld Level%4ld%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%4ld%s%4ld%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%4ld%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%4ld%s%4ld%s%4ld"
                              "%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("%4ld%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%ld\n", "DF          = ", df);
    }
  else
    {
      printf("Routine nag_reml_mixed_regsn (g02jac) failed, with error "
        "message:\n%s\n", fail.message);
    }

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