/* ode-initval2/rk2imp.c
 * 
 * Copyright (C) 2009, 2010 Tuomo Keskitalo
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/* Based on rk2imp.c by Gerard Jungman */

/* Runge-Kutta 2, Gaussian implicit. Also known as implicit midpoint
   rule. Error estimation is carried out by the step doubling
   method.
 */

/* Reference: Ascher, U.M., Petzold, L.R., Computer methods for
   ordinary differential and differential-algebraic equations, SIAM,
   Philadelphia, 1998.
 */

#include <config.h>
#include <stdlib.h>
#include <string.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_odeiv2.h>

#include "odeiv_util.h"
#include "rksubs.c"
#include "modnewton1.c"

/* Stage of method */
#define RK2IMP_STAGE 1

typedef struct
{
  gsl_matrix *A;                /* Runge-Kutta coefficients */
  double *y_onestep;            /* Result with one step */
  double *y_twostep;            /* Result with two half steps */
  double *ytmp;                 /* Temporary work space */
  double *y_save;               /* Backup space */
  double *YZ;                   /* Runge-Kutta points */
  double *fYZ;                  /* Derivatives at YZ */
  gsl_matrix *dfdy;             /* Jacobian matrix */
  double *dfdt;                 /* time derivative of f */
  modnewton1_state_t *esol;     /* nonlinear equation solver */
  double *errlev;               /* desired error level of y */
  const gsl_odeiv2_driver *driver;      /* pointer to driver object */
}
rk2imp_state_t;

static void *
rk2imp_alloc (size_t dim)
{
  rk2imp_state_t *state = (rk2imp_state_t *) malloc (sizeof (rk2imp_state_t));

  if (state == 0)
    {
      GSL_ERROR_NULL ("failed to allocate space for rk2imp_state",
                      GSL_ENOMEM);
    }

  state->A = gsl_matrix_alloc (RK2IMP_STAGE, RK2IMP_STAGE);

  if (state->A == 0)
    {
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for A", GSL_ENOMEM);
    }

  state->y_onestep = (double *) malloc (dim * sizeof (double));

  if (state->y_onestep == 0)
    {
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for y_onestep", GSL_ENOMEM);
    }

  state->y_twostep = (double *) malloc (dim * sizeof (double));

  if (state->y_twostep == 0)
    {
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for y_onestep", GSL_ENOMEM);
    }

  state->ytmp = (double *) malloc (dim * sizeof (double));

  if (state->ytmp == 0)
    {
      free (state->y_twostep);
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for ytmp", GSL_ENOMEM);
    }

  state->y_save = (double *) malloc (dim * sizeof (double));

  if (state->y_save == 0)
    {
      free (state->ytmp);
      free (state->y_twostep);
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for y_save", GSL_ENOMEM);
    }

  state->YZ = (double *) malloc (dim * RK2IMP_STAGE * sizeof (double));

  if (state->YZ == 0)
    {
      free (state->y_save);
      free (state->ytmp);
      free (state->y_twostep);
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for YZ", GSL_ENOMEM);
    }

  state->fYZ = (double *) malloc (dim * RK2IMP_STAGE * sizeof (double));

  if (state->fYZ == 0)
    {
      free (state->YZ);
      free (state->y_save);
      free (state->ytmp);
      free (state->y_twostep);
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for fYZ", GSL_ENOMEM);
    }

  state->dfdt = (double *) malloc (dim * sizeof (double));

  if (state->dfdt == 0)
    {
      free (state->fYZ);
      free (state->YZ);
      free (state->y_save);
      free (state->ytmp);
      free (state->y_twostep);
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for dfdt", GSL_ENOMEM);
    }

  state->dfdy = gsl_matrix_alloc (dim, dim);

  if (state->dfdy == 0)
    {
      free (state->dfdt);
      free (state->fYZ);
      free (state->YZ);
      free (state->y_save);
      free (state->ytmp);
      free (state->y_twostep);
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for dfdy", GSL_ENOMEM);
    }

  state->esol = modnewton1_alloc (dim, RK2IMP_STAGE);

  if (state->esol == 0)
    {
      gsl_matrix_free (state->dfdy);
      free (state->dfdt);
      free (state->fYZ);
      free (state->YZ);
      free (state->y_save);
      free (state->ytmp);
      free (state->y_twostep);
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for esol", GSL_ENOMEM);
    }

  state->errlev = (double *) malloc (dim * sizeof (double));

  if (state->errlev == 0)
    {
      modnewton1_free (state->esol);
      gsl_matrix_free (state->dfdy);
      free (state->dfdt);
      free (state->fYZ);
      free (state->YZ);
      free (state->y_save);
      free (state->ytmp);
      free (state->y_twostep);
      free (state->y_onestep);
      gsl_matrix_free (state->A);
      free (state);
      GSL_ERROR_NULL ("failed to allocate space for errlev", GSL_ENOMEM);
    }

  state->driver = NULL;

  return state;
}

static int
rk2imp_apply (void *vstate, size_t dim, double t, double h,
              double y[], double yerr[],
              const double dydt_in[], double dydt_out[],
              const gsl_odeiv2_system * sys)
{
  /* Makes a Gaussian implicit 4th order Runge-Kutta step with size h
     and estimates the local error of the step by step doubling.
   */

  rk2imp_state_t *state = (rk2imp_state_t *) vstate;

  double *const y_onestep = state->y_onestep;
  double *const y_twostep = state->y_twostep;
  double *const ytmp = state->ytmp;
  double *const y_save = state->y_save;
  double *const YZ = state->YZ;
  double *const fYZ = state->fYZ;
  gsl_matrix *const dfdy = state->dfdy;
  double *const dfdt = state->dfdt;
  double *const errlev = state->errlev;

  const modnewton1_state_t *esol = state->esol;

#ifdef DEBUG
  {
    size_t di;

    printf ("rk2imp_apply: t=%.5e, h=%.5e, y:", t, h);

    for (di = 0; di < dim; di++)
      {
        printf ("%.5e ", y[di]);
      }
    printf ("\n");
  }
#endif

  /* Runge-Kutta coefficients */
  gsl_matrix *A = state->A;
  const double b[] = { 1.0 };
  const double c[] = { 0.5 };
  gsl_matrix_set (A, 0, 0, 0.5);

  if (esol == NULL)
    {
      GSL_ERROR ("no non-linear equation solver speficied", GSL_EINVAL);
    }

  /* Get desired error levels via gsl_odeiv2_control object through driver
     object, which is a requirement for this stepper.
   */

  if (state->driver == NULL)
    {
      return GSL_EFAULT;
    }
  else
    {
      size_t i;

      for (i = 0; i < dim; i++)
        {
          if (dydt_in != NULL)
            {
              gsl_odeiv2_control_errlevel (state->driver->c, y[i],
                                           dydt_in[i], h, i, &errlev[i]);
            }
          else
            {
              gsl_odeiv2_control_errlevel (state->driver->c, y[i],
                                           0.0, h, i, &errlev[i]);
            }
        }
    }

  /* Evaluate Jacobian for modnewton1 */

  {
#ifdef DEBUG
    printf ("-- evaluate jacobian\n");
#endif
    int s = GSL_ODEIV_JA_EVAL (sys, t, y, dfdy->data, dfdt);

    if (s != GSL_SUCCESS)
      {
#ifdef DEBUG
        printf ("-- FAIL at jacobian function evaluation\n");
#endif
        return s;
      }
#ifdef DEBUG
    size_t M;
    size_t N;
    size_t i;
    size_t j;

    M = dfdy->size1;
    N = dfdy->size2;

    for (i = 0; i < M; i++)
      {
        for (j = 0; j < N; j++)
          {
            double aij = gsl_matrix_get (dfdy, i, j);
            printf ("(%3lu,%3lu)[%lu,%lu]: %22.18g\n", M, N, i, j, aij);
          }
      }
#endif
  }

  /* Calculate a single step with size h */

  {
    int s = modnewton1_init ((void *) esol, A, h, dfdy, sys);

    if (s != GSL_SUCCESS)
      {
        return s;
      }

#ifdef DEBUG
    size_t M;
    size_t N;
    size_t i;
    size_t j;

    printf ("-- modnewton1_init IhAJ:\n");

    M = esol->IhAJ->size1;
    N = esol->IhAJ->size2;

    for (i = 0; i < M; i++)
      {
        for (j = 0; j < N; j++)
          {
            double aij = gsl_matrix_get (esol->IhAJ, i, j);
            printf ("(%3lu,%3lu)[%lu,%lu]: %22.18g\n", M, N, i, j, aij);
          }
      }

    printf ("-- modnewton1_init p:\n");

    M = esol->p->size;

    for (i = 0; i < M; i++)
      {
        double pi = gsl_permutation_get (esol->p, i);
        printf ("(%3lu)[%lu]: %22.18g\n", M, i, pi);
      }
#endif
  }

  {
    int s = modnewton1_solve ((void *) esol, A, c, t, h, y,
                              sys, YZ, errlev);
#ifdef DEBUG
    printf ("-- modnewton1_solve s=%d\n", s);
#endif

    if (s != GSL_SUCCESS)
      {
        return s;
      }
  }

#ifdef DEBUG
  {
    size_t di;

    printf ("YZ:");

    for (di = 0; di < dim * RK2IMP_STAGE; di++)
      {
        printf ("%.5e ", YZ[di]);
      }
    printf ("\n");
  }
#endif

  {
    int s = GSL_ODEIV_FN_EVAL (sys, t + c[0] * h, YZ, fYZ);

    if (s != GSL_SUCCESS)
      {
        return s;
      }
  }

#ifdef DEBUG
  {
    size_t di;

    printf ("fYZ:");

    for (di = 0; di < dim * RK2IMP_STAGE; di++)
      {
        printf ("%.5e ", fYZ[di]);
      }
    printf ("\n");
  }
#endif

  {
    int s = rksubs (y_onestep, h, y, fYZ, b, RK2IMP_STAGE, dim);

    if (s != GSL_SUCCESS)
      {
        return s;
      }
  }

  /* Error estimation by step doubling */

  {
    int s = modnewton1_init ((void *) esol, A, h / 2.0, dfdy, sys);

    if (s != GSL_SUCCESS)
      {
        return s;
      }
  }

  /* 1st half step */

  {
    int s = modnewton1_solve ((void *) esol, A, c, t, h / 2.0, y,
                              sys, YZ, errlev);

    if (s != GSL_SUCCESS)
      return s;
  }

  {
    int s = GSL_ODEIV_FN_EVAL (sys, t + c[0] * h / 2.0, YZ, fYZ);

    if (s != GSL_SUCCESS)
      {
        return s;
      }
  }

  {
    int s = rksubs (ytmp, h / 2.0, y, fYZ, b, RK2IMP_STAGE, dim);

    if (s != GSL_SUCCESS)
      {
        return s;
      }
  }

  /* Save original y values in case of error */

  DBL_MEMCPY (y_save, y, dim);

  /* 2nd half step */

  {
    int s = modnewton1_solve ((void *) esol, A, c, t + h / 2.0, h / 2.0,
                              ytmp, sys, YZ, errlev);

    if (s != GSL_SUCCESS)
      {
        return s;
      }
  }

  {
    int s = GSL_ODEIV_FN_EVAL (sys, t + h / 2.0 + c[0] * h / 2.0, YZ, fYZ);
    if (s != GSL_SUCCESS)
      {
        return s;
      }
  }

  {
    /* Note: rk2imp returns y using the results from two half steps
       instead of the single step since the results are freely
       available and more precise.
     */

    int s = rksubs (y_twostep, h / 2.0, ytmp, fYZ, b, RK2IMP_STAGE, dim);

    if (s != GSL_SUCCESS)
      {
        DBL_MEMCPY (y, y_save, dim);
        return s;
      }
  }

  DBL_MEMCPY (y, y_twostep, dim);

  /* Error estimation */

  {
    size_t i;
    for (i = 0; i < dim; i++)
      {
        yerr[i] = ODEIV_ERR_SAFETY * 0.5 *
          fabs (y_twostep[i] - y_onestep[i]) / 3.0;
      }
  }

#ifdef DEBUG
  {
    size_t i;
    printf ("-- error estimates: ");
    for (i = 0; i < dim; i++)
      {
        printf ("%.5e ", yerr[i]);
      }
    printf ("\n");
  }
#endif

  /* Derivatives at output */

  if (dydt_out != NULL)
    {
      int s = GSL_ODEIV_FN_EVAL (sys, t + h, y, dydt_out);

      if (s != GSL_SUCCESS)
        {
          /* Restore original values */
          DBL_MEMCPY (y, y_save, dim);

          return s;
        }
    }

  return GSL_SUCCESS;
}

static int
rk2imp_set_driver (void *vstate, const gsl_odeiv2_driver * d)
{
  rk2imp_state_t *state = (rk2imp_state_t *) vstate;

  state->driver = d;

  return GSL_SUCCESS;
}

static int
rk2imp_reset (void *vstate, size_t dim)
{
  rk2imp_state_t *state = (rk2imp_state_t *) vstate;

  DBL_ZERO_MEMSET (state->y_onestep, dim);
  DBL_ZERO_MEMSET (state->y_twostep, dim);
  DBL_ZERO_MEMSET (state->ytmp, dim);
  DBL_ZERO_MEMSET (state->y_save, dim);
  DBL_ZERO_MEMSET (state->YZ, dim);
  DBL_ZERO_MEMSET (state->fYZ, dim);

  return GSL_SUCCESS;
}

static unsigned int
rk2imp_order (void *vstate)
{
  rk2imp_state_t *state = (rk2imp_state_t *) vstate;
  state = 0;                    /* prevent warnings about unused parameters */
  return 2;
}

static void
rk2imp_free (void *vstate)
{
  rk2imp_state_t *state = (rk2imp_state_t *) vstate;

  free (state->errlev);
  modnewton1_free (state->esol);
  gsl_matrix_free (state->dfdy);
  free (state->dfdt);
  free (state->fYZ);
  free (state->YZ);
  free (state->y_save);
  free (state->ytmp);
  free (state->y_twostep);
  free (state->y_onestep);
  gsl_matrix_free (state->A);
  free (state);
}

static const gsl_odeiv2_step_type rk2imp_type = {
  "rk2imp",                     /* name */
  1,                            /* can use dydt_in? */
  1,                            /* gives exact dydt_out? */
  &rk2imp_alloc,
  &rk2imp_apply,
  &rk2imp_set_driver,
  &rk2imp_reset,
  &rk2imp_order,
  &rk2imp_free
};

const gsl_odeiv2_step_type *gsl_odeiv2_step_rk2imp = &rk2imp_type;