/* nag_kalman_sqrt_filt_cov_var (g13eac) Example Program.
 *
 * Copyright 1996 Numerical Algorithms Group
 *
 * Mark 4, 1996.
 * Mark 5 revised, 1998.
 * Mark 6 revised, 2000.
 * Mark 7 revised, 2001.
 * Mark 8 revised, 2004.
 *
 */

#include <nag.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <nag_stdlib.h>
#include <nage04.h>
#include <nagf06.h>
#include <nagg05.h>
#include <nagg13.h>
#include <nagx02.h>
#include <nagx04.h>

#ifdef __cplusplus
extern "C" {
#endif
static void NAG_CALL objfun(Integer n, double theta_phi[], double *objf,
                            double g[], Nag_Comm *comm);
#ifdef __cplusplus
}
#endif

static int ex1(FILE *fpin, FILE *fpout);
static int ex2(FILE *fpout, char *outfile);

int main(int argc, char *argv[])
{
  FILE     *fpin, *fpout;
  /* Integer scalar and array declarations */
  Integer  exit_status_ex1 = 0;
  Integer  exit_status_ex2 = 0;

  /* Character string declarations */
  char     *outfile = 0;

  /* NAG structures */
  NagError fail;

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

  /* Check for command-line IO options */
  fpin = nag_example_file_io(argc, argv, "-data", NULL);
  fpout = nag_example_file_io(argc, argv, "-results", NULL);
  (void) nag_example_file_io(argc, argv, "-nag_write", &outfile);
  if (!outfile)
    {
      outfile = NAG_ALLOC(7, char);
      strcpy(outfile, "stdout");
    }

  fprintf(fpout,
          "nag_kalman_sqrt_filt_cov_var (g13eac) Example Program Results\n\n");

  /* Run example 1 */
  exit_status_ex1 = ex1(fpin, fpout);

  /* Run example 2 */
  exit_status_ex2 = ex2(fpout, outfile);

  if (outfile) NAG_FREE(outfile);
  if (fpin != stdin) fclose(fpin);
  if (fpout != stdout) fclose(fpout);
  return (exit_status_ex1 == 0 && exit_status_ex2 == 0) ? 0 : 1;
}


/* Start of the first example ... */
#define A(I, J) a[(I) *tda + J]
#define B(I, J) b[(I) *tdb + J]
#define C(I, J) c[(I) *tdc + J]
#define K(I, J) k[(I) *tdk + J]
#define Q(I, J) q[(I) *tdq + J]
#define R(I, J) r[(I) *tdr + J]
#define S(I, J) s[(I) *tds + J]
#define H(I, J) h[(I) *tdh + J]

static int ex1(FILE *fpin, FILE *fpout)
{
  /* Integer scalar and array declarations */
  Integer  exit_status = 0;
  Integer  i, j, m, n, p, istep, tda, tdb, tdc, tdk, tdq, tdr, tds, tdh;

  /* Double scalar and array declarations */
  double   tol;
  double   *a = 0, *b = 0, *c = 0, *k = 0, *q = 0, *r = 0, *s = 0, *h = 0;

  /* NAG structures */
  NagError fail;

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

  fprintf(fpout, "Example 1\n");

  /* Skip the heading in the data file */
  fscanf(fpin, "%*[^\n]");

  /* Get the problem size */
  fscanf(fpin, "%ld%ld%ld%lf", &n, &m, &p, &tol);
  if (n < 1 || m < 1 || p < 1)
    {
      fprintf(fpout, "Invalid n or m or p.\n");
      exit_status = 1;
      goto END;
    }

  tda = n;
  tdb = m;
  tdc = n;
  tdk = p;
  tdq = m;
  tdr = p;
  tds = n;
  tdh = p;

  /* Allocate arrays */
  if (!(a = NAG_ALLOC(n*tda, double)) ||
      !(b = NAG_ALLOC(n*tdb, double)) ||
      !(c = NAG_ALLOC(p*tdc, double)) ||
      !(k = NAG_ALLOC(n*tdk, double)) ||
      !(q = NAG_ALLOC(m*tdq, double)) ||
      !(r = NAG_ALLOC(p*tdr, double)) ||
      !(s = NAG_ALLOC(n*tds, double)) ||
      !(h = NAG_ALLOC(n*tdh, double)))
    {
      fprintf(fpout, "Allocation failure\n");
      exit_status = -1;
      goto END;
    }

  /* Read data */
  for (i = 0; i < n; ++i)
    for (j = 0; j < n; ++j)
      fscanf(fpin, "%lf", &S(i, j));
  for (i = 0; i < n; ++i)
    for (j = 0; j < n; ++j)
      fscanf(fpin, "%lf", &A(i, j));
  for (i = 0; i < n; ++i)
    for (j = 0; j < m; ++j)
      fscanf(fpin, "%lf", &B(i, j));
  if (q)
    {
      for (i = 0; i < m; ++i)
        for (j = 0; j < m; ++j)
          fscanf(fpin, "%lf", &Q(i, j));
    }
  for (i = 0; i < p; ++i)
    for (j = 0; j < n; ++j)
      fscanf(fpin, "%lf", &C(i, j));
  for (i = 0; i < p; ++i)
    for (j = 0; j < p; ++j)
      fscanf(fpin, "%lf", &R(i, j));

  /* Perform three iterations of the Kalman filter recursion */
  for (istep = 1; istep <= 3; ++istep)
    {
      /* Run the kalman filter */
      nag_kalman_sqrt_filt_cov_var(n, m, p, s, tds, a, tda, b, tdb, q, tdq,
                                   c, tdc, r, tdr, k, tdk, h, tdh, tol, &fail);
      if (fail.code != NE_NOERROR)
        {
          fprintf(fpout,
                  "Error from nag_kalman_sqrt_filt_cov_var (g13eac).\n%s\n",
                  fail.message);
          exit_status = 1;
          goto END;
        }
    }

  /* Print the results */
  fprintf(fpout, "\nThe square root of the state covariance matrix is\n\n");
  for (i = 0; i < n; ++i)
    {
      for (j = 0; j < n; ++j)
        fprintf(fpout, "%8.4f ", S(i, j));
      fprintf(fpout, "\n");
    }
  if (k)
    {
      fprintf(fpout, "\nThe matrix AK (the product of the Kalman gain\n");
      fprintf(fpout, "matrix with the state transition matrix) is\n\n");
      for (i = 0; i < n; ++i)
        {
          for (j = 0; j < p; ++j)
            fprintf(fpout, "%8.4f ", K(i, j));
          fprintf(fpout, "\n");
        }
    }

 END:
  if (a) NAG_FREE(a);
  if (b) NAG_FREE(b);
  if (c) NAG_FREE(c);
  if (k) NAG_FREE(k);
  if (q) NAG_FREE(q);
  if (r) NAG_FREE(r);
  if (s) NAG_FREE(s);
  if (h) NAG_FREE(h);

  return exit_status;
}
/* ... end of the first example */


/* Start of the second example ... */
#define NY 2000

/* This illustrates the use of the kalman filter to estimate the
   parameters of an ARMA(1,1) time series model.
   Note : theta_phi[0] contains theta (moving average coefficient), and
   theta_phi[1] contains phi (autoregressive coefficient)
 */
static int ex2(FILE *fpout, char *outfile)
{
  /* Integer scalar and array declarations */
  Integer       exit_status = 0;
  Integer       n, lr;
  Integer       lstate;
  Integer       *state = 0;

  /* Double scalar and array declarations */
  double        sy[NY];
  double        *theta_phi = 0, *g = 0, *bl = 0, *bu = 0, *r = 0;
  double        objf, var;

  /* NAG structures and data types */
  Nag_BoundType bound;
  Nag_Comm      comm;
  Nag_E04_Opt   options;
  NagError      fail;
  Nag_ModeRNG   mode;

  /* Choose the base generator */
  Nag_BaseRNG   genid = Nag_Basic;
  Integer       subid = 0;

  /* Set the seed */
  Integer       seed[] = { 1762543 };
  Integer       lseed = 1;

  /* Set the autoregressive coefficients for the (randomly generated)
     time series */
  Integer       ip = 1;
  double        sphi[] = { 0.4 };

  /* Set the moving average coefficients for the (randomly generated)
     time series */
  Integer       iq = 1;
  double        stheta[] = { 0.9 };

  /* Number of terms in the (randomly generated) time series */
  Integer       nterms = NY;

  /* Mean and variance of the (randomly generated) time series */
  double        mean = 0.0, vara = 1.0;

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

  fprintf(fpout, "\n\nExample 2\n\n");

  /* Get the length of the state array */
  lstate = -1;
  nag_rand_init_repeatable(genid, subid, seed, lseed, state, &lstate, &fail);
  if (fail.code != NE_NOERROR)
    {
      fprintf(fpout, "Error from nag_rand_init_repeatable (g05kfc).\n%s\n",
              fail.message);
      exit_status = 1;
      goto END;
    }

  /* Calculate the size of the reference vector */
  lr = (ip > iq + 1)?ip:iq + 1;
  lr += ip+iq+6;

  /* Allocate arrays */
  if (!(state = NAG_ALLOC(lstate, Integer)) ||
      !(r = NAG_ALLOC(lr, double)))
    {
      fprintf(fpout, "Allocation failure\n");
      exit_status = -1;
      goto END;
    }

  /* Initialise the generator to a repeatable sequence */
  nag_rand_init_repeatable(genid, subid, seed, lseed, state, &lstate, &fail);
  if (fail.code != NE_NOERROR)
    {
      fprintf(fpout, "Error from nag_rand_init_repeatable (g05kfc).\n%s\n",
              fail.message);
      exit_status = 1;
      goto END;
    }

  /* Generate a time series with nterms terms */
  mode = Nag_InitializeAndGenerate;
  nag_rand_arma(mode, nterms, mean, ip, sphi, iq, stheta, vara, r, lr, state,
                &var, sy, &fail);
  if (fail.code != NE_NOERROR)
    {
      fprintf(fpout, "Error from nag_rand_arma (g05phc).\n%s\n", fail.message);
      exit_status = 1;
      goto END;
    }

  /* Number of parameters to estimate */
  n = ip + iq;

  /* Allocate arrays */
  if (!(theta_phi = NAG_ALLOC(n, double)) ||
      !(g = NAG_ALLOC(n, double)) ||
      !(bl = NAG_ALLOC(n, double)) ||
      !(bu = NAG_ALLOC(n, double)))
    {
      fprintf(fpout, "Allocation failure\n");
      exit_status = -1;
      goto END;
    }

  /* Make an initial guess of the parameters */
  theta_phi[0] = 0.5;
  theta_phi[1] = 0.5;

  /* Set the bounds */
  bound = Nag_Bounds;
  bl[0] = -1.0;
  bu[0] = 1.0;
  bl[1] = -1.0;
  bu[1] = 1.0;
  comm.user = &sy[0];

  /* nag_opt_init (e04xxc).
   * Initialization function for option setting
   */
  nag_opt_init(&options);
  strcpy(options.outfile, outfile);
  options.print_level = Nag_NoPrint;
  options.list = Nag_FALSE;

  /* nag_opt_bounds_no_deriv (e04jbc).
   * Bound constrained nonlinear minimization (no derivatives
   * required)
   */
  if (strcmp(outfile, "stdout"))
    fclose(fpout);
  nag_opt_bounds_no_deriv(n, objfun, bound, bl, bu, theta_phi, &objf, g,
                          &options, &comm, &fail);
  if (strcmp(outfile, "stdout"))
    {
      fpout = fopen(outfile, "a");
    }
  if (fail.code != NE_NOERROR && fail.code != NW_COND_MIN)
    {
      fprintf(fpout, "Error from nag_opt_bounds_no_deriv (e04jbc).\n%s\n",
              fail.message);
      exit_status = 1;
      goto END;
    }

  /* nag_opt_free (e04xzc).
   * Memory freeing function for use with option setting
   */
  nag_opt_free(&options, "all", &fail);
  if (fail.code != NE_NOERROR)
    {
      fprintf(fpout, "Error from nag_opt_free (e04xzc).\n%s\n", fail.message);
      exit_status = 1;
      goto END;
    }

  /* Display the results */
  fprintf(fpout, "The estimates are :  theta = %7.3f, phi = %7.3f \n",
          theta_phi[0], theta_phi[1]);

 END:
  if (state) NAG_FREE(state);
  if (theta_phi) NAG_FREE(theta_phi);
  if (g) NAG_FREE(g);
  if (bl) NAG_FREE(bl);
  if (bu) NAG_FREE(bu);
  if (r) NAG_FREE(r);

  return exit_status;
}

/* Define  objective function used in the non-linear optimisation routine ... */

static void NAG_CALL objfun(Integer n, double theta_phi[], double *objf,
                            double g[], Nag_Comm *comm)
{
  /* Routine to evaluate objective function. */

  Integer ione = 1, itwo = 2, k, m1 = 1, n1 = 2, nsteps = NY, nsum = 0, p1 = 1;
  double  a[2][2], ak[2][1], b[2][1], c[1][2], h[1][1], hs, k11, logdet = 0.0;
  double  phi, q[1][1], r[1][1];        /* There is no measurement noise */
  double  s[2][2], ss = 0.0, temp1, temp2, theta, tol = 0.0;
  double  v, xp[2], *y;

  y = comm->user;

  xp[0] = 0.0;   /* The expectation of the mean of an
                  * ARMA(1,1) is 0.0 */
  xp[1] = 0.0;
  q[0][0] = 1.0;

  c[0][0] = 1.0;
  c[0][1] = 0.0;

  r[0][0] = 0.0; /* There is no measurement noise */
  theta = theta_phi[0];
  phi = theta_phi[1];

  b[0][0] = 1.0;
  b[1][0] = -theta;

  a[0][0] = phi;
  a[1][0] = 0.0;
  a[0][1] = 1.0;
  a[1][1] = 0.0;

  /* set value for cholesky factor of state covariance matrix  */
  temp1 = 1.0 + (theta * theta) - (2.0 * theta * phi);
  temp2 = 1.0 - (phi * phi);
  k11 = temp1/temp2;
  s[0][0] = sqrt(k11);
  s[0][1] = 0.0;
  s[1][0] = -theta /s[0][0];
  s[1][1] = theta * sqrt(1.0 - (1.0/k11));


  /* iterate kalman filter for number of observations  */
  for (k = 1; k <= nsteps; ++k)
    {
      /* nag_kalman_sqrt_filt_cov_var (g13eac).
       * One iteration step of the time-varying Kalman filter
       * recursion using the square root covariance implementation
       */
      nag_kalman_sqrt_filt_cov_var(n1, m1, p1, &s[0][0], itwo, &a[0][0],
                                   itwo, &b[0][0], ione, &q[0][0], ione,
                                   &c[0][0], itwo, &r[0][0], ione, &ak[0][0],
                                   ione, &h[0][0], ione, tol, NAGERR_DEFAULT);

      v = y[k-1] - c[0][0]*xp[0];
      hs = h[0][0] * h[0][0];
      logdet = logdet +  log(hs);
      ss = ss + (v * v/ hs);
      nsum = nsum + 1;

      xp[0] = a[0][0]* xp[0] +  a[0][1] * xp[1] + ak[0][0] * v;
      xp[1] = ak[1][0] * v;
    }
  *objf = nsum * log(ss/nsum) + logdet;
}
/* ... end of the objective function definition */
/* ... end of the second example */