#include <cvode/cvode.h>
#include <nvector/nvector_serial.h>
#include <stdio.h>
#include <math.h>
#include <sunlinsol/sunlinsol_dense.h>
#include <sunmatrix/sunmatrix_dense.h>

/* User-defined vector and matrix accessor macros: Ith, IJth */

/* These macros are defined in order to write code which exactly matches
   the mathematical problem description given above.

   Ith(v,i) references the ith component of the vector v, where i is in
   the range [1..NEQ] and NEQ is defined below. The Ith macro is defined
   using the N_VIth macro in nvector.h. N_VIth numbers the components of
   a vector starting from 0.

   IJth(A,i,j) references the (i,j)th element of the dense matrix A, where
   i and j are in the range [1..NEQ]. The IJth macro is defined using the
   SM_ELEMENT_D macro. SM_ELEMENT_D numbers rows and columns of
   a dense matrix starting from 0. */

#define Ith(v, i) NV_Ith_S(v, i - 1) /* i-th vector component i=1..NEQ */


/* Problem Constants */
#define NEQ   2               /* number of equations  */
#define Y1    SUN_RCONST(0.45) /* initial y components */
#define Y2    SUN_RCONST(0.5)
#define RTOL  SUN_RCONST(1.0e-6) /* scalar relative tolerance            */
#define ATOL1 SUN_RCONST(1.0e-6) /* vector absolute tolerance components */
#define ATOL2 SUN_RCONST(1.0e-6)
#define T0    SUN_RCONST(0.0)  /* initial time           */
#define T1    SUN_RCONST(3600) /* first output time = 1 hour     */
#define TADD  SUN_RCONST(3600) /* time increment = 1 hour    */
#define NOUT  72               /* number of output times such that tmax = 3 days */


/* Functions Called by the Solver */

static int f(sunrealtype t, N_Vector y, N_Vector ydot, void* user_data);

static int g(sunrealtype t, N_Vector y, sunrealtype* gout, void* user_data);

static int Jac(sunrealtype t, N_Vector y, N_Vector fy, SUNMatrix J,
               void* user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3);

/* Private functions to output results */

// static void PrintOutput(sunrealtype t, sunrealtype y1, sunrealtype y2,
//                         sunrealtype y3);
// static void PrintRootInfo(int root_f1, int root_f2);

/* Private function to check function return values */

static int check_retval(void* returnvalue, const char* funcname, int opt);

/* Private function to check computed solution */

static int check_ans(N_Vector y, sunrealtype t, sunrealtype rtol, N_Vector atol);

static int root_fn(sunrealtype t, N_Vector y, sunrealtype *gout, void *user_data) {
  sunrealtype y1, y2;

  y1 = Ith(y, 1);
  y2 = Ith(y, 2);

  *gout = SUN_RCONST(10) * y2 -y1;
  return 0;
}

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

int main(void)
{
  SUNContext sunctx;
  sunrealtype t, tout;
  N_Vector y;
  N_Vector abstol;
  SUNMatrix A;
  SUNLinearSolver LS;
  void* cvode_mem;
  int retval, iout;
  int retvalr;
  int rootsfound[1];
  FILE *FID, *TFILE, *YFILE;
  int qcur;
  sunrealtype tcur, tstep, z;
  N_Vector eweight, ele;

  y         = NULL;
  abstol    = NULL;
  A         = NULL;
  LS        = NULL;
  cvode_mem = NULL;

  /* Create the SUNDIALS context */
  retval = SUNContext_Create(SUN_COMM_NULL, &sunctx);
  if (check_retval(&retval, "SUNContext_Create", 1)) { return (1); }

  /* Initial conditions */
  y = N_VNew_Serial(NEQ, sunctx);
  eweight = N_VNew_Serial(NEQ, sunctx);
  ele = N_VNew_Serial(NEQ, sunctx);
  if (check_retval((void*)y, "N_VNew_Serial", 0)) { return (1); }

  /* Initialize y */
  Ith(y, 1) = Y1;
  Ith(y, 2) = Y2;

  /* Set the vector absolute tolerance */
  abstol = N_VNew_Serial(NEQ, sunctx);
  if (check_retval((void*)abstol, "N_VNew_Serial", 0)) { return (1); }

  Ith(abstol, 1) = ATOL1;
  Ith(abstol, 2) = ATOL2;

  /* Call CVodeCreate to create the solver memory and specify the
   * Backward Differentiation Formula */
  cvode_mem = CVodeCreate(CV_BDF, sunctx);
  if (check_retval((void*)cvode_mem, "CVodeCreate", 0)) { return (1); }

  /* Call CVodeInit to initialize the integrator memory and specify the
   * user's right hand side function in y'=f(t,y), the initial time T0, and
   * the initial dependent variable vector y. */
  retval = CVodeInit(cvode_mem, f, T0, y);
  if (check_retval(&retval, "CVodeInit", 1)) { return (1); }

  /* Call CVodeSVtolerances to specify the scalar relative tolerance
   * and vector absolute tolerances */
  retval = CVodeSVtolerances(cvode_mem, RTOL, abstol);
  if (check_retval(&retval, "CVodeSVtolerances", 1)) { return (1); }

  /* Call CVodeRootInit to specify the root function g with 2 components */
  retval = CVodeRootInit(cvode_mem, 1, root_fn);
  if (check_retval(&retval, "CVodeRootInit", 1)) { return (1); }

  // CVodeSetMaxOrd(cvode_mem, 1);
  // CVodeSetMaxNumSteps(cvode_mem, -1);

  /* Create dense SUNMatrix for use in linear solves */
  A = SUNDenseMatrix(NEQ, NEQ, sunctx);
  if (check_retval((void*)A, "SUNDenseMatrix", 0)) { return (1); }

  /* Create dense SUNLinearSolver object for use by CVode */
  LS = SUNLinSol_Dense(y, A, sunctx);
  if (check_retval((void*)LS, "SUNLinSol_Dense", 0)) { return (1); }

  /* Attach the matrix and linear solver */
  retval = CVodeSetLinearSolver(cvode_mem, LS, A);
  if (check_retval(&retval, "CVodeSetLinearSolver", 1)) { return (1); }

  /* Set the user-supplied Jacobian routine Jac */
  // retval = CVodeSetJacFn(cvode_mem, Jac);
  // if (check_retval(&retval, "CVodeSetJacFn", 1)) { return (1); }

  /* In loop, call CVode, print results, and test for error.
     Break out of loop when NOUT preset output times have been reached.  */

  /* Open file for printing statistics */
  FID = fopen("solver_stats.csv", "w");
  TFILE = fopen("time.csv", "w");
  YFILE = fopen("y.csv", "w");

  iout = 0;
  tout = T1;
  while (1)
  {
    retval = CVode(cvode_mem, tout, y, &t, CV_NORMAL);
    CVodeGetCurrentTime(cvode_mem, &tcur);
    CVodeGetLastStep(cvode_mem, &tstep);
    CVodeGetCurrentOrder(cvode_mem, &qcur);
    CVodeGetErrWeights(cvode_mem, eweight);
    CVodeGetEstLocalErrors(cvode_mem, ele);
    printf("At asked_t = %0.4e y1 = %0.6e current_t = %0.4e step = %.1e order = %d err1 = %0.1e \n", 
      t, Ith(y, 1), tcur, tstep, qcur, Ith(ele, 1));
    // fprintf(TFILE, "%.6e,", t);
    // fprintf(YFILE, "%.6e,", Ith(y, 1));

    if (retval == CV_ROOT_RETURN)
    {
      retvalr = CVodeGetRootInfo(cvode_mem, rootsfound);
      if (check_retval(&retvalr, "CVodeGetRootInfo", 1)) { return (1); }
      printf("Root Detected\n");
      CVodeReInit(cvode_mem, t, y);
      // PrintRootInfo(rootsfound[0], rootsfound[1]);
    }

    if (check_retval(&retval, "CVode", 1)) { break; }
    if (retval == CV_SUCCESS)
    {
      iout++;
      tout += TADD;
    }

    retval = CVodePrintAllStats(cvode_mem, FID, SUN_OUTPUTFORMAT_CSV);

    if (iout == NOUT) { break; }
  }
  fclose(FID);
  fclose(TFILE);
  fclose(YFILE);

  /* Print final statistics to the screen */
  printf("\nFinal Statistics:\n");
  retval = CVodePrintAllStats(cvode_mem, stdout, SUN_OUTPUTFORMAT_TABLE);

  /* check the solution error */
  // retval = check_ans(y, t, RTOL, abstol);

  /* Free memory */
  N_VDestroy(y);            /* Free y vector */
  N_VDestroy(abstol);       /* Free abstol vector */
  CVodeFree(&cvode_mem);    /* Free CVODE memory */
  SUNLinSolFree(LS);        /* Free the linear solver memory */
  SUNMatDestroy(A);         /* Free the matrix memory */
  SUNContext_Free(&sunctx); /* Free the SUNDIALS context */

  return (retval);
}

/*
 *-------------------------------
 * Functions called by the solver
 *-------------------------------
 */

/*
 * ODE RHS
 */

static int f(sunrealtype t, N_Vector y, N_Vector ydot, void* user_data)
{
  sunrealtype y1, y2, yd1, yd2;

  y1 = Ith(y, 1);
  y2 = Ith(y, 2);

  yd1 = Ith(ydot, 1) = fmax(0, SUN_RCONST(10) * y2 -y1);
  yd2 = Ith(ydot, 2) = (SUN_RCONST(0.8)*sin(t/SUN_RCONST(3600)/SUN_RCONST(24)) - y2)/SUN_RCONST(3600);

  return (0);
}

/*
 * Check function return value...
 *   opt == 0 means SUNDIALS function allocates memory so check if
 *            returned NULL pointer
 *   opt == 1 means SUNDIALS function returns an integer value so check if
 *            retval < 0
 *   opt == 2 means function allocates memory so check if returned
 *            NULL pointer
 */

static int check_retval(void* returnvalue, const char* funcname, int opt)
{
  int* retval;

  /* Check if SUNDIALS function returned NULL pointer - no memory allocated */
  if (opt == 0 && returnvalue == NULL)
  {
    fprintf(stderr, "\nSUNDIALS_ERROR: %s() failed - returned NULL pointer\n\n",
            funcname);
    return (1);
  }

  /* Check if retval < 0 */
  else if (opt == 1)
  {
    retval = (int*)returnvalue;
    if (*retval < 0)
    {
      fprintf(stderr, "\nSUNDIALS_ERROR: %s() failed with retval = %d\n\n",
              funcname, *retval);
      return (1);
    }
  }

  /* Check if function returned NULL pointer - no memory allocated */
  else if (opt == 2 && returnvalue == NULL)
  {
    fprintf(stderr, "\nMEMORY_ERROR: %s() failed - returned NULL pointer\n\n",
            funcname);
    return (1);
  }

  return (0);
}


  ############################

  To unsubscribe from the SUNDIALS-USERS list:
  write to: mailto:SUNDIALS-USERS-SIGNOFF-REQUEST@LISTSERV.LLNL.GOV