/* -----------------------------------------------------------------------------
 * Programmer(s): David J. Gardner @ LLNL
 * -----------------------------------------------------------------------------
 * SUNDIALS Copyright Start
 * Copyright (c) 2002-2022, Lawrence Livermore National Security
 * and Southern Methodist University.
 * All rights reserved.
 *
 * See the top-level LICENSE and NOTICE files for details.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 * SUNDIALS Copyright End
 * -----------------------------------------------------------------------------
 * This example problem is adapted from:
 *
 * H. Ranocha, M. Sayyari, L. Dalcin, M. Parsani, and D.I. Ketcheson,
 * "Relaxation Runge-Kutta Methods: Fully-Discrete Explicit Entropy-Stable
 * Schemes for the Compressible Euler and Navier-Stokes Equations," SIAM Journal
 * on Scientific Computing, 42(2), 2020, https://doi.org/10.1137/19M1263480.
 * -----------------------------------------------------------------------------
 * This example evolves system
 *
 *   du/dt = -exp(v)
 *   dv/dt =  exp(u)
 *
 * for t in the interval [0, 5] with the initial condition
 *
 *   u(0) = 1.0
 *   v(0) = 0.5
 *
 * The system has the analytic solution
 *
 *   u = log(e + e^(3/2)) - log(b)
 *   v = log(a * e^(a * t)) - log(b)
 *
 * where log is the natural logarithm, a = sqrt(e) + e, and
 * b = sqrt(e) + e^(a * t).
 *
 * The conserved exponential entropy for the system is given by
 * ent(u,v) = exp(u) + exp(v) with the Jacobian
 * ent'(u,v) = [ de/du de/dv ]^T = [ exp(u) exp(v) ]^T.
 *
 * The problem is advanced in time with an explicit or implicit relaxed
 * Runge-Kutta method from ARKStep to ensure conservation of the entropy.
 * ---------------------------------------------------------------------------*/

/* Header files */
#include <math.h>
#include <stdio.h>

/* SUNDIALS headers */
#include <arkode/arkode_arkstep.h>
#include <nvector/nvector_serial.h>
#include <sunlinsol/sunlinsol_dense.h>
#include <sunmatrix/sunmatrix_dense.h>

/* Value of the natural number e */
#define EVAL 2.718281828459045235360287471352662497757247093699959574966

/* Convince macros for calling precision-specific math functions */
#if defined(SUNDIALS_DOUBLE_PRECISION)
#define EXP(x)  (exp((x)))
#define SQRT(x) (sqrt((x)))
#define LOG(x)  (log((x)))
#elif defined(SUNDIALS_SINGLE_PRECISION)
#define EXP(x)  (expf((x)))
#define SQRT(x) (sqrtf((x)))
#define LOG(x)  (logf((x)))
#elif defined(SUNDIALS_EXTENDED_PRECISION)
#define EXP(x)  (expl((x)))
#define SQRT(x) (sqrtl((x)))
#define LOG(x)  (logl((x)))
#endif

/* Convince macros for using precision-specific format specifiers */
#if defined(SUNDIALS_EXTENDED_PRECISION)
#define GSYM "Lg"
#define ESYM "Le"
#define FSYM "Lf"
#else
#define GSYM "g"
#define ESYM "e"
#define FSYM "f"
#endif

/* ----------------------- *
 * User-supplied functions *
 * ----------------------- */

/* ODE RHS function */
int f(sunrealtype t, N_Vector y, N_Vector ydot, void* user_data);

/* ODE RHS Jacobian function */
int Jac(sunrealtype t, N_Vector y, N_Vector fy, SUNMatrix J, void* user_data,
        N_Vector tmp1, N_Vector tmp2, N_Vector tmp3);

/* Entropy function */
int Ent(N_Vector y, sunrealtype* e, void* user_data);

/* Entropy Jacobian function */
int JacEnt(N_Vector y, N_Vector J, void* user_data);

/* ----------------- *
 * Utility functions *
 * ----------------- */

/* Analytic solution */
int ans(sunrealtype t, N_Vector y);

/* Check for an unrecoverable (negative) return flag from a SUNDIALS function */
int check_flag(int flag, const char* funcname);

/* Check if a function returned a NULL pointer */
int check_ptr(void* ptr, const char* funcname);

/* ------------ *
 * Main Program *
 * ------------ */

int main(int argc, char* argv[])
{
  /* Error-checking flag */
  int flag;

  /* Initial and final times */
  sunrealtype t0 = SUN_RCONST(0.0);
  sunrealtype tf = SUN_RCONST(5.0);

  /* Relative and absolute tolerances */
  sunrealtype reltol = SUN_RCONST(1.0e-6);
  sunrealtype abstol = SUN_RCONST(1.0e-10);

  /* SUNDIALS context, vector, matrix, and linear solver objects */
  SUNContext ctx     = NULL;
  N_Vector y         = NULL;
  N_Vector ytrue     = NULL;
  SUNMatrix A        = NULL;
  SUNLinearSolver LS = NULL;

  /* Pointer to vector data array */
  sunrealtype* ydata;
  sunrealtype* ytdata;

  /* Initial, current, and change in entropy value */
  sunrealtype ent0, ent, ent_err;

  /* Solution errors */
  sunrealtype u_err, v_err;

  /* ARKODE memory structure */
  void* arkode_mem = NULL;

  /* ARKODE statistics */
  long int nst, nst_a, nfe, nfi;
  long int nrf, nrbf, nre, nrje, nrnlsi, nrnlsf;
  long int nsetups, nje, nfeLS, nni, ncfn, netf;

  /* Output time and file */
  sunrealtype t;
  FILE* UFID;

  /* Command line options */
  int relax    = 1;                      /* enable relaxation */
  int implicit = 1;                      /* implicit          */
  sunrealtype fixed_h = SUN_RCONST(0.0); /* adaptive stepping */

  /* -------------------- *
   * Output Problem Setup *
   * -------------------- */

  if (argc > 1) relax = atoi(argv[1]);
  if (argc > 2) implicit = atoi(argv[2]);
  if (argc > 3) fixed_h = atof(argv[3]);

  printf("\nConserved Exponential Entropy problem:\n");
  if (implicit)
  {
    printf("   method     = DIRK\n");
  }
  else
  {
    printf("   method     = ERK\n");
  }
  printf("   reltol     = %.1" ESYM "\n", reltol);
  printf("   abstol     = %.1" ESYM "\n", abstol);
  if (fixed_h > SUN_RCONST(0.0))
  {
    printf("   fixed h    = %.1" ESYM "\n", fixed_h);
  }
  if (relax)
  {
    printf("   relaxation = ON\n");
  }
  else
  {
    printf("   relaxation = OFF\n");
  }
  printf("\n");

  /* ------------ *
   * Setup ARKODE *
   * ------------ */

  /* Create the SUNDIALS context object for this simulation */
  flag = SUNContext_Create(NULL, &ctx);
  if (check_flag(flag, "SUNContext_Create")) return 1;

  /* Create serial vector and set the initial condition values */
  y = N_VNew_Serial(2, ctx);
  if (check_ptr(y, "N_VNew_Serial")) return 1;

  ydata = N_VGetArrayPointer(y);
  if (check_ptr(ydata, "N_VGetArrayPointer")) return 1;

  ydata[0] = SUN_RCONST(1.0);
  ydata[1] = SUN_RCONST(0.5);

  ytrue = N_VClone(y);
  if (check_ptr(ytrue, "N_VClone")) return 1;

  ytdata = N_VGetArrayPointer(ytrue);
  if (check_ptr(ytdata, "N_VGetArrayPointer")) return 1;

  /* Initialize ARKStep */
  if (implicit)
  {
    arkode_mem = ARKStepCreate(NULL, f, t0, y, ctx);
  }
  else
  {
    arkode_mem = ARKStepCreate(f, NULL, t0, y, ctx);
  }
  if (check_ptr(arkode_mem, "ARKStepCreate")) return 1;

  /* Specify tolerances */
  flag = ARKStepSStolerances(arkode_mem, reltol, abstol);
  if (check_flag(flag, "ARKStepSStolerances")) return 1;

  if (relax)
  {
    /* Enable relaxation methods */
    flag = ARKStepSetRelaxFn(arkode_mem, Ent, JacEnt);
    if (check_flag(flag, "ARKStepSetRelaxFn")) return 1;
  }

  if (implicit)
  {
    /* Create dense matrix and linear solver */
    A = SUNDenseMatrix(2, 2, ctx);
    if (check_ptr(A, "SUNDenseMatrix")) return 1;

    LS = SUNLinSol_Dense(y, A, ctx);
    if (check_ptr(LS, "SUNLinSol_Dense")) return 1;

    /* Attach the matrix and linear solver */
    flag = ARKStepSetLinearSolver(arkode_mem, LS, A);
    if (check_flag(flag, "ARKStepSetLinearSolver")) return 1;

    /* Set Jacobian routine */
    flag = ARKStepSetJacFn(arkode_mem, Jac);
    if (check_flag(flag, "ARKStepSetJacFn")) return 1;

    /* Select a Butcher table with non-negative b values */
    flag = ARKStepSetTableName(arkode_mem, "ARKODE_ARK2_DIRK_3_1_2",
                               "ARKODE_ERK_NONE");
    if (check_flag(flag, "ARKStepSetTableName")) return 1;
  }

  if (fixed_h > SUN_RCONST(0.0))
  {
    flag = ARKStepSetFixedStep(arkode_mem, fixed_h);
    if (check_flag(flag, "ARKStepSetFixedStep")) return 1;
  }

  /* Open output stream for results, output comment line */
  UFID = fopen("ark_conserved_exp_entropy_ark.txt", "w");
  fprintf(UFID, "# vars: t u v entropy u_err v_err entropy_error\n");

  /* --------------- *
   * Advance in Time *
   * --------------- */

  /* Initial time */
  t = t0;

  /* Output the initial condition and entropy */
  flag = Ent(y, &ent0, NULL);
  if (check_flag(flag, "Ent")) return 1;

  fprintf(UFID,
          "%23.16" ESYM " %23.16" ESYM " %23.16" ESYM " %23.16" ESYM
          " %23.16" ESYM " %23.16" ESYM " %23.16" ESYM "\n",
          t0, ydata[0], ydata[1], ent0, SUN_RCONST(0.0), SUN_RCONST(0.0),
          SUN_RCONST(0.0));

  printf(" step   t              u              v              e            "
         "  delta e\n");
  printf(" --------------------------------------------------------------------"
         "-----------\n");
  printf("%5d %14.6" ESYM " %14.6" ESYM " %14.6" ESYM " %14.6" ESYM " %14.6" ESYM
         "\n",
         0, t, ydata[0], ydata[1], ent0, SUN_RCONST(0.0));

  while (t < tf)
  {
    /* Evolve in time */
    flag = ARKStepEvolve(arkode_mem, tf, y, &t, ARK_ONE_STEP);
    if (check_flag(flag, "ARKStepEvolve")) break;

    /* Output solution and errors */
    flag = Ent(y, &ent, NULL);
    if (check_flag(flag, "Ent")) return 1;

    flag = ans(t, ytrue);
    if (check_flag(flag, "ans")) return 1;

    ent_err = ent - ent0;
    u_err   = ydata[0] - ytdata[0];
    v_err   = ydata[1] - ytdata[1];

    /* Output to the screen periodically */
    flag = ARKStepGetNumSteps(arkode_mem, &nst);
    check_flag(flag, "ARKStepGetNumSteps");

    if (nst % 40 == 0)
    {
      printf("%5ld %14.6" ESYM " %14.6" ESYM " %14.6" ESYM " %14.6" ESYM " %14.6" ESYM
             "\n",
             nst, t, ydata[0], ydata[1], ent, ent_err);
    }

    /* Write all steps to file */
    fprintf(UFID,
            "%23.16" ESYM " %23.16" ESYM " %23.16" ESYM " %23.16" ESYM
            " %23.16" ESYM " %23.16" ESYM " %23.16" ESYM "\n",
            t, ydata[0], ydata[1], ent, u_err, v_err, ent_err);
  }

  printf(" --------------------------------------------------------------------"
         "-----------\n");
  fclose(UFID);

  /* ------------ *
   * Output Stats *
   * ------------ */

  /* Get final statistics on how the solve progressed */
  flag = ARKStepGetNumSteps(arkode_mem, &nst);
  check_flag(flag, "ARKStepGetNumSteps");

  flag = ARKStepGetNumStepAttempts(arkode_mem, &nst_a);
  check_flag(flag, "ARKStepGetNumStepAttempts");

  flag = ARKStepGetNumErrTestFails(arkode_mem, &netf);
  check_flag(flag, "ARKStepGetNumErrTestFails");

  flag = ARKStepGetNumRhsEvals(arkode_mem, &nfe, &nfi);
  check_flag(flag, "ARKStepGetNumRhsEvals");

  printf("\nFinal Solver Statistics:\n");
  printf("   Internal solver steps = %li (attempted = %li)\n", nst, nst_a);
  printf("   Total number of error test failures = %li\n", netf);
  printf("   Total RHS evals:  Fe = %li,  Fi = %li\n", nfe, nfi);

  if (implicit)
  {
    flag = ARKStepGetNumNonlinSolvIters(arkode_mem, &nni);
    check_flag(flag, "ARKStepGetNumNonlinSolvIters");

    flag = ARKStepGetNumNonlinSolvConvFails(arkode_mem, &ncfn);
    check_flag(flag, "ARKStepGetNumNonlinSolvConvFails");

    flag = ARKStepGetNumLinSolvSetups(arkode_mem, &nsetups);
    check_flag(flag, "ARKStepGetNumLinSolvSetups");

    flag = ARKStepGetNumJacEvals(arkode_mem, &nje);
    check_flag(flag, "ARKStepGetNumJacEvals");

    flag = ARKStepGetNumLinRhsEvals(arkode_mem, &nfeLS);
    check_flag(flag, "ARKStepGetNumLinRhsEvals");

    printf("   Total number of Newton iterations = %li\n", nni);
    printf("   Total number of linear solver convergence failures = %li\n", ncfn);
    printf("   Total linear solver setups = %li\n", nsetups);
    printf("   Total number of Jacobian evaluations = %li\n", nje);
    printf("   Total RHS evals for setting up the linear system = %li\n", nfeLS);
  }

  if (relax)
  {
    flag = ARKStepGetNumRelaxFnEvals(arkode_mem, &nre);
    check_flag(flag, "ARKStepGetNumRelaxFnEvals");

    flag = ARKStepGetNumRelaxJacEvals(arkode_mem, &nrje);
    check_flag(flag, "ARKStepGetNumRelaxJacEvals");

    flag = ARKStepGetNumRelaxFails(arkode_mem, &nrf);
    check_flag(flag, "ARKStepGetNumRelaxFails");

    flag = ARKStepGetNumRelaxBoundFails(arkode_mem, &nrbf);
    check_flag(flag, "ARKStepGetNumRelaxBoundFails");

    flag = ARKStepGetNumRelaxSolveFails(arkode_mem, &nrnlsf);
    check_flag(flag, "ARKStepGetNumRelaxSolveFails");

    flag = ARKStepGetNumRelaxSolveIters(arkode_mem, &nrnlsi);
    check_flag(flag, "ARKStepGetNumRelaxSolveIters");

    printf("   Total Relaxation Fn evals    = %li\n", nre);
    printf("   Total Relaxation Jac evals   = %li\n", nrje);
    printf("   Total Relaxation fails       = %li\n", nrf);
    printf("   Total Relaxation bound fails = %li\n", nrbf);
    printf("   Total Relaxation NLS fails   = %li\n", nrnlsf);
    printf("   Total Relaxation NLS iters   = %li\n", nrnlsi);
  }
  printf("\n");

  /* -------- *
   * Clean up *
   * -------- */

  /* Free ARKStep integrator and SUNDIALS objects */
  ARKStepFree(&arkode_mem);
  SUNLinSolFree(LS);
  SUNMatDestroy(A);
  N_VDestroy(y);
  N_VDestroy(ytrue);
  SUNContext_Free(&ctx);

  return flag;
}

/* ----------------------- *
 * User-supplied functions *
 * ----------------------- */

/* ODE RHS function f(t,y). */
int f(sunrealtype t, N_Vector y, N_Vector ydot, void* user_data)
{
  sunrealtype* ydata = N_VGetArrayPointer(y);
  sunrealtype* fdata = N_VGetArrayPointer(ydot);

  fdata[0] = -EXP(ydata[1]);
  fdata[1] = EXP(ydata[0]);

  return 0;
}

/* ODE RHS Jacobian function J(t,y) = df/dy. */
int Jac(sunrealtype t, N_Vector y, N_Vector fy, SUNMatrix J, void* user_data,
        N_Vector tmp1, N_Vector tmp2, N_Vector tmp3)
{
  sunrealtype* ydata = N_VGetArrayPointer(y);
  sunrealtype* Jdata = SUNDenseMatrix_Data(J);

  /* column 0 */
  Jdata[0] = SUN_RCONST(0.0);
  Jdata[1] = EXP(ydata[0]);

  /* column 1 */
  Jdata[2] = -EXP(ydata[1]);
  Jdata[3] = SUN_RCONST(0.0);

  return 0;
}

/* Entropy function e(y) */
int Ent(N_Vector y, sunrealtype* e, void* user_data)
{
  sunrealtype* ydata = N_VGetArrayPointer(y);

  *e = EXP(ydata[0]) + EXP(ydata[1]);

  return 0;
}

/* Entropy function Jacobian Je(y) = de/dy */
int JacEnt(N_Vector y, N_Vector J, void* user_data)
{
  sunrealtype* ydata = N_VGetArrayPointer(y);
  sunrealtype* jdata = N_VGetArrayPointer(J);

  jdata[0] = EXP(ydata[0]);
  jdata[1] = EXP(ydata[1]);

  return 0;
}

/* ----------------- *
 * Utility functions *
 * ----------------- */

/* Analytic solution */
int ans(sunrealtype t, N_Vector y)
{
  sunrealtype a, b;
  sunrealtype* ydata = N_VGetArrayPointer(y);

  a = SQRT(EVAL) + EVAL;
  b = SQRT(EVAL) + EXP(a * t);

  ydata[0] = LOG(EVAL + EXP(SUN_RCONST(1.5))) - LOG(b);
  ydata[1] = LOG(a * EXP(a * t)) - LOG(b);

  return 0;
}

/* Check for an unrecoverable (negative) return flag from a SUNDIALS function */
int check_flag(int flag, const char* funcname)
{
  if (flag < 0)
  {
    fprintf(stderr, "ERROR: %s() returned %d\n", funcname, flag);
    return 1;
  }
  return 0;
}

/* Check if a function returned a NULL pointer */
int check_ptr(void* ptr, const char* funcname)
{
  if (!ptr)
  {
    fprintf(stderr, "ERROR: %s() returned NULL\n", funcname);
    return 1;
  }
  return 0;
}
