/* nag_ode_bvp_coll_nlin_solve (d02tlc) Example Program.
 *
 * NAGPRODCODE Version.
 *
 * Copyright 2016 Numerical Algorithms Group.
 * 
 * Mark 26, 2016.
 */

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

typedef struct
{
  double omega;
  double sqrofr;
} func_data;

#ifdef __cplusplus
extern "C"
{
#endif
  static void NAG_CALL ffun(double x, const double y[], Integer neq,
                            const Integer m[], double f[], Nag_Comm *comm);
  static void NAG_CALL fjac(double x, const double y[], Integer neq,
                            const Integer m[], double dfdy[], Nag_Comm *comm);
  static void NAG_CALL gafun(const double ya[], Integer neq,
                             const Integer m[], Integer nlbc, double ga[],
                             Nag_Comm *comm);
  static void NAG_CALL gbfun(const double yb[], Integer neq,
                             const Integer m[], Integer nrbc, double gb[],
                             Nag_Comm *comm);
  static void NAG_CALL gajac(const double ya[], Integer neq,
                             const Integer m[], Integer nlbc, double dgady[],
                             Nag_Comm *comm);
  static void NAG_CALL gbjac(const double yb[], Integer neq,
                             const Integer m[], Integer nrbc, double dgbdy[],
                             Nag_Comm *comm);
  static void NAG_CALL guess(double x, Integer neq, const Integer m[],
                             double y[], double dym[], Nag_Comm *comm);
#ifdef __cplusplus
}
#endif

int main(void)
{
  /* Scalars */
  double one = 1.0;
  Integer exit_status = 0, neq = 3, mmax = 3, nlbc = 3, nrbc = 3;
  double dx, ermx, r, omega;
  Integer i, iermx, ijermx, j, k, licomm, lrcomm, mxmesh, ncol, ncont, nmesh;
  /* Arrays */
  static double ruser[7] = { -1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0 };
  double rdum[1];
  Integer idum[2];
  double *mesh = 0, *tol = 0, *rcomm = 0, *y = 0;
  Integer *ipmesh = 0, *icomm = 0, *m = 0;
  func_data fd;
  /* Nag Types */
  Nag_Comm comm;
  NagError fail;

  INIT_FAIL(fail);

  printf("nag_ode_bvp_coll_nlin_solve (d02tlc) Example Program Results\n\n");

  /* For communication with user-supplied functions: */
  comm.user = ruser;

  /* Skip heading in data file */
  scanf("%*[^\n] ");
  scanf("%" NAG_IFMT "%" NAG_IFMT "%" NAG_IFMT "%*[^\n] ", &ncol, &nmesh,
        &mxmesh);

  if (!(mesh = NAG_ALLOC(mxmesh, double)) ||
      !(m = NAG_ALLOC(neq, Integer)) ||
      !(tol = NAG_ALLOC(neq, double)) ||
      !(y = NAG_ALLOC(neq * mmax, double)) ||
      !(ipmesh = NAG_ALLOC(mxmesh, Integer)))
  {
    printf("Allocation failure\n");
    exit_status = -1;
    goto END;
  }

  /* Set problem equation orders. */
  m[0] = 1;
  m[1] = 3;
  m[2] = 2;
  scanf("%lf%*[^\n] ", &omega);
  for (i = 0; i < neq; i++)
    scanf("%lf", &tol[i]);
  scanf("%*[^\n] ");

  dx = one / (double) (nmesh - 1);
  mesh[0] = 0.0;
  ipmesh[0] = 1;
  for (i = 1; i < nmesh - 1; i++) {
    mesh[i] = mesh[i - 1] + dx;
    ipmesh[i] = 2;
  }
  mesh[nmesh - 1] = one;
  ipmesh[nmesh - 1] = 1;

  /* Query to get size of rcomm and icomm (by setting lrcomm=0) */
  /* nag_ode_bvp_coll_nlin_setup (d02tvc).
   * Ordinary differential equations, general nonlinear boundary value problem,
   * setup for nag_ode_bvp_coll_nlin_solve (d02tlc).
   */
  nag_ode_bvp_coll_nlin_setup(neq, m, nlbc, nrbc, ncol, tol, mxmesh, nmesh,
                              mesh, ipmesh, rdum, 0, idum, 2, &fail);
  if (fail.code == NE_NOERROR) {
    lrcomm = idum[0];
    licomm = idum[1];
    printf("lrcomm = %" NAG_IFMT ", licomm = %" NAG_IFMT "\n", lrcomm,
           licomm);
    if (!(rcomm = NAG_ALLOC(lrcomm, double))
        || !(icomm = NAG_ALLOC(licomm, Integer)))
    {
      printf("Allocation failure\n");
      exit_status = -2;
      goto END;
    }

    /* Initialize again using  nag_ode_bvp_coll_nlin_setup (d02tvc). */
    nag_ode_bvp_coll_nlin_setup(neq, m, nlbc, nrbc, ncol, tol, mxmesh, nmesh,
                                mesh, ipmesh, rcomm, lrcomm, icomm, licomm,
                                &fail);
  }
  if (fail.code != NE_NOERROR) {
    printf("Error from nag_ode_bvp_coll_nlin_setup (d02tvc).\n%s\n",
           fail.message);
    exit_status = 1;
    goto END;
  }

  /* Number of continuation steps (last r=100**ncont, sqrofr=10**ncont) */
  scanf("%" NAG_IFMT "%*[^\n] ", &ncont);
  /* Initialize problem continuation parameter. */
  scanf("%lf%*[^\n] ", &r);

  /* Set data required for the user-supplied functions */
  fd.omega = omega;
  fd.sqrofr = sqrt(r);
  /* Associate the data structure with comm.p */
  comm.p = (Pointer) &fd;

  for (j = 0; j < ncont; j++) {
    printf("\n Tolerance = %8.1e  r = %10.3e\n", tol[0], r);
    /* Solve problem. */

    /* nag_ode_bvp_coll_nlin_solve (d02tlc).
     * Ordinary differential equations, general nonlinear boundary value 
     * problem, collocation technique.
     */
    nag_ode_bvp_coll_nlin_solve(ffun, fjac, gafun, gbfun, gajac, gbjac, guess,
                                rcomm, icomm, &comm, &fail);
    if (fail.code != NE_NOERROR) {
      printf("Error from nag_ode_bvp_coll_nlin_solve (d02tlc).\n%s\n",
             fail.message);
      exit_status = 2;
      goto END;
    }

    /* Extract mesh */

    /* nag_ode_bvp_coll_nlin_diag (d02tzc).
     * Ordinary differential equations, general nonlinear boundary value
     * problem, diagnostics for nag_ode_bvp_coll_nlin_solve (d02tlc).
     */
    nag_ode_bvp_coll_nlin_diag(mxmesh, &nmesh, mesh, ipmesh, &ermx, &iermx,
                               &ijermx, rcomm, icomm, &fail);
    if (fail.code != NE_NOERROR) {
      printf("Error from nag_ode_bvp_coll_nlin_diag (d02tzc).\n%s\n",
             fail.message);
      exit_status = 3;
      goto END;
    }

    /* Print mesh and error statistics. */
    printf("\n Used a mesh of %4" NAG_IFMT "  points\n", nmesh);
    printf(" Maximum error = %10.2e in interval %4" NAG_IFMT " for component"
           " %4" NAG_IFMT "\n\n\n", ermx, iermx, ijermx);
    printf(" Mesh points:\n");
    for (i = 0; i < nmesh; i++) {
      printf("%4" NAG_IFMT "(%1" NAG_IFMT ")", i + 1, ipmesh[i]);
      printf("%12.4e%s", mesh[i], i % 4 == 3 ? "\n" : " ");
    }
    /* Print solution components on mesh. */
    printf("\n\n      x        f        f'       g\n");
    for (i = 0; i < nmesh; i++) {

      /* nag_ode_bvp_coll_nlin_interp (d02tyc).
       * Ordinary differential equations, general nonlinear boundary value 
       * problem, interpolation for nag_ode_bvp_coll_nlin_solve (d02tlc).
       */
      nag_ode_bvp_coll_nlin_interp(mesh[i], y, neq, mmax, rcomm, icomm,
                                   &fail);
      if (fail.code != NE_NOERROR) {
        printf("Error from nag_ode_bvp_coll_nlin_interp (d02tyc).\n%s\n",
               fail.message);
        exit_status = 4;
        goto END;
      }

      printf("%8.3f ", mesh[i]);
      for (k = 0; k < neq; k++)
        printf("%9.4f", y[k]);
      printf("\n");
    }
    if (j == ncont - 1) {
      goto END;
    }
    /* Modify continuation parameter. */
    r = 100.0 * r;
    fd.sqrofr = sqrt(r);

    /* Select mesh for continuation and call continuation primer routine. */
    for (i = 1; i < nmesh - 1; i++) {
      ipmesh[i] = 2;
    }
    /* nag_ode_bvp_coll_nlin_contin (d02txc).
     * Ordinary differential equations, general nonlinear boundary value 
     * problem, continuation facility for nag_ode_bvp_coll_nlin_solve (d02tlc).
     */
    nag_ode_bvp_coll_nlin_contin(mxmesh, nmesh, mesh, ipmesh, rcomm, icomm,
                                 &fail);
    if (fail.code != NE_NOERROR) {
      printf("Error from nag_ode_bvp_coll_nlin_contin (d02txc).\n%s\n",
             fail.message);
      exit_status = 5;
      goto END;
    }
  }

END:
  NAG_FREE(mesh);
  NAG_FREE(m);
  NAG_FREE(tol);
  NAG_FREE(rcomm);
  NAG_FREE(y);
  NAG_FREE(ipmesh);
  NAG_FREE(icomm);
  return exit_status;
}

static void NAG_CALL ffun(double x, const double y[], Integer neq,
                          const Integer m[], double f[], Nag_Comm *comm)
{
  func_data *fd = (func_data *) comm->p;
  double y22 = y[1 + 2 * neq];
  double y31 = y[2 + 1 * neq];
  if (comm->user[0] == -1.0) {
    printf("(User-supplied callback ffun, first invocation.)\n");
    comm->user[0] = 0.0;
  }
  f[0] = y[1];
  f[1] = -(y[0] * y22 + y[2] * y31) * fd->sqrofr;
  f[2] = (y[1] * y[2] - y[0] * y31) * fd->sqrofr;
}

static void NAG_CALL fjac(double x, const double y[], Integer neq,
                          const Integer m[], double dfdy[], Nag_Comm *comm)
{
  func_data *fd = (func_data *) comm->p;
  double y22 = y[1 + 2 * neq];
  double y31 = y[2 + 1 * neq];

  if (comm->user[1] == -1.0) {
    printf("(User-supplied callback fjac, first invocation.)\n");
    comm->user[1] = 0.0;
  }
  dfdy[0 + 1 * neq + 0 * neq * neq] = 1.0;
  dfdy[1 + 0 * neq + 0 * neq * neq] = -y22 * fd->sqrofr;
  dfdy[1 + 1 * neq + 2 * neq * neq] = -y[0] * fd->sqrofr;
  dfdy[1 + 2 * neq + 0 * neq * neq] = -y31 * fd->sqrofr;
  dfdy[1 + 2 * neq + 1 * neq * neq] = -y[2] * fd->sqrofr;
  dfdy[2 + 0 * neq + 0 * neq * neq] = -y31 * fd->sqrofr;
  dfdy[2 + 1 * neq + 0 * neq * neq] = -y[2] * fd->sqrofr;
  dfdy[2 + 2 * neq + 0 * neq * neq] = y[1] * fd->sqrofr;
  dfdy[2 + 2 * neq + 1 * neq * neq] = -y[0] * fd->sqrofr;
}

static void NAG_CALL gafun(const double ya[], Integer neq, const Integer m[],
                           Integer nlbc, double ga[], Nag_Comm *comm)
{
  func_data *fd = (func_data *) comm->p;
  if (comm->user[2] == -1.0) {
    printf("(User-supplied callback gafun, first invocation.)\n");
    comm->user[2] = 0.0;
  }
  ga[0] = ya[0];
  ga[1] = ya[1];
  ga[2] = ya[2] - fd->omega;
}

static void NAG_CALL gbfun(const double yb[], Integer neq, const Integer m[],
                           Integer nrbc, double gb[], Nag_Comm *comm)
{
  func_data *fd = (func_data *) comm->p;
  if (comm->user[3] == -1.0) {
    printf("(User-supplied callback gbfun, first invocation.)\n");
    comm->user[3] = 0.0;
  }
  gb[0] = yb[0];
  gb[1] = yb[1];
  gb[2] = yb[2] + fd->omega;
}

static void NAG_CALL gajac(const double ya[], Integer neq, const Integer m[],
                           Integer nlbc, double dgady[], Nag_Comm *comm)
{
  if (comm->user[4] == -1.0) {
    printf("(User-supplied callback gajac, first invocation.)\n");
    comm->user[4] = 0.0;
  }
  dgady[0 + 0 * neq] = 1.0;
  dgady[1 + 1 * neq] = 1.0;
  dgady[2 + 2 * neq] = 1.0;
}

static void NAG_CALL gbjac(const double yb[], Integer neq, const Integer m[],
                           Integer nrbc, double dgbdy[], Nag_Comm *comm)
{
  if (comm->user[5] == -1.0) {
    printf("(User-supplied callback gbjac, first invocation.)\n");
    comm->user[5] = 0.0;
  }
  dgbdy[0 + 0 * neq] = 1.0;
  dgbdy[1 + 1 * neq] = 1.0;
  dgbdy[2 + 2 * neq] = 1.0;
}

static void NAG_CALL guess(double x, Integer neq, const Integer m[],
                           double y[], double dym[], Nag_Comm *comm)
{
  func_data *fd = (func_data *) comm->p;
  double xh = x - 0.5;
  double xx1 = x * (x - 1.0);

  if (comm->user[6] == -1.0) {
    printf("(User-supplied callback guess, first invocation.)\n");
    comm->user[6] = 0.0;
  }
  y[0 + 0 * neq] = -xh * pow(xx1, 2);
  y[1 + 0 * neq] = -xx1 * (5.0 * xx1 + 1.0);
  y[1 + 1 * neq] = -2.0 * xh * (10.0 * xx1 + 1.0);
  y[1 + 2 * neq] = -12.0 * (5.0 * xx1 + x);
  y[2 + 0 * neq] = -8.0 * fd->omega * pow(xh, 3);
  y[2 + 1 * neq] = -24.0 * fd->omega * pow(xh, 2);

  dym[0] = y[1];
  dym[1] = -120.0 * xh;
  dym[2] = -56.0 * fd->omega * xh;
}