/* nag_tsa_kalman_unscented_state (g13ekc) Example Program.
 *
 * Copyright 2024 Numerical Algorithms Group.
 *
 * Mark 30.1, 2024.
 */

#include <math.h>
#include <nag.h>
#include <stdio.h>

#define LY(I, J) ly[(J)*my + (I)]
#define LX(I, J) lx[(J)*mx + (I)]
#define ST(I, J) st[(J)*mx + (I)]
#define XT(I, J) xt[(J)*mx + (I)]
#define FXT(I, J) fxt[(J)*mx + (I)]
#define YT(I, J) yt[(J)*mx + (I)]
#define HYT(I, J) hyt[(J)*my + (I)]

typedef struct g13_problem_data {
  double delta, a, r, d;
  double phi_rt, phi_lt;
} g13_problem_data;

#ifdef __cplusplus
extern "C" {
#endif
static void NAG_CALL f(Integer mx, Integer n, const double *xt, double *fxt,
                       Nag_Comm *comm, Integer *info);
static void NAG_CALL h(Integer mx, Integer my, Integer n, const double *yt,
                       double *hyt, Nag_Comm *comm, Integer *info);
#ifdef __cplusplus
}
#endif

static void read_problem_dat(Integer t, Nag_Comm *comm);

int main(void) {
  /* Integer scalar and array declarations */
  Integer mx = 3, my = 2;
  Integer i, ntime, t, j;
  Integer exit_status = 0;

  /* NAG structures and types */
  NagError fail;
  Nag_Comm comm;

  /* Double scalar and array declarations */
  double *lx = 0, *ly = 0, *st = 0, *x = 0, *y = 0;

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

  printf("nag_tsa_kalman_unscented_state (g13ekc) Example Program Results\n\n");

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

  if (!(lx = NAG_ALLOC(mx * mx, double)) ||
      !(ly = NAG_ALLOC(my * my, double)) || !(x = NAG_ALLOC(mx, double)) ||
      !(st = NAG_ALLOC(mx * mx, double)) || !(y = NAG_ALLOC(my, double))) {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END2;
  }

  /* Read in the Cholesky factorization of the covariance matrix for the
     process noise */
  for (i = 0; i < mx; i++) {
    for (j = 0; j <= i; j++) {
      scanf("%lf", &LX(i, j));
    }
    scanf("%*[^\n] ");
  }

  /* Read in the Cholesky factorization of the covariance matrix for the
     observation noise */
  for (i = 0; i < my; i++) {
    for (j = 0; j <= i; j++) {
      scanf("%lf", &LY(i, j));
    }
    scanf("%*[^\n] ");
  }

  /* Read in the initial state vector */
  for (i = 0; i < mx; i++) {
    scanf("%lf", &x[i]);
  }
  scanf("%*[^\n] ");

  /* Read in the Cholesky factorization of the initial state covariance
     matrix */
  for (i = 0; i < mx; i++) {
    for (j = 0; j <= i; j++) {
      scanf("%lf", &ST(i, j));
    }
    scanf("%*[^\n] ");
  }

  /* Read in the number of time points to run the system for */
  scanf("%" NAG_IFMT "%*[^\n] ", &ntime);

  /* Read in any problem specific data that is constant */
  read_problem_dat(0, &comm);

  /* Title for first set of output */
  printf("   Time  ");
  for (i = 0; i < (11 * mx - 16) / 2; i++)
    putchar(' ');
  printf("Estimate of State\n ");
  for (i = 0; i < 7 + 11 * mx; i++)
    putchar('-');
  printf("\n");

  /* Loop over each time point */
  for (t = 0; t < ntime; t++) {

    /* Read in any problem specific data that is time dependent */
    read_problem_dat(t + 1, &comm);

    /* Read in the observed data for time t */
    for (i = 0; i < my; i++) {
      scanf("%lf", &y[i]);
    }
    scanf("%*[^\n] ");

    /* Call Unscented Kalman Filter routine (g13ekc) */
    nag_tsa_kalman_unscented_state(mx, my, y, lx, ly, f, h, x, st, &comm,
                                   &fail);
    if (fail.code != NE_NOERROR) {
      printf("Error from nag_tsa_kalman_unscented_state (g13ekc).\n%s\n",
             fail.message);
      exit_status = 1;
      goto END;
    }

    /* Display the current state estimates */
    printf(" %3" NAG_IFMT "    ", t + 1);
    for (i = 0; i < mx; i++) {
      printf(" %10.3f", x[i]);
    }
    printf("\n");
  }

  printf("\n");
  printf("Estimate of Cholesky Factorization of the State\n");
  printf("Covariance Matrix at the Last Time Point\n");
  for (i = 0; i < mx; i++) {
    for (j = 0; j <= i; j++) {
      printf(" %10.2e", ST(i, j));
    }
    printf("\n");
  }

END:

  /* clean up any memory allocated in comm */
  read_problem_dat(-1, &comm);

END2:
  NAG_FREE(lx);
  NAG_FREE(ly);
  NAG_FREE(st);
  NAG_FREE(x);
  NAG_FREE(y);

  return (exit_status);
}

static void NAG_CALL f(Integer mx, Integer n, const double *xt, double *fxt,
                       Nag_Comm *comm, Integer *info) {
  double t1, t3;
  Integer i;
  g13_problem_data *pdat;

  /* Cast the void point in comm back to point at the data structure */
  pdat = (g13_problem_data *)comm->p;

  t1 = 0.5 * pdat->r * (pdat->phi_rt + pdat->phi_lt);
  t3 = (pdat->r / pdat->d) * (pdat->phi_rt - pdat->phi_lt);

  for (i = 0; i < n; i++) {
    FXT(0, i) = XT(0, i) + cos(XT(2, i)) * t1;
    FXT(1, i) = XT(1, i) + sin(XT(2, i)) * t1;
    FXT(2, i) = XT(2, i) + t3;
  }
  /* Set info nonzero to terminate execution for any reason. */
  *info = 0;
}

static void NAG_CALL h(Integer mx, Integer my, Integer n, const double *yt,
                       double *hyt, Nag_Comm *comm, Integer *info) {
  Integer i;
  g13_problem_data *pdat;

  /* Cast the void point in comm back to point at the data structure */
  pdat = (g13_problem_data *)comm->p;

  for (i = 0; i < n; i++) {
    HYT(0, i) = pdat->delta - YT(0, i) * cos(pdat->a) - YT(1, i) * sin(pdat->a);
    HYT(1, i) = YT(2, i) - pdat->a;

    /* Make sure that the theta is in the same range as the observed data,
       which in this case is [0, 2*pi) */
    if (HYT(1, i) < 0.0)
      HYT(1, i) += 2 * X01AAC;
  }
  /* Set info nonzero to terminate execution for any reason. */
  *info = 0;
}

static void read_problem_dat(Integer t, Nag_Comm *comm) {
  /* Read in any data specific to the f and h functions */
  Integer tt;
  g13_problem_data *pdat;

  if (t == 0) {
    /* Allocate some memory to hold the data */
    pdat = NAG_ALLOC(1, g13_problem_data);

    /* Read in the data that is constant across all time points */
    scanf("%lf%lf%lf%lf%*[^\n] ", &(pdat->r), &(pdat->d), &(pdat->delta),
          &(pdat->a));
    /* Set the void pointer in comm to point to the data structure */
    comm->p = (void *)pdat;

  } else if (t > 0) {
    /* Cast the void point in comm back to point at the data structure */
    pdat = (g13_problem_data *)comm->p;

    /* Read in data for time point t */
    scanf("%" NAG_IFMT "%lf%lf%*[^\n] ", &tt, &(pdat->phi_rt), &(pdat->phi_lt));
    if (tt != t) {
      /* Sanity check */
      printf("Expected to read in data for time point %" NAG_IFMT "\n", t);
      printf("Data that was read in was for time point %" NAG_IFMT "\n", tt);
    }

  } else {
    /* Clean up */

    /* Cast the void point in comm back to point at the data structure */
    pdat = (g13_problem_data *)comm->p;

    if (pdat)
      NAG_FREE(pdat);
  }
}